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 b4b7afb9a09..e3980bd8f3b 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -1306,6 +1306,8 @@ jobs: done openai-sdk-compat: name: OpenAI ${{ matrix.openai-version }} compatibility + needs: rust-bridge + if: ${{ !cancelled() }} permissions: contents: read pull-requests: read @@ -1342,10 +1344,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/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md index c79632e9195..e1f606a2eb7 100644 --- a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -40,6 +40,8 @@ env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" } Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway. +The `lite` CLI can set up this provider each time it starts Codex, with no `config.toml` edit. Export `LITELLM_PROXY_URL` and `LITELLM_PROXY_API_KEY`, then run `lite codex --subscription`. Run `lite config set codex_subscription true` to make it the default. + For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway). ## Native compatibility and Pi 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..7a07350e2d9 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 ( @@ -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, ) @@ -113,6 +115,7 @@ from litellm.types.utils import ( LLMResponseTypes, SpecialEnums, ) +from litellm.utils import get_utc_datetime class _ManagedFileRetrieve(Protocol): @@ -507,10 +510,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(): @@ -2387,12 +2392,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 +2411,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/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index b91b4a970e8..c25da7cc187 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -15,7 +15,11 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, cast from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import TypeAdapter +from litellm_enterprise.proxy.management_endpoints.project_model_quota import ( + describe_missing_quota, + missing_project_model_quota, +) +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -27,6 +31,15 @@ from litellm.proxy.management.teams.dependencies import get_team_access from litellm.proxy.management.users.service import org_admin_org_ids from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects +from litellm.proxy.management_helpers.model_quota_policy import ( + ENFORCE_PROJECT_MODEL_QUOTA_SETTING, + REQUIRE_BOTH_LIMITS, + ModelQuotaPolicy, + RouteOf, + no_route, + parse_model_quota_setting, + router_model_routes, +) from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) @@ -276,26 +289,6 @@ def _check_team_project_limits( ) -def _project_models_missing_positive_quota( - models: list[str] | None, - rpm_limits: Mapping[str, object] | None, - tpm_limits: Mapping[str, object] | None, -) -> list[str]: - """Return the models that lack a positive `rpm` AND `tpm` quota. - - A valid quota is a positive integer; null, zero, and negative are rejected - because downstream rate limiters treat a non-positive limit as immediately - exhausted (every request blocked). - """ - - def _is_positive(value: object) -> bool: - return isinstance(value, int) and not isinstance(value, bool) and value > 0 - - rpm = rpm_limits or {} - tpm = tpm_limits or {} - return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))] - - def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]: return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset() @@ -322,64 +315,79 @@ def _raise_on_project_models_expanding_at_request_time( raise HTTPException( status_code=400, detail={ - "error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead." + "error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' is enabled. List concrete model names instead." }, ) +def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None: + try: + return parse_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING)) + except ValidationError as e: + raise HTTPException( + status_code=500, + detail={"error": f"general_settings.{ENFORCE_PROJECT_MODEL_QUOTA_SETTING} is invalid: {e}"}, + ) from e + + def _raise_on_missing_project_model_quota( - data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset() + data: NewProjectRequest | UpdateProjectRequest, + access_group_names: frozenset[str] = frozenset(), + policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + route_of: RouteOf = no_route, ) -> None: - """Require a positive `rpm`/`tpm` quota for every model on project CREATE. + """Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE. `model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request model's `set_model_info` validator, so they are read from there. - - Only invoked when `general_settings.enforce_project_model_quota` is enabled - (default off), so it is opt-in and does not change behavior for existing users. """ _raise_on_project_models_expanding_at_request_time(data.models, access_group_names) metadata = data.metadata or {} - missing = _project_models_missing_positive_quota( - data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit") + missing: Final = missing_project_model_quota( + data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, route_of ) if not missing: return raise HTTPException( status_code=400, detail={ - "error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model." + "error": f"models added to project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit." }, ) def _raise_on_missing_project_model_quota_on_update( - data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset() + data: UpdateProjectRequest, + existing_project: object, + access_group_names: frozenset[str] = frozenset(), + policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + route_of: RouteOf = no_route, ) -> None: - """Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE. + """Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE. `/project/update` replaces `models` and `metadata` when they are provided, so the check runs on what the project WILL look like: a partial update that doesn't touch models/quota keeps the existing values, while one that adds a model or clears a - model's quota must leave every resulting model with a positive limit. - - Only invoked when `general_settings.enforce_project_model_quota` is enabled - (default off), so it is opt-in and does not change behavior for existing users. + model's quota must leave every resulting model with the limits it needs. """ resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or []) resulting_metadata = ( data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {}) ) _raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names) - missing = _project_models_missing_positive_quota( - resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit") + missing: Final = missing_project_model_quota( + resulting_models, + resulting_metadata.get("model_rpm_limit"), + resulting_metadata.get("model_tpm_limit"), + policy, + route_of, ) if not missing: return raise HTTPException( status_code=400, detail={ - "error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model." + "error": f"models would be left on the project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit." }, ) @@ -530,6 +538,7 @@ async def new_project( """ from litellm.proxy.proxy_server import ( general_settings, + general_settings_view, litellm_proxy_admin_name, llm_router, premium_user, @@ -578,10 +587,6 @@ async def new_project( data=data, ) - # Opt-in (default off): require rpm/tpm for every model added to the project. - if general_settings.get("enforce_project_model_quota", False): - _raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router)) - # Check if user has permission to create projects for this team # only team admins can create projects for their team has_permission = await _check_user_permission_for_project( @@ -600,6 +605,15 @@ async def new_project( }, ) + quota_policy: Final = _project_model_quota_policy(general_settings_view()) + if quota_policy is not None: + _raise_on_missing_project_model_quota( + data, + _router_access_group_names(llm_router), + quota_policy, + router_model_routes(llm_router, data.team_id), + ) + # Generate project_id if not provided if data.project_id is None: data.project_id = str(uuid.uuid4()) @@ -721,6 +735,7 @@ async def update_project( """ from litellm.proxy.proxy_server import ( general_settings, + general_settings_view, litellm_proxy_admin_name, llm_router, premium_user, @@ -828,10 +843,14 @@ async def update_project( data=data, ) - # Opt-in (default off): require rpm/tpm for every model the update would leave on the project. - if general_settings.get("enforce_project_model_quota", False): + quota_policy: Final = _project_model_quota_policy(general_settings_view()) + if quota_policy is not None: _raise_on_missing_project_model_quota_on_update( - data, existing_project, _router_access_group_names(llm_router) + data, + existing_project, + _router_access_group_names(llm_router), + quota_policy, + router_model_routes(llm_router, target_team_id), ) # Prepare update data diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py new file mode 100644 index 00000000000..e36ff9b53dc --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py @@ -0,0 +1,49 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +from litellm.proxy.management_helpers.model_quota_policy import ( + ModelQuotaPolicy, + QuotaLimit, + RouteOf, + required_limits, + unmet_limits, +) + +_NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({}) + + +@dataclass(frozen=True, slots=True) +class MissingModelQuota: + model: str + limits: tuple[QuotaLimit, ...] + + +def _limits_set_for( + model: str, rpm_set: Mapping[str, object], tpm_set: Mapping[str, object] +) -> Mapping[QuotaLimit, object]: + return MappingProxyType({"rpm": rpm_set.get(model), "tpm": tpm_set.get(model)}) + + +def missing_project_model_quota( + models: Sequence[str], + rpm_limits: Mapping[str, object] | None, + tpm_limits: Mapping[str, object] | None, + policy: ModelQuotaPolicy, + route_of: RouteOf, +) -> tuple[MissingModelQuota, ...]: + rpm_set: Final = rpm_limits or _NO_LIMITS_SET + tpm_set: Final = tpm_limits or _NO_LIMITS_SET + gaps: Final = tuple( + MissingModelQuota( + model, + unmet_limits(required_limits(policy, model, route_of(model)), _limits_set_for(model, rpm_set, tpm_set)), + ) + for model in models + ) + return tuple(gap for gap in gaps if gap.limits) + + +def describe_missing_quota(missing: Sequence[MissingModelQuota]) -> str: + return ", ".join(f"{gap.model} ({' and '.join(f'model_{limit}_limit' for limit in gap.limits)})" for gap in missing) 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-rust/.agents/skills/rust-error/SKILL.md b/litellm-rust/.agents/skills/rust-error/SKILL.md new file mode 100644 index 00000000000..fd5d6222287 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/SKILL.md @@ -0,0 +1,36 @@ +--- +name: rust-error +description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings +--- + +# Rust errors + +`crates/http/src/error.rs` is the reference implementation. Read the reference for the part you are touching before editing, each one carries the reasoning and worked examples behind its rules + +- [shape](references/shape.md) + - One top-level `Error` per crate in `src/error.rs`. A public enum by default, an opaque struct over a private `Repr` enum when the variants must stay private. No `#[non_exhaustive]` + - A variant is a failure mode a caller handles differently. Provider names and guidance are data, never variants +- [messages and causes](references/messages-and-causes.md) + - Every message is written once, in `#[error(...)]`, and describes its own layer only. Call sites pass typed facts, never a string + - The cause is the `source()`, never part of the message. Keep it typed, never stringify or drop it + - `reqwest::Error` enters our types only through `litellm_http::Error`, which strips its URL + - Anything shown to someone else is rendered with `litellm_error::Report`, never `to_string()` +- [conversions](references/conversions.md) + - One `From` per source and target pair, only when the conversion needs no context. It classifies and sanitizes once + - Nest a lower crate's error in one variant. Translate when the meaning changes, never mirror its variants +- [classification](references/classification.md) + - Callers ask predicates named after their decision, they do not match nested variants. Predicates match exhaustively + - Classification states facts. Retry and fallback policy lives with the caller that knows idempotency +- [boundaries](references/boundaries.md) + - An error becomes text for someone else only in one mapping function per error type, in the gateway crates, the Python bridge and host-python. The inference gateway and the bridge decide by the same predicates + - A failure before the provider was called is a decline. Both ask `declined()` first and hand it back to the host unchanged +- [gateway](references/gateway.md) + - One `Error` with `IntoResponse` per gateway crate that serves a router. Every failure flows through the same `From` chain, is classified once in `status()` and rendered by surface and framing +- [pyo3](references/pyo3.md) + - Two `create_exception!` classes with the payload in `args`, `ValueError` for the caller's fault and `RuntimeError` otherwise, decoded in `failures.py` +- [testing](references/testing.md) + - Test predicates, chain shape and redaction on errors built the way production builds them +- [recipes](references/recipes.md) + - Common tasks walked through end to end + +The workspace uses thiserror 2 for every error type, and no library crate uses `anyhow` or `Box` in a public signature. The [derive docs](https://docs.rs/thiserror/latest/thiserror/) are the authority on macro behavior, and the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html) on what `Display` and `source()` mean diff --git a/litellm-rust/.agents/skills/rust-error/references/boundaries.md b/litellm-rust/.agents/skills/rust-error/references/boundaries.md new file mode 100644 index 00000000000..876ebfdc184 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/boundaries.md @@ -0,0 +1,45 @@ +# Boundaries + +A boundary is where an error becomes something another program reads: an HTTP response, a Python exception, a stored or logged message. Each boundary has one mapping function per error type it receives, and nothing outside those functions formats an error for someone else + +| Boundary | Mapping | Reference | +|---|---|---| +| Inference HTTP response | `Error::status()` and its renderers in `crates/gateway-inference/src/error.rs` | [gateway](gateway.md) | +| Other gateway surfaces | the `IntoResponse` of each `crates/gateway-*` crate's `Error` | [gateway](gateway.md) | +| Python bridge, inference and OCR routes | `route_failure` in `crates/python-bridge/src/errors.rs` and in `crates/python-bridge/src/routes/ocr/errors.rs` | [pyo3](pyo3.md) | +| Python bridge, everything else | one `From for PyErr` or mapping function per crate error, all in `crates/python-bridge` | [pyo3](pyo3.md) | +| Python host runtime | `From for PyErr` in `crates/host-python/src/error.rs` | [pyo3](pyo3.md) | + +## Asking the nested error + +The inference gateway and the Python bridge's route mapping both receive a nested route error and decide by the same predicates in the same order, so a new failure lands correctly on both once its predicates answer right + +| Predicate | Gateway | Bridge | +|---|---|---| +| `declined()` | 501 | `RustBridgeDeclined`, without the route's public mapping | +| `upstream_response()` | the provider's status and body | `RustUpstreamError` | +| `is_missing_credential()` | 401 | no class of its own, a `RuntimeError` because `is_request()` answers false | +| `is_request()` | 400 | `ValueError` | +| none of these | 500 | `RuntimeError` | + +If you find yourself adding an arm for your variant at a boundary, check whether a predicate is missing first + +The two boundaries are not identical. The gateway has variants of its own for body parsing, an unknown model and an unsupported path, and the OCR bridge has Python parity arms such as `FileNotFoundError` for a missing file. Those arms match the boundary's own variants or the top level of a public enum, never a nested error's variants. A missing API key is 401 at the gateway and `RuntimeError` at both bridges, pinned by tests to what Python's own path does + +## Declines + +A failure before anything reached the provider is a decline, and `declined()` on the route error names its reason. It means another host may serve the same request as it was received, so each boundary hands it to the host unchanged: the gateway answers 501 and the Python bridge raises `RustBridgeDeclined`, never a public exception or a `ValueError`. See [gateway](gateway.md) and [pyo3](pyo3.md) for each host's side + +Anything after the send is terminal, never a decline. `declined()` answers `None` for a response fault and for a post-call hook failure even when the wrapped error was a decline, and every route that transforms the provider's answer collapses its errors with `Error::after_send`, which keeps a provider answer and a response fault as they are and nests anything else as the source of an invalid response, so a decline can never describe a call that was billed. Ask `declined()`, do not match the variant behind it, because that after-send cut lives in the predicate + +## What reaches the client + +A boundary renders anything that is not a provider answer with `litellm_error::Report`, so the reader sees every layer of the chain once. That is why messages are written for a reader, each layer leaves its cause to `source()`, and reqwest URLs are stripped on entry. A boundary that must not show the chain, such as the bridge's TLS settings error that names the setting but not the file, renders fixed text of its own and says so in its test + +Provider bodies are passed through intact for gateway and Python error mapping. `litellm_http::Error::upstream(status, body)` retains the complete body, but its `Display`, `Debug` and `Report` diagnostics show only the first 256 characters and mark the cut. Do not log the raw body or `RustUpstreamError` args. OCR's `Provider` error also keeps the whole body for Python's `get_error_class`, bounded by the connection's response size limit + +Never format `Debug` of an error into a response or a Python exception. A type whose sources may hold secrets derives `veil::Redact`, so its `Debug` is safe to log + +## Logging + +Propagate with context, and let the terminal handler own the one log or trace event for a failed call. Do not log at every `?`, and do not log an error you also return. Diagnostic tracing excludes raw error strings and bodies, see [`rust-tracing`](../../rust-tracing/SKILL.md) diff --git a/litellm-rust/.agents/skills/rust-error/references/classification.md b/litellm-rust/.agents/skills/rust-error/references/classification.md new file mode 100644 index 00000000000..e32463847db --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/classification.md @@ -0,0 +1,66 @@ +# Classification + +## Ask, don't match + +Downstream crates ask a lower error questions instead of matching its nested variants. `litellm_http::Error` answers these + +| Predicate | Caller decides | +|---|---| +| `dispatched()` | whether retrying elsewhere could repeat a billed call | +| `is_request()` | the caller's input is at fault, so 400 and no reroute | +| `is_response()` | the provider answered with something we cannot use | +| `upstream_response()` | the provider's own status and body to pass through | +| `status()` | an HTTP status to report | +| `is_connect()` | the connection was never made, for a host's own connect-error mapping | +| `tls_source()` | which TLS setting to blame in a config error | +| `is_computed_header()` | the caller forwarded a header the signer owns, which `litellm_inference::Error` turns into a decline | + +Upper types add the questions their callers ask. `litellm_inference::Error` and the OCR error answer `is_missing_credential()`, and `litellm_llms::Error` and `litellm_inference::Error` answer `declined()`, which the boundaries ask first instead of reaching into `Auth` or `Declined`, see [boundaries](boundaries.md) + +```rust +// reaches through two crates' variants and breaks when either changes +Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => StatusCode::UNAUTHORIZED, + +// asks the question it cares about +Self::Route(error) if error.is_missing_credential() => StatusCode::UNAUTHORIZED, +``` + +A predicate is an exhaustive `match` with no `_` arm, never `matches!`, so a new variant does not compile until it has answered every question on purpose + +```rust +pub fn is_request(&self) -> bool { + match &*self.0 { + Repr::Request(_) | Repr::BlockedUrl | Repr::DownloadTooLarge => true, + Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => false, + // ... every other variant, listed by name + } +} +``` + +Add a predicate when a real caller needs a new distinction, and name it after the fact the caller needs, in the caller's words, not after the variant's spelling. An upper type's predicates call through to the nested error, as in `Self::Http(error) => error.is_request()` and `Self::Auth(error) => error.is_missing_credential()` + +A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `Unsupported` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and the old `Unsupported` route variant + +## Facts, not policy + +A predicate reports what happened. The decision to retry, reroute or fall back depends on things the error cannot know, such as whether the request is idempotent, whether its body can be replayed, the deadline and the attempt budget. That decision lives with the caller. `dispatched()` exists for the router's reroute decision, so every transport failure answers it even before that caller is wired + +A timeout is the main trap. A request that timed out may already have reached the provider and been billed, so `Timeout` answers `dispatched() == true` and is not safe to retry as if it were a connect failure. Only a failure that provably happened before any byte reached the provider answers `false`: connecting, building or signing the request, and media downloads, which fetch the caller's document before the provider call is made + +```rust +pub fn dispatched(&self) -> bool { + match &*self.0 { + Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => true, + Repr::Connect(_) | Repr::RequestBody(_) | Repr::Download(_) | Repr::BlockedUrl => false, + // ... every other variant, listed by name + } +} +``` + +When the transport cannot tell whether a request reached the provider, it must answer `true`. Wrongly answering `false` risks a duplicate billed call, wrongly answering `true` only costs a reroute + +Keep one predicate per question. `status()` currently mixes the provider's status, the media download's status and a synthesized 408 for timeouts, so a caller that wants "did the provider answer" must use `upstream_response()` instead. Do not add more meanings to an existing predicate, add a new one + +## Equality + +`Arc` errors compare with `Arc::ptr_eq`, so a clone equals its original and two separately built errors never compare equal. That exists only so a wrapping enum such as `litellm_inference::Error` can keep its `PartialEq` derive. Do not give a new type pointer equality for any other reason, and do not rely on it in tests. Assert predicates instead, see [testing](testing.md) diff --git a/litellm-rust/.agents/skills/rust-error/references/conversions.md b/litellm-rust/.agents/skills/rust-error/references/conversions.md new file mode 100644 index 00000000000..e40b76b123f --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/conversions.md @@ -0,0 +1,92 @@ +# Conversions + +## One From per source and target pair + +A `From for Ours` exists when the conversion is obvious, infallible and needs no context from the call site. It does all classification and sanitizing for that pair, so every call site is `?` or `.map_err(Ours::from)` + +```rust +impl From for Error { + fn from(error: reqwest::Error) -> Self { + let error = error.without_url(); + if error.is_timeout() { + Repr::Timeout(error).into() + } else if error.is_connect() || error.is_builder() { + Repr::Connect(error).into() + } else { + Repr::Network(error).into() + } + } +} +``` + +A call site never re-derives that classification + +```rust +// every route used to repeat this, and each drifted a little +.map_err(|err| { + if err.is_connect() || err.is_builder() { + Error::Transport(Connect(err.to_string())) + } else { + Error::Transport(Network(err.to_string())) + } +})? + +// now +.map_err(litellm_http::Error::from)? +``` + +The same foreign type may convert into several of our types, one `From` each. What is not allowed is two places converting the same pair differently + +## When the call site knows something the From cannot + +One `io::Error` can mean "reading the CA bundle failed" or "resolving the media host failed". A global `From` cannot know which, so there is none. Name the operation at the call site instead + +```rust +.map_err(Repr::Resolve)? +.map_err(|source| Repr::Read { path, tls_source, source })? +``` + +Do not drop meaningful context just to use `?` + +## Reclassify once, in the shared From + +When a lower error needs a different meaning for every caller, do it in the upper type's `From`, not per route + +```rust +impl From for Error { + fn from(error: litellm_http::Error) -> Self { + if error.is_computed_header() { + return Self::Declined(Decline::new("request forwards a header AWS SigV4 computes")); + } + Self::Http(error) + } +} +``` + +All four inference routes classify a header the signer computes the same way, because none of them does it themselves + +## Nest, translate or add context, never mirror + +| The upper layer | Do | Example | +|---|---|---| +| has nothing to add | nest in one transparent variant and delegate predicates | `Self::Http(error) => error.is_request()` | +| changes what the failure means | translate to an upper variant, keep the cause attached when it helps | computed header -> `Declined` | +| knows something the lower one does not | a variant with the extra fields plus the cause as its source | `Read { path, tls_source, source }` | + +Never re-declare a lower crate's variants in the upper type with an upper-layer name and a 1:1 mapping function. OCR once did that for every media download failure, `MediaError::BlockedUrl => Error::BlockedDocumentUrl` and so on, which doubled every change and let the two lists drift. + +## Third-party and host faults + +Converting a fault from a crate we do not own into a string variant is still flattening. Prefer a variant that carries the fault, and when an existing `ErrorDetail` constructor fits, use it + +```rust +// flattens +.map_err(|error| Error::InvalidRequest(error.to_string().into())) + +// keeps the cause typed +.map_err(|error| Error::InvalidRequest(ErrorDetail::invalid("header", error))) +``` + +## Through a foreign error type + +A trait we implement for another crate sometimes forces its error type on us, such as `io::Error` from a body stream or `sqlx::Error` from a driver. Wrap ours as that type's source (`io::Error::other(error)`, `azure_core::Error::new(kind, error)`) and recover it with `downcast_ref` on the other side, never `with_message(error.to_string())` or `from_msg`. The error must come back out as the same value it went in as, so one failure classifies the same way on every path diff --git a/litellm-rust/.agents/skills/rust-error/references/gateway.md b/litellm-rust/.agents/skills/rust-error/references/gateway.md new file mode 100644 index 00000000000..5da34aa84b2 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/gateway.md @@ -0,0 +1,56 @@ +# Gateway responses + +The gateway is where a Rust error becomes an HTTP response. This page is about the path an error takes to get there and the one place it is classified. The gateway crates' own conventions, such as which envelope each API surface uses, which extractors exist and how the router is composed, live in the [`rust-gateway`](../../rust-gateway/SKILL.md) skill + +## The flow + +```text +reqwest / provider failure + -> litellm_http::Error From, strips the URL, classifies once + -> litellm_inference::Error::Http(_) From, one transparent variant, predicates delegate + -> gateway_inference::Error::Route(_) From, one variant per nested crate + -> handler returns Result<_, Error> every `?` ends here + -> IntoResponse status() picks the code, a renderer builds the body +``` + +Each hop is one `From` that nests the lower error in one variant, see [conversions](conversions.md). No hop re-derives what the lower crate already decided, and nothing between the provider and `IntoResponse` formats an error for the client. That is what makes "one `Error` with `IntoResponse` per gateway crate that serves a router" work: a handler can `?` any failure from any depth and the response comes out right + +## Classify once + +`status()` is pure and runs once per error. Variants the gateway declares itself, such as a bad body, an unknown model or an oversized body, are matched at the top of its own enum. A nested route error goes to `route_status` and an OCR error to `ocr_status`, which ask the predicates [boundaries](boundaries.md) lists + +| Answer | Status | +|---|---| +| `declined()` | 501, for a route error, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider | +| `upstream_response()` | the provider's own status, 502 when it is not a valid code | +| `is_missing_credential()` | 401 | +| `is_request()` | 400, for a route error | +| `http_status_code()` | the transport's status, for an OCR error, such as 408 for a timeout | +| otherwise | 500 | + +Do not reach into a nested enum's variants from the gateway. When a new nested failure needs a different status, add a predicate on the nested error, see [classification](classification.md) + +## Render by surface and framing + +Rendering is separate from classification. The surface picks the envelope and the moment picks the framing, a JSON body before the headers went out and an `event: error` SSE frame after. Both framings of one surface render the same status and `Report` message, so a streaming failure and a JSON failure of the same error agree on everything but framing + +| Surface | JSON | SSE | +|---|---|---| +| Anthropic, `/v1/messages` | `body(request_id)` | `sse_frame()` | +| OpenAI, chat, responses and the rest | `openai_response()` | `responses_sse_frame()`, the Responses API `error` event | + +The message of anything that is not a provider answer is the `litellm_error::Report` of the error, and in the OpenAI and Anthropic JSON envelopes the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Responses `error` event has the fixed type `error` and the status as its `code`. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else. The OpenAI envelope always wraps, because that is what the Python proxy returns there + +## Every way in uses the same Error + +A request can fail before the handler runs, in the handler, or after the headers went out. All three end in the same `Error`, so they get the same classification and envelope + +- An extractor sets `type Rejection = Error`, so a bad body is the surface's envelope rather than axum's plain text +- A path the gateway knows but does not serve returns `Error::Unsupported(path)` from one handler, so the client gets a 501 in the surface's envelope rather than a bare 404 +- A stream that fails after its headers went out hands the surface's SSE frame of `Error::from(error)` to `litellm_host_http::Sse::new`. `litellm_host_http::Error` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the `From` maps the second to `Internal` + +The axum hooks behind each line are in [`rust-gateway`](../../rust-gateway/SKILL.md) + +## Testing + +`status()` and `body()` are pure, so unit tests call them on errors built the way production builds them, one `#[case]` per fault, as `status_follows_who_is_at_fault`, `an_anthropic_provider_body_passes_through` and `other_failures_get_the_anthropic_envelope_with_their_message` in `crates/gateway-inference/src/error.rs` do. Assert the status and the fields the surface promises, not the whole body, except where passing the body through untouched is the contract. Driving a whole router belongs to the gateway crate's own tests, see [`rust-gateway`](../../rust-gateway/references/testing.md) diff --git a/litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md b/litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md new file mode 100644 index 00000000000..80017423a42 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md @@ -0,0 +1,112 @@ +# Messages and causes + +## Messages belong to the variant + +Every message template is written once, in the variant's `#[error(...)]`. A caller passes the typed facts the template needs and nothing else + +```rust +// none of these may appear at a call site +Error::Network("WebSocket handshake returned a wrong Sec-WebSocket-Accept".into()) +Error::Request(format!("bad header {name}")) +Error::Signature(error.to_string()) + +// the call site names what happened, the variant owns the words +Repr::Handshake(HandshakeFault::Accept).into() +Error::signature(error) +Error::NonAsciiHeader { name } +``` + +When a check fails for a distinct reason, that reason is a variant or a field, never a string inside a generic variant. A `Message(String)` variant, or a `From` / `From<&str>` into an error, is the hole every stringly call site goes through, so neither exists + +Format arguments can call functions and reach into fields, so a hand-written `Display` is rarely needed + +```rust +#[error("could not read {}", path.display())] +#[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))] +``` + +Keep messages short and lowercase, with no trailing period, unless Python parity pins the exact text. A reporter appends `: {cause}` after them. Never format `Debug` (`{:?}`) of an error, or of a type with no `Display`, into a message + +## The cause is the source, never the message + +This is the rule from the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html): an error's `Display` describes that layer only, and the error underneath is returned by `source()`. thiserror wires `source()` from a field marked `#[source]`, a `#[from]` field, or a field named `source` + +```rust +// prints the cause twice under any reporter +#[error("could not read {}: {source}", path.display())] +Read { path: PathBuf, source: io::Error }, +#[error("invalid JSON: {0}")] +Json(#[from] serde_json::Error), + +// this layer's words, the cause is the source +#[error("could not read {}", path.display())] +Read { path: PathBuf, source: io::Error }, +#[error("invalid JSON")] +Json(#[from] serde_json::Error), +``` + +| Field | thiserror does | Use for | +|---|---|---| +| `source: E` or `#[source] cause: E` | exposes it as `source()` | a cause with context fields next to it | +| `#[from] E` | implies `#[source]` and adds `From` | context-free wrapping, see [conversions](conversions.md) | +| `#[error(transparent)]` | forwards both `Display` and `source()` to the inner value | nesting another of our errors that adds no text | +| no attribute, any other name | nothing, the field is plain data | facts the message renders, such as a path or a status | + +`#[error(transparent)]` is an invisible layer, not an extra chain node, so a transparent wrapper's `source()` is the inner value's `source()` + +A field that holds an error but is neither rendered nor a source hides that error from every reader. Every error-typed field is a source. A fieldless reason enum, such as `HandshakeFault` or `PagesFault`, is data and not a cause, so the variant renders it in its message + +`source()` is one chain, so an error with several causes, such as `CredentialChain` trying each credential in turn, makes the most relevant one its source and renders the others in its message, each with `Report` + +## Keep causes typed + +Store the lower-level error itself, never its string, and never drop it with `map_err(|_| ...)` + +| Flattening or dropping loses | Because | +|---|---| +| the chain | a `String` has no `source()` | +| downcasting | the type is gone | +| redaction | `reqwest::Error`'s `Display` includes the request URL, which can carry credentials in the path or query | + +If the cause is not `Clone` and the error must be, hold it as `#[source] litellm_error::Source`, a cloneable `Arc` that is transparent to `Display` and `source()`, or make the crate's error opaque. `litellm_error::find::(&error)` walks the chain and looks through `Source`, so it is how anyone reaches a typed cause. A third-party `Box` converts with `Source::from`, and an `anyhow::Error` a dependency returns goes through `Box` first, `Source::from(Box::::from(error))`, and a dependency's own boxed alias such as `tokenizers::Error` is held as it is. A source boxed to keep the variant small is a `Box` node of the chain, so it is reached as `find::>`. Never put `#[error(transparent)]` on a `Source`, the erased error would then be no node of the chain and `find` could not reach it. Do not stringify or drop a cause to get `Clone` or `PartialEq` + +```rust +// loses the type and leaks the URL +.map_err(|error| Error::Network(error.to_string())) +// loses everything +.map_err(|_| Error::Transport) + +// keeps the cause, the From impl strips the URL +.map_err(litellm_http::Error::from) +``` + +A cause that must not be shown at all, such as a secret manager's response, is still a source. The type derives [`veil::Redact`](https://docs.rs/veil) with `#[redact]` on the field, so `Debug` hides it, and the boundary decides how much of the chain to show, see [boundaries](boundaries.md) + +## Strip reqwest URLs on entry + +Every `reqwest::Error` enters our types through `From for litellm_http::Error`, which calls `without_url()`. A crate that talks HTTP holds `litellm_http::Error`, never `reqwest::Error`. When a dependency's own error wraps a `reqwest::Error`, as vaultrs and rustify do, one `From` takes the reqwest error out and converts it, as `secrets-hashicorp` does. That way so every later `Display`, `Debug` or chain walk is safe from that URL + +`without_url()` only removes the URL reqwest attached. It does not scrub headers, response bodies, provider messages, file paths or anything a variant stores next to it. Those need their own policy, see [boundaries](boundaries.md) + +## Rendering + +Anything that shows an error to someone else renders the whole chain with `litellm_error::Report`, which writes each layer once joined by `: `, the single-line form of the unstable [`std::error::Report`](https://doc.rust-lang.org/std/error/struct.Report.html). `to_string()` of an error is only its top layer and is never what a client, a Python exception or a log line shows + +```rust +let message = litellm_error::Report::new(&error).to_string(); +// "could not reach the provider: error sending request: client error (Connect): tcp connect error: Connection refused" +``` + +## #[from] limits + +`#[from]` only works on a variant whose only fields are the source and an optional backtrace. A variant that needs context such as a path, a phase or a `TlsSource` takes its cause through an explicit `map_err` or a named constructor + +```rust +std::fs::read(&path).map_err(|source| Repr::Read { + path: path.clone(), + tls_source: TlsSource::CaBundle, + source, +})?; +``` + +`#[backtrace]` and automatic `provide()` need a nightly compiler and are not used here diff --git a/litellm-rust/.agents/skills/rust-error/references/pyo3.md b/litellm-rust/.agents/skills/rust-error/references/pyo3.md new file mode 100644 index 00000000000..1c6ffb96a82 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/pyo3.md @@ -0,0 +1,76 @@ +# Python bridge + +`crates/python-bridge` turns a Rust error into a Python exception through one mapping function per error type. A `#[pyfunction]` returns `Err(PyErr)` and pyo3 raises it, so those functions are the whole boundary + +| Error | Mapping | +|---|---| +| `litellm_inference::Error` from a route | `errors::route_failure`, with `route_error_to_pyerr` as its shorthand for routes that set no marker | +| the OCR error | `route_failure` and `to_pyerr` in `routes/ocr/errors.rs` | +| token counter, cache, spend, secrets, argument checks | one `From for PyErr` or mapping function next to the code that registers the functions | + +A route host never builds a `PyErr` from a route error itself. It calls the mapping, and the mapping is the only place that knows the exception classes and their attributes + +## Exception classes + +| Class | Meaning | Payload | +|---|---|---| +| `RustUpstreamError` | the provider call was issued and the provider answered with a failure | `args == (status, body)` and a `headers` attribute of `(name, value)` pairs, always set | +| `RustBridgeDeclined` | the route stopped before calling the provider, so the host may run its own path | `args == (message,)` | +| a standard class from `pyo3::exceptions` | whichever Python class means the same thing, such as `ValueError` or `TypeError` for the caller's input, `FileNotFoundError` or `OSError` for a file, `KeyError` or `OverflowError` where Python raises them, and `RuntimeError` for anything else | the rendered chain, or a parity text where Python pins one | + +The two native classes come from `create_exception!` in `crates/python-bridge/src/errors.rs` and are exported with `#[pymodule_export]`, which is how the [pyo3 guide](https://github.com/PyO3/pyo3/blob/main/guide/src/exception.md) makes an exception importable. Their docstrings are the contract Python reads. `errors::upstream_error(py, status, body, headers)` is the one constructor of `RustUpstreamError`, so the payload cannot differ between routes + +Python decodes a `RustUpstreamError` in `litellm/rust_bridge/failures.py`, which validates `args` and `headers` with a `TypeAdapter` and rebuilds the provider's `httpx.Response`. `litellm/rust_bridge/runtime.py` catches `RustBridgeDeclined` and reads its reason. `_raise_upstream` in `runtime.py` still reads `args` a second time by hand for the `attempt` path, which is the one decode to fold into `failures.py`. A new fact goes into `args` or `headers`, and the decoder learns to read it in the same change + +Two routes also flag a request fault with a boolean attribute, read by their Python route module: `ocr_request_format_error` in `host/routes/ocr.py` and `messages_request_error` in `host/routes/messages.py`. OCR additionally sets `status_code` and `message` for Python parity, and `native_settings_error`, which only the OCR host reads back to tell a settings failure from the caller's exception. Do not add another attribute, nothing on the Python side types them + +## Why args, not a typed subclass + +The bridge builds with pyo3's `abi3-py310` feature. Under abi3, `#[pyclass(extends = PyException)]` is only allowed on Python 3.12 and newer, so a class with typed fields is not available. Even where it is allowed, pyo3 passes an empty tuple to the native base's `tp_new`, so `.args` and `str()` would need manual overrides. The payload therefore travels in `args`, which a `create_exception!` class fills from `new_err` + +## Mapping order + +A core decline is returned first, as `RustBridgeDeclined` with the decline's reason as its only arg, without calling the route's public mapping. Then a Python exception raised inside the call, such as a secret manager callback's, is returned as it is. Then the provider's answer passes through as `RustUpstreamError`. On `/v1/messages` a request the route rejected becomes a `ValueError` flagged with `messages_request_error`. Everything else is split by `is_request()` into `ValueError` or `RuntimeError`, with the message rendered by `litellm_error::Report`. A missing API key is a `RuntimeError` because `is_request()` answers false for it, and `missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors` pins that + +OCR puts Python parity arms in front of that order, each matching a top-level variant of its own public enum + +| Arm | Exception | +|---|---| +| the provider answered | `RustUpstreamError` with the provider's headers | +| `RequestFormat` | `ValueError` flagged with `ocr_request_format_error` | +| `FileRead` whose cause is `NotFound` | `FileNotFoundError` naming the path | +| other `FileRead` | `OSError` with the io message | + +Both routes ask `litellm_auth::Error::is_request()` through their own `is_request()`, so a missing key or API base is a `RuntimeError` on both, and a missing setting the caller can pass, such as a region, is a `ValueError` + +Use the standard class whose Python meaning matches. `PyFileNotFoundError` for a missing file is right, a `RuntimeError` whose text says the file is missing is not. `import_exception!` exists for an exception defined in Python, none is used today because `failures.py` maps to LiteLLM's public exceptions through `litellm.exception_type` + +## Declining + +`RustBridgeDeclined` means no provider call was made and the host may run its own path for the same call. The tokenizer, embeddings and token counter routes raise it from their own pre-checks, and `route_failure` raises it for a route error whose `declined()` answers, so a decline from deep inside a transformation reaches Python the same way. Never raise it for a failure after the request went out, the host would repeat a billed call. `a_core_decline_hands_the_call_back_to_python_before_the_public_mapping` and `route_failures_are_terminal` in `errors.rs` pin both sides + +Route hosts pass a terminal failure through the Python route module's `map_failure`, which runs `litellm.exception_type` and returns the public exception a Python caller expects. A decline must skip that step. Once wrapped it looks like a provider failure to the dispatcher and the host no longer falls back, so `route_failure` returns it before it calls `public` + +## Rules + +- Never format `Debug` of an error into a `PyErr` +- Never build a `PyErr` from a string at a call site. A failed argument check is a variant of a typed error, as `ProjectionError` in `coercion.rs` is, mapped once +- A Python exception raised inside the call stays that exception. Return it as is, never wrap it in a Rust message or stringify it into a new `ValueError`. The one policy exception is a caller-supplied callable, whose failure `callable::wrap_failure` reports as a `RuntimeError` with the original as `__cause__` +- A `panic!`, `unwrap` or `expect` inside a `#[pyfunction]` or `#[getter]` reaches Python as `pyo3_runtime.PanicException`, which is not an `Exception` subclass. Model the failure as a value instead +- When the mapping itself calls Python, as `setattr` does, propagate its failure with `?` rather than discarding it with `.ok()` + +## Testing + +Build the Rust error the way production builds it, map it, and assert the contract Python sees. `Python::initialize()` once, then `Python::attach` + +```rust +Python::initialize(); +Python::attach(|py| { + let failure = route_error_to_pyerr(litellm_http::Error::upstream(429, "slow down".into()).into()); + assert!(failure.is_instance_of::(py)); + let args: (u16, String) = failure.value(py).getattr("args").unwrap().extract().unwrap(); + assert_eq!(args, (429, "slow down".into())); +}); +``` + +One `#[case]` per fault class, asserting `is_instance_of` for the class and `extract` for the payload. The Python side of the contract is tested in `tests/unit/rust_bridge`, through `failures.py` diff --git a/litellm-rust/.agents/skills/rust-error/references/recipes.md b/litellm-rust/.agents/skills/rust-error/references/recipes.md new file mode 100644 index 00000000000..80ba1512f77 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/recipes.md @@ -0,0 +1,64 @@ +# Recipes + +## A new failure in litellm_http + +Say media downloads must reject a `Content-Type` outside an allow-list + +1. Add a `Repr` variant with typed facts and the message, `#[error("media type {0:?} is not allowed")] DisallowedMediaType(String)` +2. Return it where the check fails, `return Err(Repr::DisallowedMediaType(content_type.to_owned()).into())` +3. Add an arm to every predicate's exhaustive `match`, the compiler lists them, the way its neighbor `DownloadTooLarge` answers. A media download happens before the provider call, so it is not in `dispatched()`. The caller chose the URL, so it is in `is_request()` +4. Leave both boundaries alone. With those answers the gateway returns 400 and Python raises `ValueError` +5. Test it with a mock server answering `Content-Type: application/x-msdownload` and assert `error.is_request()` + +## A provider-specific failure + +A second OCR provider needs a missing-key error whose text names its environment variable and how to set it + +1. Look for an existing provider-neutral variant whose template fits. `litellm_auth::Error::MissingApiKey { provider, environment_variable }` renders `Missing {provider} API Key - Set ...` +2. If none fits, find what varies between providers with the same failure, here the environment variable and the guidance after the dash, and add one neutral variant with those as fields, `#[error("Missing {environment_variable} - {guidance}")] MissingCredential { environment_variable: &'static str, guidance: &'static str }` +3. Pass the provider's values from its transformation. The next provider passes its own and adds no variant, and `is_missing_credential()` lists `MissingCredential { .. }` once + +## Wrapping a new lower crate in litellm_inference + +Say routes start calling `litellm_cache` and its errors must surface + +1. Add one variant, `#[error(transparent)] Cache(#[from] litellm_cache::Error)` +2. Extend every predicate's `match` with an arm that calls through to the cache error's own predicates, adding a predicate there if it has none +3. Do not copy the cache error's variants into `litellm_inference::Error` + +## A foreign error that needs context + +Reading a file fails with `io::Error`, and the message must name the file + +1. A variant with the context and the cause, `Read { path: PathBuf, source: io::Error }` with `#[error("could not read {}", path.display())]` +2. Build it with `map_err(|source| Repr::Read { path: path.clone(), source })`, not a `From` +3. If the error must be `Clone`, the crate's error is opaque or the field is `litellm_error::Source` + +## Migrating a stringly error + +Call sites look like `Error::Network(error.to_string())`, `Error::Request(format!(...))` or `.map_err(|_| Error::Unavailable)` + +1. List every call site of the string variant and group them by what actually failed +2. Each group becomes a variant holding the typed cause or facts. Groups no caller tells apart share one +3. Where a foreign error was stringified with classification logic around it, move that logic into the `From` and turn the call sites into `?` +4. Delete the string variant. If it cannot be deleted yet, nothing new may construct it +5. Re-run the boundary tests, since messages a client sees may change + +## Moving a crate to source chains + +A variant reads `#[error("invalid JSON: {0}")] Json(#[from] serde_json::Error)` or stores `cause` and renders `{cause}` + +1. Drop the cause from the template, `#[error("invalid JSON")]` +2. Make the cause the source, `#[from]`, `#[source]` or a field named `source` +3. Find what renders this error for someone else and make it use `litellm_error::Report` +4. Update tests that asserted the cause in `to_string()` to assert it in the report + +## Making a public enum opaque + +Downstream crates match nested variants, or a variant needs a non-`Clone` cause + +1. Rename the enum to `pub(crate) enum Repr` and add `pub struct Error(Arc)` with `#[error(transparent)]` and `From` +2. For each downstream `match`, find the question it asks and add a predicate. Replace the `match` with the predicate +3. For each downstream construction, add a constructor taking typed data, or move the failure into this crate +4. Swap test assertions to predicates and add the `#[cfg(test)] repr()` accessor for in-crate tests +5. Implement `PartialEq` with `Arc::ptr_eq` only if a wrapping enum's derive needs it diff --git a/litellm-rust/.agents/skills/rust-error/references/shape.md b/litellm-rust/.agents/skills/rust-error/references/shape.md new file mode 100644 index 00000000000..4d51c2cee6b --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/shape.md @@ -0,0 +1,105 @@ +# Error shape + +## One top-level Error per crate + +A crate's errors live in `src/error.rs` and are re-exported from `lib.rs` as `Error`. A caller of any public function in the crate should only ever see `Result<_, crate::Error>`, so one `?` and one set of predicates cover the whole crate + +This deliberately departs from the advice to put an error type next to each unit of fallibility ([Sabrina Jewson, "Modular Errors"](https://sabrinajewson.org/blog/errors)). Our callers do not branch on which function failed, they ask the predicates in [classification](classification.md), and the boundaries map one type per crate. Splitting by function would multiply those predicates and mappings without giving any caller a decision it lacks. Layering, the other half of that advice, is kept: each crate nests the one below it, see [conversions](conversions.md) + +That one type takes one of two forms + +| Form | Use when | Example | +|---|---|---| +| Public enum | Callers outside the crate match variants, and every variant holds our own data or another crate's `Error` that is already `Clone` and comparable | `litellm_inference::Error`, `litellm_auth::Error` | +| Opaque struct over a private `Repr` | The variants hold third-party causes, or the variants must be free to change without breaking callers | `litellm_http::Error` | + +Both forms are still one top-level `Error`. The opaque form just moves the enum behind `Repr`, where it is `pub(crate)` + +```rust +// public enum +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum Error { + #[error("missing required field: {0}")] + MissingField(&'static str), + #[error(transparent)] + Http(litellm_http::Error), +} + +// opaque struct +#[derive(Clone, Debug, thiserror::Error)] +#[error(transparent)] +pub struct Error(Arc); + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Repr { + #[error("could not reach the provider")] + Connect(#[source] reqwest::Error), + #[error("upstream request failed with status {status}: {body}")] + Status { status: u16, body: String }, +} +``` + +Start with the public enum. Move to opaque when a variant needs a non-`Clone` third-party cause, or when downstream crates keep matching variants in ways that duplicate classification. `reqwest::Error`, `std::io::Error` and `serde_json::Error` are built the opaque way + +In the opaque form other crates cannot name a variant, so they cannot build an error with their own message. Public constructors exist only for errors another crate genuinely produces, take typed data and never a message, as in `Error::upstream(status, body)`, `Error::computed_header(name)`, `Error::unsignable_body()` and `Error::signature(cause)` + +`Arc` makes cloning a refcount bump, so `Repr` never has to be `Clone`. `Box` is enough when the goal is only a small `Result`. `Arc` is for when callers clone + +Every crate here is internal to the workspace, so do not mark errors or variants `#[non_exhaustive]`. It only forces a wildcard arm on the other workspace crates and turns off the exhaustiveness check that makes a new variant visit every `match`. When the variants must be free to change, use the opaque form + +## What else may sit next to Error + +These are not a second top-level error and need no folding in + +- Leaf data types whose fields are typed and whose message is in `#[error(...)]`, such as `litellm_http::request::HeaderError`. A wrapping crate nests it as one variant +- Plain data used by predicates, such as `litellm_http::TlsSource` +- A private fault enum one variant carries, such as `PemFault` inside `Repr::InvalidPem` + +## Variants are failure modes + +A public variant is something a caller handles differently: a phase, a status code, retry versus no retry, fault on the request versus the environment, or a message Python parity pins exactly. Failures no caller tells apart share one variant and differ in their data + +```rust +// one variant per provider, so every new provider adds a variant and a match arm +#[error("Missing Azure API Base - set AZURE_API_BASE")] +MissingAzureApiBase, +#[error("Missing Vertex API Base - set VERTEXAI_API_BASE")] +MissingVertexApiBase, + +// one failure mode, the provider and its guidance are data +#[error("Missing {provider} API Base - {guidance}")] +MissingApiBase { provider: &'static str, guidance: &'static str }, +``` + +An exact parity message does not justify a provider-specific variant when caller-supplied data can produce the same text + +Inside a private `Repr` a variant costs no API, so one per distinct reason is fine when each one carries a different message. When several only differ in a word and share every predicate answer, prefer one variant with a fieldless reason enum, so predicates list one arm instead of four, as `Repr::Handshake(HandshakeFault)` does + +```rust +#[error("WebSocket handshake returned {0}")] +Handshake(HandshakeFault), + +#[derive(Debug, thiserror::Error)] +pub(crate) enum HandshakeFault { + #[error("a wrong Sec-WebSocket-Accept")] + Accept, + #[error("an invalid Upgrade header")] + Upgrade, + #[error("an invalid Connection header")] + Connection, + #[error("an invalid subprotocol")] + Subprotocol, +} +``` + +## When to split + +Split into a separate type only when different functions fail in disjoint ways and different callers see them, so no caller matches variants its function can never return. A type per function overdoes it + +Name a split type after what went wrong, not after the function that returns it. `ConnectError` reads better than `OpenSessionError`, and a unit struct is fine for a single failure mode + +A split type still converts into the crate's `Error` when a caller further up needs one `Result` type + +`litellm-llms` has one recorded exception, the OCR error in `src/base_llm/ocr/error.rs`, kept apart until OCR folds into `litellm_llms::Error`, see `crates/inference-ocr/AGENTS.md`. Do not add another + +A utility crate whose modules are unrelated, such as `core-utils`, has no crate-wide failure, so each module that can fail owns one error next to its code instead of a `src/error.rs`. The rest of this skill applies to each of them diff --git a/litellm-rust/.agents/skills/rust-error/references/testing.md b/litellm-rust/.agents/skills/rust-error/references/testing.md new file mode 100644 index 00000000000..4aa22b4a293 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-error/references/testing.md @@ -0,0 +1,37 @@ +# Testing errors + +Assert behavior through predicates and accessors, on errors produced the way production produces them + +```rust +// guesses the message and the variant, and cannot be written once the type is opaque +assert_eq!(error, Error::Transport(TransportError::Http { status: 429, body: "slow down".into() })); + +// asserts what callers rely on +assert_eq!(error.upstream_response(), Some((429, "slow down"))); +assert!(error.dispatched()); +assert!(!error.is_request()); +``` + +## Producing the error + +In order of preference + +1. A real failing call. Bind a `TcpListener`, drop it, then connect to its address for a connect failure. Accept and never answer for a timeout +2. A real cause, such as `tungstenite::Error::ConnectionClosed` or a `serde_json::from_str("{")` error, passed through our `From` +3. A public constructor, such as `Error::upstream(500, "boom".into())` + +Inside the owning crate a test may match the private representation through the `#[cfg(test)] fn repr(&self) -> &Repr` accessor in `crates/http/src/error.rs`, as the `media.rs` and `tls.rs` tests do. Anywhere else, reach the cause the test expects with `litellm_error::find::(&error)`. A transparent layer is no node of the chain, so `find` reaches foreign causes underneath it but never the nested error of ours itself. Ask that error's predicates through the outer type instead, which is what callers do too + +## What to cover + +| Property | Test shape | +|---|---| +| classification | one `#[case]` per input class, asserting the predicate | +| redaction | build from a URL holding `api_key=secret`, assert the message lacks `secret` and the path | +| cause kept | `litellm_error::find::(&error)` returns the expected cause | +| rendered once | the cause's text appears exactly once in `litellm_error::Report::new(&error).to_string()` and not in `to_string()` | +| boundary mapping | the gateway status and body, or the Python exception type and its `status` attribute | + +Assert exact message text only where it is a contract, such as a message Python parity pins. Everywhere else assert that the message contains the facts it must carry. Matching a fieldless variant with `matches!` is fine, full `assert_eq!` on an error that holds a cause is not, since causes compare by pointer + +A test that returns `Result` reports only `Debug` of the error on failure. Prefer `expect` with a short reason, or `unwrap_err` followed by assertions on the predicates diff --git a/litellm-rust/.agents/skills/rust-gateway/SKILL.md b/litellm-rust/.agents/skills/rust-gateway/SKILL.md new file mode 100644 index 00000000000..66b93ee128d --- /dev/null +++ b/litellm-rust/.agents/skills/rust-gateway/SKILL.md @@ -0,0 +1,20 @@ +--- +name: rust-gateway +description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test +--- + +# Rust gateway + +`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching + +- [surfaces](references/surfaces.md) + - Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose + - The error type string is derived from the status and provider bodies pass through bounded +- [requests](references/requests.md) + - Extractors reject with the crate's `Error`, handlers take the body as a `Result` when the error needs the request id + - Every alias of a path is registered together, a known but unserved path goes to one handler + - Streams relay through `litellm_host_http::Sse`, with the error frame rendered by the crate's `Error` +- [testing](references/testing.md) + - Build the router through `tests/support` and drive it with `oneshot`, assert the fields the surface promises + +How an error is classified before it reaches a response is the [`rust-error`](../rust-error/SKILL.md) skill, in particular its [gateway](../rust-error/references/gateway.md) page. Axum's own conventions are in its [docs](https://docs.rs/axum/latest/axum/) diff --git a/litellm-rust/.agents/skills/rust-gateway/references/requests.md b/litellm-rust/.agents/skills/rust-gateway/references/requests.md new file mode 100644 index 00000000000..395f29f5c64 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-gateway/references/requests.md @@ -0,0 +1,34 @@ +# Requests + +## Extractors reject with the crate's Error + +A custom extractor sets `type Rejection = Error`, so a bad body gets the surface's envelope rather than axum's plain-text rejection. Never take axum's own `Json`, `Query` or `Path` in a handler, their rejection is plain text. The rejection is the variant's source, never its string. `JsonObject` and `InferenceBody` in `crates/gateway-inference/src/request.rs` are the pattern + +```rust +impl FromRequest for JsonObject { + type Rejection = Error; + + async fn from_request(request: Request, state: &S) -> Result { + let body = Bytes::from_request(request, state).await.map_err(|error| { + if error.status() == StatusCode::PAYLOAD_TOO_LARGE { + Error::BodyTooLarge + } else { + BodyError::Read(error).into() + } + })?; + Ok(Self(object(&body)?)) + } +} +``` + +When the handler needs something from the request in the error body, such as the request id on `/v1/messages`, take the extractor as `Result` so the handler maps the rejection itself, as `messages::create` does. `Infallible` is the right rejection for an extractor that only reads an optional header, as `RequestId` does + +## Paths + +Every alias the Python proxy registers for a route is registered together in `crates/gateway-inference/src/lib.rs`, such as the bare path, its `/v1` form and the `/engines/{model}` and `/openai/deployments/{model}` forms, otherwise one alias silently behaves differently from the others. A path the gateway knows but does not serve, such as `/v1/completions`, is routed to `request::unsupported`, which returns `Error::Unsupported(path)` so the client gets a 501 in the surface's envelope rather than a bare 404 + +`crates/gateway` merges the mounted routers and adds the shared inbound middleware as `route_layer`s, so authentication runs only for matched routes. Endpoint paths and parsing stay in the mounted crate, see the two crates' `AGENTS.md` + +## Streaming + +Once the response headers are sent the status cannot change. A streaming handler builds `litellm_host_http::Sse::new(Json, |error| Bytes::from(Error::from(error).sse_frame()))`, where the closure renders a mid-stream failure as an `event: error` frame with the same status and message a JSON failure on that surface carries, `sse_frame()` for Anthropic and `responses_sse_frame()` for the Responses API's `error` event. `litellm_host_http::Error` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the crate's `From` maps the second to `Internal` diff --git a/litellm-rust/.agents/skills/rust-gateway/references/surfaces.md b/litellm-rust/.agents/skills/rust-gateway/references/surfaces.md new file mode 100644 index 00000000000..08d4bbca4ea --- /dev/null +++ b/litellm-rust/.agents/skills/rust-gateway/references/surfaces.md @@ -0,0 +1,20 @@ +# Envelopes follow the API surface + +Each surface mirrors what the Python proxy returns there, so the shapes differ on purpose + +| Surface | Envelope | Where | +|---|---|---| +| `/v1/messages` | Anthropic, `{"type": "error", "error": {"type", "message"}, "request_id"}` | `Error::body()` and `sse_frame()` in `crates/gateway-inference/src/error.rs` | +| OpenAI-shaped routes | `{"error": {"message", "type", "param", "code"}}` | `Error::openai_response()` in the same file | +| `/v1/responses` stream | the Responses `error` event, `{"type": "error", "code", "message", "param"}` | `Error::responses_sse_frame()` in the same file | +| MCP | `{"detail": message}` | `crates/gateway-mcp/src/error.rs` | +| UI | `{"error": {"message"}}`, plus `Retry-After` on 429 | `crates/gateway-ui/src/error.rs` | +| auth rejections | plain text, the top-level message only | `crates/gateway-auth/src/error.rs` | + +Every surface takes its status from one `status()` on its crate's `Error` and renders the message of anything that is not a provider answer with `litellm_error::Report`. In `gateway-inference` the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else, and the OpenAI envelope always wraps, as the Python proxy does. A provider body is bounded by `litellm_http::Error::upstream` when the error is built + +Two surfaces show less than the chain on purpose, and a test in each pins it. An auth rejection goes to a caller who has not authenticated, so it never shows what a secret manager or key store said. An MCP upstream or authorization failure shows only its own layer, because rmcp renders the transport error, URL included, into its own `Display`, and an authorizer's cause is the same backend detail auth rejections hide. The inference surfaces show only the top layer of an auth rejection for the same reason + +`/v1/messages` adds the request id to the envelope when the provider did not set one, which is why its handler takes the body as `Result` and renders the rejection itself, see [requests](requests.md) + +A new surface gets its own renderer next to the others, named after the shape it produces, and its status still comes from `Error::status()`. Picking the status is the [`rust-error`](../../rust-error/references/gateway.md) skill's concern, not the surface's diff --git a/litellm-rust/.agents/skills/rust-gateway/references/testing.md b/litellm-rust/.agents/skills/rust-gateway/references/testing.md new file mode 100644 index 00000000000..10c1f01c472 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-gateway/references/testing.md @@ -0,0 +1,9 @@ +# Testing + +Test HTTP contracts in the crate that owns the endpoint: status codes, forwarded headers, error envelopes and streaming behavior. Inference and provider behavior is tested in the `inference-` and `llms` crates, and mounting and server lifecycle in `crates/gateway` + +`crates/gateway-inference/tests/support/mod.rs` builds a `Router` with injected deployments, pools and settings through `app()` and its variants, and `post()` drives it with `tower::ServiceExt::oneshot`. `json()` reads the response body back. Use these instead of building a router or a client by hand in a test + +Assert the status and the fields the surface promises, not the whole body. A test that pins a whole envelope breaks on every unrelated field, one that checks `status`, `error.type` and the presence of `request_id` catches the regressions that matter + +`status()` and `body()` of the error type are pure, so a classification case belongs in the error module's unit tests, see the [`rust-error`](../../rust-error/references/gateway.md) skill. Router tests cover what only the full path shows, such as an extractor rejection arriving in the right envelope or an SSE error frame after the headers went out diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index e62a9463cb5..9e7a55eb232 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -4,6 +4,8 @@ For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.a For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) +For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md) + For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all` @@ -24,13 +26,3 @@ Use [`#[rstest]`](https://docs.rs/rstest/latest/rstest/attr.rstest.html) for new - Never loop over inputs (`for`, `.iter().for_each`, `.all`) inside a test body. Give each input its own `#[case::name(...)]`, or use `#[values(...)]` for a cross product - Exception: a test pinning a Rust table against a repo-owned data file (for example `include_str!` of a JSON config) may iterate that file's entries - -## Error definitions - -- A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs` -- Put message templates in the variant's `#[error(...)]` declaration. Callers pass only the small typed arguments needed to fill them, never `Error::Variant(format!(...))` or a preformatted message. Keep the smallest set of neutral variants that callers need to distinguish. Different wording or providers do not justify new variants -- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each. A failure mode is something a caller handles differently (phase, status code, retry, a message Python parity pins exactly). Failures no caller tells apart share one variant and differ only in its message -- Keep shared error enums minimal and provider-neutral. Provider names, credential types, configuration fields, and setup guidance belong in caller-supplied data, not dedicated variants or hardcoded shared messages. Reuse a variant for the same failure mode across providers, such as `MissingApiBase { provider: "Azure", guidance: "..." }`. An exact parity message does not justify a provider-specific variant when caller-supplied context can preserve it -- Wrap a lower-level error as a variant with `#[from]` or `#[source]` instead of flattening it to a string -- Exception: split into separate types when different functions fail in disjoint ways, especially when different callers see them. A shared enum would force every caller to match variants its function can never return -- Name a split type after what went wrong (a unit struct is fine for a single failure mode), not after the function that returns it diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 63137949f38..9c45c5c4b1c 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2342,7 +2342,7 @@ dependencies = [ "hmac 0.13.0", "http 1.4.2", "jiff", - "reqwest 0.13.5", + "reqwest", "rustc_version", "rustls 0.23.42", "rustls-pki-types", @@ -2400,7 +2400,7 @@ dependencies = [ "pin-project", "prost", "prost-types", - "reqwest 0.13.5", + "reqwest", "rustc_version", "serde", "serde_json", @@ -2896,7 +2896,6 @@ dependencies = [ "tokio", "tokio-rustls 0.26.4", "tower-service", - "webpki-roots", ] [[package]] @@ -3386,10 +3385,13 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "bytes", + "futures-util", "litellm-auth-types", + "litellm-error", "litellm-http", "moka", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -3406,6 +3408,7 @@ dependencies = [ "azure_core", "azure_identity", "litellm-auth-types", + "litellm-error", "moka", "rstest", "serde_json", @@ -3424,6 +3427,7 @@ dependencies = [ "google-cloud-auth", "http 1.4.2", "litellm-auth-types", + "litellm-error", "moka", "rstest", "serde", @@ -3438,6 +3442,7 @@ dependencies = [ name = "litellm-auth-types" version = "0.1.0" dependencies = [ + "litellm-error", "rstest", "serde", "serde_json", @@ -3451,6 +3456,7 @@ dependencies = [ name = "litellm-cache" version = "0.1.0" dependencies = [ + "litellm-error", "rstest", "serde", "serde_json", @@ -3473,9 +3479,10 @@ dependencies = [ "litellm-cache-response", "litellm-cache-testing", "litellm-http", - "reqwest 0.12.28", + "reqwest", "rstest", "serde_json", + "thiserror 2.0.19", "tokio", "url", "wiremock", @@ -3494,6 +3501,7 @@ dependencies = [ "serde-pickle", "serde_json", "tempfile", + "thiserror 2.0.19", "tokio", ] @@ -3506,11 +3514,13 @@ dependencies = [ "litellm-auth-types", "litellm-cache", "litellm-cache-testing", + "litellm-error", "litellm-http", "percent-encoding", - "reqwest 0.12.28", + "reqwest", "rstest", "serde_json", + "thiserror 2.0.19", "tokio", "wiremock", ] @@ -3533,9 +3543,10 @@ dependencies = [ "futures-util", "litellm-cache", "litellm-cache-testing", + "litellm-error", "litellm-http", "qdrant-client", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -3553,11 +3564,13 @@ version = "0.1.0" dependencies = [ "litellm-cache", "litellm-cache-testing", + "litellm-error", "r2d2", "redis", "redis-test", "rstest", "serde_json", + "thiserror 2.0.19", "tokio", ] @@ -3610,8 +3623,9 @@ dependencies = [ "litellm-auth-aws", "litellm-cache", "litellm-cache-testing", + "litellm-error", "litellm-http", - "reqwest 0.12.28", + "reqwest", "rstest", "serde_json", "tokio", @@ -3632,11 +3646,13 @@ dependencies = [ "litellm-cache", "litellm-cache-redis", "litellm-cache-testing", + "litellm-error", "redis", "redis-test", "rstest", "serde_json", "sha2 0.10.9", + "thiserror 2.0.19", "tokio", "uuid", ] @@ -3646,6 +3662,7 @@ name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ "litellm-auth", + "litellm-error", "litellm-host", "litellm-host-python", "proptest", @@ -3653,6 +3670,7 @@ dependencies = [ "rstest", "serde_json", "strum", + "thiserror 2.0.19", ] [[package]] @@ -3668,6 +3686,7 @@ name = "litellm-config" version = "0.1.0" dependencies = [ "litellm-auth-types", + "litellm-error", "litellm-router-types", "rstest", "serde", @@ -3685,6 +3704,7 @@ name = "litellm-core-utils" version = "0.1.0" dependencies = [ "fancy-regex 0.19.2", + "litellm-error", "litellm-llms-types", "litellm-tracing", "rstest", @@ -3712,6 +3732,7 @@ dependencies = [ "criterion", "proptest", "rstest", + "thiserror 2.0.19", ] [[package]] @@ -3741,6 +3762,14 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-error" +version = "0.1.0" +dependencies = [ + "rstest", + "thiserror 2.0.19", +] + [[package]] name = "litellm-framer" version = "0.1.0" @@ -3768,6 +3797,7 @@ dependencies = [ "litellm-auth-types", "litellm-circuit-breaker", "litellm-config", + "litellm-error", "litellm-gateway-auth", "litellm-gateway-inference", "litellm-gateway-mcp", @@ -3801,6 +3831,7 @@ dependencies = [ "futures-util", "litellm-auth-types", "litellm-config", + "litellm-error", "litellm-secrets", "rstest", "serde", @@ -3824,6 +3855,7 @@ dependencies = [ "litellm-auth", "litellm-cache-memory", "litellm-cache-response", + "litellm-error", "litellm-gateway-auth", "litellm-host", "litellm-host-http", @@ -3843,6 +3875,7 @@ dependencies = [ "serde_json", "thiserror 2.0.19", "tokio", + "tokio-tungstenite", "tower", "wiremock", ] @@ -3853,6 +3886,7 @@ version = "0.1.0" dependencies = [ "base64 0.22.1", "litellm-auth-types", + "litellm-error", "litellm-gateway-auth", "rand 0.8.7", "rstest", @@ -3870,6 +3904,7 @@ dependencies = [ "http 1.4.2", "litellm-auth-types", "litellm-config", + "litellm-error", "litellm-secrets", "moka", "rmcp", @@ -3884,6 +3919,7 @@ dependencies = [ "tower", "url", "uuid", + "veil", ] [[package]] @@ -3893,9 +3929,11 @@ dependencies = [ "axum", "axum-login", "base64 0.22.1", + "futures-util", "governor", "jsonwebtoken", "litellm-auth-types", + "litellm-error", "litellm-gateway-auth", "rand 0.8.7", "rstest", @@ -3920,6 +3958,7 @@ dependencies = [ "litellm-coroutine", "rstest", "serde_json", + "thiserror 2.0.19", "tokio", ] @@ -3957,6 +3996,7 @@ dependencies = [ "bytes", "futures-util", "litellm-coroutine", + "litellm-error", "litellm-host", "parking_lot", "pyo3", @@ -3977,10 +4017,9 @@ dependencies = [ "http 1.4.2", "hyper-util", "litellm-core-utils", + "litellm-error", "rcgen", - "reqwest 0.12.28", - "reqwest 0.13.5", - "rmcp", + "reqwest", "rstest", "rustls 0.23.42", "rustls-native-certs", @@ -3990,8 +4029,10 @@ dependencies = [ "tempfile", "thiserror 2.0.19", "tokio", + "tokio-rustls 0.26.4", "tokio-tungstenite", "tracing", + "url", "veil", "webpki-roots", ] @@ -4007,13 +4048,14 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-error", "litellm-framer", "litellm-host", "litellm-http", "litellm-llms", "litellm-secrets", "litellm-tracing", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4034,6 +4076,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-error", "litellm-host", "litellm-host-native", "litellm-http", @@ -4061,6 +4104,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-error", "litellm-framer", "litellm-host", "litellm-host-native", @@ -4072,7 +4116,7 @@ dependencies = [ "litellm-router-types", "litellm-secrets", "litellm-tracing", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4092,6 +4136,7 @@ dependencies = [ "litellm-auth-aws", "litellm-auth-gcp", "litellm-core-utils", + "litellm-error", "litellm-host", "litellm-host-native", "litellm-http", @@ -4109,6 +4154,7 @@ dependencies = [ "time", "tokio", "tracing", + "url", "wiremock", ] @@ -4124,6 +4170,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", + "litellm-error", "litellm-host", "litellm-host-native", "litellm-http", @@ -4133,7 +4180,7 @@ dependencies = [ "litellm-llms-types", "litellm-secrets", "litellm-tracing", - "reqwest 0.12.28", + "reqwest", "rstest", "serde_json", "tokio", @@ -4161,6 +4208,7 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-core-utils", + "litellm-error", "litellm-http", "litellm-inference", "litellm-inference-testing", @@ -4189,6 +4237,7 @@ dependencies = [ "litellm-auth-azure", "litellm-auth-gcp", "litellm-core-utils", + "litellm-error", "litellm-framer", "litellm-host", "litellm-http", @@ -4196,7 +4245,7 @@ dependencies = [ "litellm-python-compat", "litellm-router-types", "litellm-secrets", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4229,6 +4278,7 @@ version = "0.1.0" dependencies = [ "indexmap 2.14.0", "jsonschema", + "litellm-error", "litellm-llms-types", "rstest", "schemars 1.2.2", @@ -4255,6 +4305,7 @@ dependencies = [ "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core-utils", + "litellm-error", "litellm-host", "litellm-host-python", "litellm-http", @@ -4278,13 +4329,14 @@ dependencies = [ "pyo3", "pyo3-async-runtimes", "redis", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", "serde_with", "sqlx", "strum", + "tempfile", "thiserror 2.0.19", "tokio", "tokio-tungstenite", @@ -4362,6 +4414,7 @@ dependencies = [ "google-cloud-auth", "google-cloud-kms-v1", "litellm-core-utils", + "litellm-error", "litellm-http", "litellm-python-compat", "litellm-secrets-aws", @@ -4371,7 +4424,7 @@ dependencies = [ "litellm-secrets-hashicorp", "litellm-secrets-types", "moka", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", @@ -4413,13 +4466,14 @@ dependencies = [ "litellm-http", "litellm-secrets-types", "percent-encoding", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "url", "veil", "wiremock", ] @@ -4436,13 +4490,14 @@ dependencies = [ "moka", "percent-encoding", "rcgen", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", "tempfile", "thiserror 2.0.19", "tokio", + "url", "veil", "wiremock", ] @@ -4463,12 +4518,13 @@ dependencies = [ "litellm-secrets-types", "moka", "percent-encoding", - "reqwest 0.12.28", + "reqwest", "rstest", "serde", "serde_json", "thiserror 2.0.19", "tokio", + "url", "veil", "wiremock", ] @@ -4478,7 +4534,10 @@ name = "litellm-secrets-hashicorp" version = "0.1.0" dependencies = [ "litellm-core-utils", + "litellm-error", + "litellm-http", "litellm-secrets-types", + "reqwest", "rstest", "rustify", "rustify_derive", @@ -4512,6 +4571,7 @@ name = "litellm-spend-clickhouse" version = "0.1.0" dependencies = [ "flate2", + "litellm-error", "litellm-http", "litellm-storage-clickhouse", "rstest", @@ -4532,6 +4592,7 @@ name = "litellm-storage-clickhouse" version = "0.1.0" dependencies = [ "flate2", + "litellm-error", "litellm-http", "rstest", "serde", @@ -4551,7 +4612,8 @@ version = "0.1.0" dependencies = [ "flate2", "futures-util", - "reqwest 0.12.28", + "litellm-http", + "reqwest", "rstest", "semver", "serde", @@ -4574,6 +4636,7 @@ dependencies = [ "criterion", "indexmap 2.14.0", "itoa", + "litellm-error", "litellm-token-counter-fast", "litellm-token-counter-huggingface", "litellm-token-counter-tiktoken", @@ -4590,6 +4653,7 @@ name = "litellm-token-counter-fast" version = "0.1.0" dependencies = [ "base64 0.22.1", + "litellm-error", "rand 0.8.7", "rstest", "rustc-hash", @@ -4615,6 +4679,7 @@ name = "litellm-token-counter-tiktoken" version = "0.1.0" dependencies = [ "base64 0.22.1", + "litellm-error", "once_cell", "rstest", "rustc-hash", @@ -5936,49 +6001,6 @@ version = "1.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" -[[package]] -name = "reqwest" -version = "0.12.28" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" -dependencies = [ - "base64 0.22.1", - "bytes", - "futures-core", - "futures-util", - "h2 0.4.15", - "http 1.4.2", - "http-body 1.1.0", - "http-body-util", - "hyper 1.10.1", - "hyper-rustls 0.27.9", - "hyper-util", - "js-sys", - "log", - "mime_guess", - "percent-encoding", - "pin-project-lite", - "quinn", - "rustls 0.23.42", - "rustls-pki-types", - "serde", - "serde_json", - "serde_urlencoded", - "sync_wrapper", - "tokio", - "tokio-rustls 0.26.4", - "tokio-util", - "tower", - "tower-http", - "tower-service", - "url", - "wasm-bindgen", - "wasm-bindgen-futures", - "wasm-streams 0.4.2", - "web-sys", - "webpki-roots", -] - [[package]] name = "reqwest" version = "0.13.5" @@ -6000,6 +6022,7 @@ dependencies = [ "js-sys", "log", "mime", + "mime_guess", "percent-encoding", "pin-project-lite", "quinn", @@ -6019,7 +6042,7 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams 0.5.0", + "wasm-streams", "web-sys", ] @@ -6066,7 +6089,7 @@ dependencies = [ "pin-project-lite", "process-wrap", "rand 0.10.2", - "reqwest 0.13.5", + "reqwest", "schemars 1.2.2", "serde", "serde_json", @@ -6189,7 +6212,7 @@ dependencies = [ "async-trait", "bytes", "http 1.4.2", - "reqwest 0.13.5", + "reqwest", "rustify_derive", "serde", "serde_json", @@ -7221,7 +7244,7 @@ dependencies = [ "memchr", "parse-display", "pin-project-lite", - "reqwest 0.13.5", + "reqwest", "serde", "serde_json", "serde_with", @@ -7470,11 +7493,7 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" dependencies = [ "futures-util", "log", - "rustls 0.23.42", - "rustls-native-certs", - "rustls-pki-types", "tokio", - "tokio-rustls 0.26.4", "tungstenite", ] @@ -7819,8 +7838,6 @@ dependencies = [ "httparse", "log", "rand 0.8.7", - "rustls 0.23.42", - "rustls-pki-types", "sha1 0.10.7", "thiserror 1.0.69", "utf-8", @@ -7860,7 +7877,7 @@ dependencies = [ "futures", "pin-project", "rand 0.10.2", - "reqwest 0.13.5", + "reqwest", "serde", "serde_json", "time", @@ -8074,7 +8091,7 @@ dependencies = [ "async-trait", "derive_builder", "http 1.4.2", - "reqwest 0.13.5", + "reqwest", "rustify", "rustify_derive", "serde", @@ -8221,19 +8238,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "wasm-streams" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" -dependencies = [ - "futures-util", - "js-sys", - "wasm-bindgen", - "wasm-bindgen-futures", - "web-sys", -] - [[package]] name = "wasm-streams" version = "0.5.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index dd8d8704370..b7b22cfd32c 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -10,6 +10,7 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-config = { path = "crates/config" } +litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } litellm-router-types = { path = "crates/router-types" } litellm-tracing = { path = "crates/tracing" } @@ -90,7 +91,7 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } rand = "0.8" macro_rules_attribute = "0.2.3" schemars = "1" -reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] } +reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] } redis = "1.8.0" redis-test = "1.0.4" testcontainers-modules = "0.15.0" @@ -100,10 +101,10 @@ rstest = "0.26.1" wiremock = "0.6.5" rstest_reuse = "0.7.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustls-native-certs = "0.8" rustify = "=0.7.0" rustify_derive = "=0.5.5" vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] } -rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] } @@ -118,7 +119,7 @@ thiserror = "2.0" tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } -tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } +tokio-tungstenite = { version = "0.24", default-features = false, features = ["handshake"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" flate2 = "1" diff --git a/litellm-rust/clippy.toml b/litellm-rust/clippy.toml index 18663cfdffb..075e24a7849 100644 --- a/litellm-rust/clippy.toml +++ b/litellm-rust/clippy.toml @@ -10,8 +10,11 @@ disallowed-methods = [ { path = "reqwest::Client::new", reason = "take litellm_http::Client from HttpClientPool" }, { path = "reqwest::Client::builder", reason = "HttpClientConfig owns client construction" }, { path = "reqwest::ClientBuilder::danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" }, + { path = "reqwest::ClientBuilder::tls_danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" }, + { path = "reqwest::ClientBuilder::tls_danger_accept_invalid_hostnames", reason = "set HttpClientConfig::verify instead" }, { path = "reqwest::ClientBuilder::identity", reason = "set HttpClientConfig::client_certificate instead" }, { path = "reqwest::ClientBuilder::use_preconfigured_tls", reason = "HttpClientConfig owns the TLS configuration" }, + { path = "reqwest::ClientBuilder::tls_backend_preconfigured", reason = "HttpClientConfig owns the TLS configuration" }, { path = "sqlx::query", reason = "use sqlx::query! or query_file! so the SQL is checked against the migrated schema" }, { path = "sqlx::query_as", reason = "use sqlx::query_as! or query_file_as! so the SQL is checked against the migrated schema" }, { path = "sqlx::query_scalar", reason = "use sqlx::query_scalar! so the SQL is checked against the migrated schema" }, diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 65b411ffaef..e4578777a08 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -7,7 +7,9 @@ repository.workspace = true [dependencies] litellm-auth-types.workspace = true +litellm-error.workspace = true litellm-http.workspace = true +reqwest.workspace = true moka = { workspace = true, features = ["sync"] } serde.workspace = true @@ -26,5 +28,6 @@ aws-smithy-runtime-api = "1.13.0" [dev-dependencies] rstest.workspace = true litellm-http = { workspace = true, features = ["test-support"] } -reqwest.workspace = true +bytes.workspace = true +futures-util.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 8fdbba21711..81aa4b068bc 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1,4 +1,3 @@ -use std::collections::BTreeMap; use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; @@ -12,16 +11,17 @@ use aws_sigv4::http_request::{ }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; use litellm_auth_types::AwsParams; -use super::Error; use super::constants::{ AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES, }; +use crate::error::{CredentialFlow, Error, SigningStage}; const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); @@ -271,7 +271,7 @@ impl AwsAuthService { provider .provide_credentials() .await - .map_err(|error| Error::AwsProfile(error.to_string())) + .map_err(Error::credentials(CredentialFlow::Profile)) } AwsAuthFlow::AssumeRole { role, session_name } => { if is_already_running_as_role(&role, &resolved).await? { @@ -287,7 +287,7 @@ impl AwsAuthService { let credentials = provider .provide_credentials() .await - .map_err(|error| Error::AwsDefaultChain(error.to_string()))?; + .map_err(Error::credentials(CredentialFlow::DefaultChain))?; self.set_cached_credentials( key, credentials.clone(), @@ -327,7 +327,7 @@ impl AwsAuthService { provider .provide_credentials() .await - .map_err(|error| Error::AwsAssumeRole(error.to_string())) + .map_err(Error::credentials(CredentialFlow::AssumeRole)) } AwsAuthFlow::WebIdentity { token, @@ -350,12 +350,14 @@ impl AwsAuthService { .web_identity_token(token) .send() .await - .map_err(|error| Error::AwsWebIdentity(error.to_string()))?; + .map_err(Error::credentials(CredentialFlow::WebIdentity))?; let credentials = response .credentials() - .ok_or(Error::AwsMissingWebIdentityCredentials)?; - let expiration = SystemTime::try_from(*credentials.expiration()) - .map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?; + .ok_or(Error::MissingWebIdentityCredentials)?; + let expiration = + SystemTime::try_from(*credentials.expiration()).map_err(|source| { + Error::WebIdentityExpiration(litellm_error::Source::new(source)) + })?; Ok(Credentials::new( credentials.access_key_id(), credentials.secret_access_key(), @@ -376,7 +378,7 @@ impl AwsAuthService { let credentials = provider .provide_credentials() .await - .map_err(|error| Error::AwsDefaultChain(error.to_string()))?; + .map_err(Error::credentials(CredentialFlow::DefaultChain))?; self.set_cached_credentials( key, credentials.clone(), @@ -435,12 +437,12 @@ fn default_session_name() -> String { /// client header cannot change the canonical request and invalidate the /// signature. Signing everything instead makes the request 403 on a header the /// caller supplied, on a deployment that works on the Python path. -pub fn aws_signature_headers(headers: &BTreeMap) -> BTreeMap { +pub fn aws_signature_headers(headers: &HeaderMap) -> HeaderMap { headers .iter() .filter(|(name, _)| { - let name = name.to_ascii_lowercase(); - AWS_SIGNED_HEADER_NAMES.contains(&name.as_str()) + let name = name.as_str(); + AWS_SIGNED_HEADER_NAMES.contains(&name) || name.starts_with("x-amz-") || name.starts_with("x-amzn-") }) @@ -456,45 +458,62 @@ pub fn is_sigv4_computed_header(name: &str) -> bool { SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str()) } -pub fn sign_post( +#[derive(Clone, Copy, Debug)] +pub struct SigningScope<'a> { + pub region: &'a str, + pub service: &'a str, + pub credentials: &'a Credentials, + pub time: SystemTime, +} + +pub fn sign_headers( + method: &str, url: &str, body: &[u8], - headers: &BTreeMap, - region: &str, - service: &str, - credentials: &Credentials, - signing_time: SystemTime, -) -> Result, Error> { - let identity: Identity = credentials.clone().into(); + headers: &HeaderMap, + scope: SigningScope<'_>, +) -> Result { + let identity: Identity = scope.credentials.clone().into(); let params = v4::SigningParams::builder() .identity(&identity) - .region(region) - .name(service) - .time(signing_time) + .region(scope.region) + .name(scope.service) + .time(scope.time) .settings(SigningSettings::default()) .build() .map(SigningParams::from) - .map_err(|error| Error::AwsSigningParameters(error.to_string()))?; + .map_err(Error::signing(SigningStage::Parameters))?; let header_refs = headers .iter() - .map(|(name, value)| (name.as_str(), value.as_str())); - let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) - .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; + .map(|(name, value)| { + value + .to_str() + .map(|value| (name.as_str(), value)) + .map_err(|_| Error::NonAsciiHeader { + name: name.to_string(), + }) + }) + .collect::, _>>()?; + let request = SignableRequest::new( + method, + url, + header_refs.into_iter(), + SignableBody::Bytes(body), + ) + .map_err(Error::signing(SigningStage::Request))?; let (instructions, _) = sign(request, ¶ms) - .map_err(|error| Error::AwsSigning(error.to_string()))? + .map_err(Error::signing(SigningStage::Signature))? .into_parts(); - Ok(instructions + instructions .headers() .map(|(name, value)| { - let normalized_name = match name { - "authorization" => "Authorization", - "x-amz-date" => "X-Amz-Date", - "x-amz-security-token" => "X-Amz-Security-Token", - _ => name, - }; - (normalized_name.to_string(), value.to_string()) + let name = HeaderName::from_bytes(name.as_bytes()) + .map_err(Error::signing(SigningStage::Output))?; + let value = + HeaderValue::from_str(value).map_err(Error::signing(SigningStage::Output))?; + Ok((name, value)) }) - .collect()) + .collect() } /// Model-id and region parsing shared by every Bedrock route. @@ -639,6 +658,7 @@ pub fn host_supplied_credentials(params: &AwsParams) -> Option { mod tests { use super::*; use crate::constants::{AWS_REGION, BEDROCK_SERVICE}; + use reqwest::header::CONTENT_TYPE; fn no_env(_: &str) -> Option { None @@ -731,12 +751,12 @@ mod tests { ); } - fn parity_inputs() -> (String, Vec, BTreeMap) { + fn parity_inputs() -> (String, Vec, HeaderMap) { ( "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" .to_string(), br#"{"input":"hello"}"#.to_vec(), - BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]), + HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]), ) } @@ -961,9 +981,9 @@ mod tests { // cannot change the canonical request. Signing it instead makes the // request 403 the moment anything on the wire rewrites or drops it. let (url, body, mut headers) = parity_inputs(); - headers.insert("x-request-id".to_string(), "abc-123".to_string()); - headers.insert("Accept-Encoding".to_string(), "gzip".to_string()); - headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string()); + headers.insert("x-request-id", HeaderValue::from_static("abc-123")); + headers.insert("Accept-Encoding", HeaderValue::from_static("gzip")); + headers.insert("x-amzn-trace-id", HeaderValue::from_static("Root=1-abc")); let signable = aws_signature_headers(&headers); assert!(!signable.contains_key("x-request-id")); @@ -979,18 +999,21 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &signable, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - SystemTime::UNIX_EPOCH, + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: SystemTime::UNIX_EPOCH, + }, ) .expect("signs"); - let authorization = signed - .get("Authorization") + let authorization = signed["Authorization"] + .to_str() .expect("carries an authorization header"); assert!( !authorization.contains("x-request-id"), @@ -1012,29 +1035,24 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &headers, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + }, ) .expect("golden signature"); + assert_eq!(signed["X-Amz-Date"], "20240102T030405Z"); + assert_eq!(signed["X-Amz-Security-Token"], "session-token"); assert_eq!( - signed.get("X-Amz-Date").map(String::as_str), - Some("20240102T030405Z") - ); - assert_eq!( - signed.get("X-Amz-Security-Token").map(String::as_str), - Some("session-token") - ); - assert_eq!( - signed.get("Authorization").map(String::as_str), - Some( - "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" - ) + signed["Authorization"], + "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" ); } @@ -1048,19 +1066,71 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &headers, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + }, ) .expect("signature"); assert!(!signed.contains_key("X-Amz-Security-Token")); } + #[test] + fn a_header_value_outside_visible_ascii_is_named_instead_of_signed() { + let (url, body, mut headers) = parity_inputs(); + headers.insert("x-amz-meta-note", HeaderValue::from_str("café").unwrap()); + let credentials = Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"); + + let error = sign_headers( + "POST", + &url, + &body, + &headers, + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: UNIX_EPOCH, + }, + ) + .unwrap_err(); + + assert_eq!( + error, + Error::NonAsciiHeader { + name: "x-amz-meta-note".into() + } + ); + } + + #[test] + fn the_shared_auth_error_nests_the_aws_failure_and_its_cause() { + let cause = std::io::Error::new(std::io::ErrorKind::NotFound, "profile not found"); + let error = + litellm_auth_types::Error::from(Error::credentials(CredentialFlow::Profile)(cause)); + + assert!(matches!( + litellm_error::find::(&error), + Some(Error::Credentials { + flow: CredentialFlow::Profile, + .. + }) + )); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.contains("could not load AWS profile credentials"), + "{report}" + ); + assert_eq!(report.matches("profile not found").count(), 1, "{report}"); + } + #[ignore] #[tokio::test] async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box> { @@ -1068,7 +1138,7 @@ mod tests { let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?; let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec(); let headers = - BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]); + HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]); let credentials = AwsAuthService::default() .resolve_credentials( AwsAuthConfig { @@ -1080,30 +1150,32 @@ mod tests { &no_env, ) .await?; - let client = litellm_http::Client::plain_for_test(); + let client = litellm_http::Client::for_test(); let mut failures = Vec::new(); for region in ["us-west-2", "us-east-1"] { let url = format!( "https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke" ); - let signed_headers = sign_post( + let signed_headers = sign_headers( + "POST", &url, &body, &headers, - region, - BEDROCK_SERVICE, - &credentials, - SystemTime::now(), + SigningScope { + region, + service: BEDROCK_SERVICE, + credentials: &credentials, + time: SystemTime::now(), + }, )?; - let mut request = client.post(&url).body(body.clone()); - for (name, value) in &headers { - request = request.header(name, value); - } - for (name, value) in signed_headers { - request = request.header(name, value); - } - let response = request.send().await?; + let response = client + .post(&url) + .headers(headers.clone()) + .headers(signed_headers) + .body(body.clone()) + .send() + .await?; let status = response.status(); let response_body = response.text().await?; let snippet: String = response_body.chars().take(240).collect(); diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index d4c6ae8cf6c..5f88fc8a171 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -1,46 +1,73 @@ -use thiserror::Error as ThisError; +use litellm_error::Source; -#[derive(Clone, Debug, ThisError, PartialEq, Eq)] +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { - #[error("AWS profile credentials failed: {0}")] - AwsProfile(String), - #[error("AWS default credentials failed: {0}")] - AwsDefaultChain(String), - #[error("AWS role credentials failed: {0}")] - AwsAssumeRole(String), - #[error("AWS web identity credentials failed: {0}")] - AwsWebIdentity(String), - #[error("AWS web identity expiration was invalid: {0}")] - AwsWebIdentityExpiration(String), - #[error("AWS signing parameters failed: {0}")] - AwsSigningParameters(String), - #[error("AWS signable request failed: {0}")] - AwsSignableRequest(String), - #[error("AWS request signing failed: {0}")] - AwsSigning(String), + #[error("could not load AWS {flow} credentials")] + Credentials { + flow: CredentialFlow, + #[source] + source: Source, + }, #[error("AWS web identity response had no credentials")] - AwsMissingWebIdentityCredentials, + MissingWebIdentityCredentials, + #[error("AWS web identity response had an invalid expiration")] + WebIdentityExpiration(#[source] Source), + #[error("header {name} has a value that is not visible ASCII")] + NonAsciiHeader { name: String }, + #[error("{stage}")] + Signing { + stage: SigningStage, + #[source] + source: Source, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum CredentialFlow { + #[error("profile")] + Profile, + #[error("default chain")] + DefaultChain, + #[error("assumed role")] + AssumeRole, + #[error("web identity")] + WebIdentity, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum SigningStage { + #[error("SigV4 signing parameters are invalid")] + Parameters, + #[error("SigV4 cannot sign this request")] + Request, + #[error("SigV4 could not compute the signature")] + Signature, + #[error("SigV4 produced an invalid header")] + Output, +} + +impl Error { + pub(crate) fn credentials( + flow: CredentialFlow, + ) -> impl FnOnce(E) -> Self { + move |source| Self::Credentials { + flow, + source: Source::new(source), + } + } + + pub(crate) fn signing( + stage: SigningStage, + ) -> impl FnOnce(E) -> Self { + move |source| Self::Signing { + stage, + source: Source::new(source), + } + } } impl From for litellm_auth_types::Error { fn from(error: Error) -> Self { - Self::ProviderAuthentication(error.to_string()) - } -} - -#[cfg(test)] -mod tests { - use super::Error; - - #[test] - fn converts_to_shared_auth_error_without_losing_context() { - let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into())); - - assert_eq!( - error, - litellm_auth_types::Error::ProviderAuthentication( - "AWS profile credentials failed: profile not found".into() - ) - ); + Self::provider(error) } } diff --git a/litellm-rust/crates/auth-aws/src/lib.rs b/litellm-rust/crates/auth-aws/src/lib.rs index 0fe0b390110..bd0f866371f 100644 --- a/litellm-rust/crates/auth-aws/src/lib.rs +++ b/litellm-rust/crates/auth-aws/src/lib.rs @@ -5,5 +5,5 @@ mod signer; pub use aws::*; pub use aws_credential_types::Credentials; -pub use error::Error; +pub use error::{CredentialFlow, Error, SigningStage}; pub use signer::SigV4Signer; diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index 3868fdd939b..dff917d7c13 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -1,11 +1,13 @@ -use std::{collections::BTreeMap, time::SystemTime}; +use std::time::SystemTime; + +use aws_credential_types::Credentials; +use reqwest::{Request, header::HeaderMap}; use crate::{ - AwsAuthService, AwsCredentialSource, Error, aws_signature_headers, is_sigv4_computed_header, - sign_post, + AwsAuthService, AwsCredentialSource, Error, SigningScope, aws_signature_headers, + is_sigv4_computed_header, sign_headers, }; -use aws_credential_types::Credentials; -use litellm_http::outbound::{RequestSigner, UnsignedRequest}; +use litellm_http::outbound::RequestSigner; #[derive(Clone, Debug)] pub struct SigV4Signer { @@ -45,41 +47,54 @@ impl SigV4Signer { } impl RequestSigner for SigV4Signer { - fn sign( - &self, - request: UnsignedRequest<'_>, - ) -> Result, litellm_http::Error> { - if let Some((name, _)) = request - .headers - .iter() - .find(|(name, _)| is_sigv4_computed_header(name)) + fn sign(&self, request: &Request) -> Result { + if let Some(name) = request + .headers() + .keys() + .find(|name| is_sigv4_computed_header(name.as_str())) { - return Err(litellm_http::Error::ComputedHeader(name.clone())); + return Err(litellm_http::Error::computed_header(name.as_str())); } - let headers: BTreeMap = request.headers.iter().cloned().collect(); - sign_post( - request.url, - request.body, - &aws_signature_headers(&headers), - &self.region, - self.service, - &self.credentials, - (self.clock)(), + let body = match request.body() { + None => &[], + Some(body) => body + .as_bytes() + .ok_or_else(litellm_http::Error::unsignable_body)?, + }; + sign_headers( + request.method().as_str(), + request.url().as_str(), + body, + &aws_signature_headers(request.headers()), + SigningScope { + region: &self.region, + service: self.service, + credentials: &self.credentials, + time: (self.clock)(), + }, ) - .map(|signature| signature.into_iter().collect()) - .map_err(|error| litellm_http::Error::Signature(error.to_string())) + .map_err(|error| match error { + Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error), + Error::Credentials { .. } + | Error::MissingWebIdentityCredentials + | Error::WebIdentityExpiration(_) + | Error::Signing { .. } => litellm_http::Error::signature(error), + }) } } #[cfg(test)] mod tests { + use reqwest::{Body, Method, Url}; + use serde_json::{Value, json as json_value}; use std::time::{Duration, UNIX_EPOCH}; - use litellm_http::outbound::OutboundRequest; - use serde_json::{Value, json}; + use litellm_http::outbound::{json, signed}; use super::*; + const URL: &str = "https://textract.us-east-1.amazonaws.com/"; + fn fixed_clock() -> SystemTime { UNIX_EPOCH + Duration::from_secs(1_700_000_000) } @@ -93,82 +108,121 @@ mod tests { .with_clock(fixed_clock) } + fn target() -> (String, String) { + ("X-Amz-Target".into(), "Textract.DetectDocumentText".into()) + } + fn authorization(body: &Value, service: &'static str) -> String { - OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], - body, - None, + signed( + json(URL, &[target()], body, None).unwrap(), &signer(service), ) .unwrap() - .header("Authorization") - .unwrap() - .to_string() + .headers()["authorization"] + .to_str() + .unwrap() + .to_string() } #[test] fn the_signature_verifies_against_the_bytes_that_are_sent() { - let sent = OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], - &json!({"Document": {"Bytes": "aGk="}}), - None, + let sent = signed( + json( + URL, + &[target()], + &json_value!({"Document": {"Bytes": "aGk="}}), + None, + ) + .unwrap(), &signer("textract"), ) .unwrap(); - let unsigned: BTreeMap = sent + let unsigned: HeaderMap = sent .headers() .iter() - .filter(|(name, _)| !is_sigv4_computed_header(name)) - .cloned() + .filter(|(name, _)| !is_sigv4_computed_header(name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) .collect(); - let recomputed = sign_post( - sent.url(), - sent.body(), + let recomputed = sign_headers( + "POST", + sent.url().as_str(), + sent.body().unwrap().as_bytes().unwrap(), &aws_signature_headers(&unsigned), - "us-east-1", - "textract", - &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), - fixed_clock(), + SigningScope { + region: "us-east-1", + service: "textract", + credentials: &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), + time: fixed_clock(), + }, ) .unwrap(); - assert_eq!( - sent.header("Authorization"), - Some(recomputed["Authorization"].as_str()) - ); + assert_eq!(sent.headers()["authorization"], recomputed["authorization"]); } #[test] fn the_signature_depends_on_the_body_and_the_service() { - let original = authorization(&json!({"text": "card 4111"}), "textract"); + let original = authorization(&json_value!({"text": "card 4111"}), "textract"); assert_ne!( original, - authorization(&json!({"text": "card [REDACTED]"}), "textract") + authorization(&json_value!({"text": "card [REDACTED]"}), "textract") ); assert_ne!( original, - authorization(&json!({"text": "card 4111"}), "bedrock") + authorization(&json_value!({"text": "card 4111"}), "bedrock") ); assert!(original.contains("/us-east-1/textract/aws4_request")); } #[test] fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() { - let error = OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("authorization".into(), "Bearer caller".into())], - &json!({}), - None, + let error = signed( + json( + URL, + &[("authorization".into(), "Bearer caller".into())], + &json_value!({}), + None, + ) + .unwrap(), &signer("textract"), ) .unwrap_err(); - assert_eq!( - error, - litellm_http::Error::ComputedHeader("authorization".into()) - ); + assert!(error.is_computed_header()); + assert!(error.to_string().contains("authorization"), "{error}"); + } + + #[test] + fn a_forwarded_header_that_cannot_be_signed_is_the_callers_fault() { + let error = signed( + json( + URL, + &[("x-amz-meta-note".into(), "café".into())], + &json_value!({}), + None, + ) + .unwrap(), + &signer("textract"), + ) + .unwrap_err(); + + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("x-amz-meta-note"), "{report}"); + assert!(error.is_request()); + assert!(!error.dispatched()); + } + + #[test] + fn a_streaming_body_cannot_be_signed() { + let mut request = Request::new(Method::POST, Url::parse(URL).unwrap()); + *request.body_mut() = Some(Body::wrap_stream(futures_util::stream::once(async { + Ok::<_, std::io::Error>(bytes::Bytes::from_static(b"{}")) + }))); + + assert!(matches!( + signed(request, &signer("textract")), + Err(error) if !error.dispatched() && !error.is_request() + )); } } diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 1fd9d39d113..f99972f813b 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -18,5 +18,6 @@ azure_core = "1.0.0" azure_identity = { version = "1.0.0", features = ["tokio"] } [dev-dependencies] +litellm-error.workspace = true rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 64752162384..9fda67746d1 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -12,7 +12,7 @@ use azure_identity::{ }; use sha2::{Digest, Sha256}; -use litellm_auth_types::Error; +use litellm_auth_types::{Error, ErrorDetail}; use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; use super::credential_provider_cache::{ @@ -133,7 +133,9 @@ impl NativeAzureTokenAcquirer { let token = credential .get_token(&[scope.as_str()], None) .await - .map_err(|error| Error::CredentialAcquisition(error.to_string().into()))?; + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure token request", error)) + })?; let expires_on = u64::try_from(token.expires_on.unix_timestamp()) .ok() .map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds)); @@ -250,11 +252,8 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> { let Some(authority) = authority else { return Ok(()); }; - let url = url::Url::parse(authority.value()).map_err(|_| { - Error::InvalidConfiguration( - "Azure authority must be an HTTPS origin without credentials, query, or fragment" - .into(), - ) + let url = url::Url::parse(authority.value()).map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed("Azure authority parsing", error)) })?; if url.scheme() != "https" || url.host_str().is_none() @@ -264,10 +263,10 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> { || url.fragment().is_some() || !matches!(url.path(), "" | "/") { - return Err(Error::InvalidConfiguration( - "Azure authority must be an HTTPS origin without credentials, query, or fragment" - .into(), - )); + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "Azure authority".into(), + expected: "an HTTPS origin without credentials, query, or fragment", + })); } Ok(()) } @@ -377,7 +376,7 @@ fn trusted_source(sources: &[InputSource]) -> InputSource { fn mixed_sources() -> Result { Err(Error::InvalidConfiguration( - "request-controlled Azure auth inputs cannot be combined with host credentials".into(), + ErrorDetail::MixedInputSources { provider: "Azure" }, )) } @@ -444,7 +443,7 @@ fn build_credential( .map(|credential| credential as Arc), } .map_err(|error| { - Error::InvalidConfiguration(litellm_auth_types::ErrorDetail::failed( + Error::InvalidConfiguration(ErrorDetail::failed( "Azure credential initialization", error, )) @@ -662,12 +661,11 @@ mod tests { )) .unwrap_err(); - assert_eq!( - error, - litellm_auth_types::Error::InvalidConfiguration( + assert!( + litellm_error::Report::new(&error).to_string().contains( "request-controlled Azure auth inputs cannot be combined with host credentials" - .into() - ) + ), + "{error:?}" ); } @@ -695,12 +693,11 @@ mod tests { authority, )) .unwrap_err(); - assert_eq!( - error, - litellm_auth_types::Error::InvalidConfiguration( + assert!( + litellm_error::Report::new(&error).to_string().contains( "Azure authority must be an HTTPS origin without credentials, query, or fragment" - .into() - ) + ), + "{error:?}" ); } } diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 2142e22db50..25b71741342 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,8 +1,8 @@ -use litellm_auth_types::Error; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, SecretValue, Sourced, TokenProviderHandle, }; +use litellm_auth_types::{Error, ErrorDetail}; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -91,9 +91,9 @@ impl AzureAuthService { AzureCredentialPlan::Caller(caller) => { let credential = caller.acquire().await?; if credential.secret().expose().is_empty() { - return Err(Error::EmptyCallerCredential( - "Azure AD token provider returned an empty token", - )); + return Err(Error::EmptyCallerCredential { + caller: "Azure AD token provider", + }); } Ok(Some(Sourced::new(credential, InputSource::Deployment))) } @@ -106,11 +106,9 @@ impl AzureAuthService { } => { let assertion = resolve_reference(inputs, env_lookup, reference.value()) .await? - .ok_or_else(|| { - Error::CredentialAcquisition( - "Azure OIDC reference did not resolve to a value".into(), - ) - })?; + .ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "Azure OIDC reference", + }))?; let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion { tenant_id, client_id, @@ -142,7 +140,13 @@ impl AzureAuthService { Err(error) => failures.push(error), } } - Err(Error::CredentialChain(failures)) + match failures.pop() { + Some(last) => Err(Error::CredentialChain { + last: Box::new(last), + earlier: failures, + }), + None => Ok(None), + } } AzureCredentialPlan::Missing => Ok(None), } @@ -173,7 +177,12 @@ pub(crate) fn select_auth_plan( .map(|selector| Sourced::new(selector, value.source())) }) .transpose() - .map_err(|_| Error::InvalidConfiguration("invalid Azure credential selector".into()))?; + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed( + "Azure credential selector parsing", + error, + )) + })?; let federated_token_file = configured_string( &inputs.federated_token_file, AZURE_FEDERATED_TOKEN_FILE_ENV, @@ -263,9 +272,12 @@ fn select_native_plan( let selection_source = selected.source(); match selected.into_value() { - AzureCredentialType::ClientSecretCredential => Err(Error::InvalidConfiguration( - "ClientSecretCredential requires tenant_id, client_id, and client_secret".into(), - )), + AzureCredentialType::ClientSecretCredential => { + Err(Error::InvalidConfiguration(ErrorDetail::MissingSettings { + credential: "ClientSecretCredential", + settings: "tenant_id, client_id, and client_secret", + })) + } AzureCredentialType::WorkloadIdentityCredential => { Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( workload_request(tenant_id, client_id, federated_token_file, scope, authority)?, @@ -349,22 +361,21 @@ fn workload_request( authority: Option>, ) -> Result { Ok(NativeAzureRequest::WorkloadIdentity { - tenant_id: tenant_id.ok_or_else(|| { - Error::InvalidConfiguration("WorkloadIdentityCredential requires tenant_id".into()) - })?, - client_id: client_id.ok_or_else(|| { - Error::InvalidConfiguration("WorkloadIdentityCredential requires client_id".into()) - })?, - token_file_path: token_file_path.ok_or_else(|| { - Error::InvalidConfiguration( - "WorkloadIdentityCredential requires azure_federated_token_file".into(), - ) - })?, + tenant_id: tenant_id.ok_or(workload_requires("tenant_id"))?, + client_id: client_id.ok_or(workload_requires("client_id"))?, + token_file_path: token_file_path.ok_or(workload_requires("azure_federated_token_file"))?, scope, authority, }) } +fn workload_requires(settings: &'static str) -> Error { + Error::InvalidConfiguration(ErrorDetail::MissingSettings { + credential: "WorkloadIdentityCredential", + settings, + }) +} + fn configured_string( configured: &ConfigValue, environment_name: &str, @@ -410,11 +421,10 @@ async fn resolve_reference( .map_or(CredentialLookup::Missing, CredentialLookup::Found), CredentialRef::None => return Ok(None), CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => { - let resolver = inputs.credential_resolver.as_ref().ok_or_else(|| { - Error::InvalidConfiguration( - "credential reference requires a host credential resolver".into(), - ) - })?; + let resolver = inputs + .credential_resolver + .as_ref() + .ok_or(Error::InvalidConfiguration(ErrorDetail::MissingResolver))?; resolver.resolve(reference).await? } }; @@ -433,7 +443,7 @@ fn oidc_reference( let value = token.value().expose(); if token.source() == InputSource::Request && value.starts_with("oidc/") { return Err(Error::InvalidConfiguration( - "request-controlled Azure credential references are not allowed".into(), + ErrorDetail::RequestControlledReference { provider: "Azure" }, )); } if let Some(name) = value.strip_prefix("oidc/env/") { @@ -457,19 +467,17 @@ fn oidc_reference( } if value.starts_with("oidc/") { return Err(Error::InvalidConfiguration( - "unsupported OIDC reference".into(), + ErrorDetail::UnsupportedReference, )); } Ok(None) } -fn non_empty_reference(value: &str, kind: &str) -> Result { +fn non_empty_reference(value: &str, kind: &'static str) -> Result { if value.is_empty() { - return Err(Error::InvalidConfiguration( - litellm_auth_types::ErrorDetail::Empty { - subject: kind.into(), - }, - )); + return Err(Error::InvalidConfiguration(ErrorDetail::Empty { + subject: kind, + })); } Ok(value.to_string()) } @@ -518,7 +526,9 @@ mod tests { expires_on: None, }) } else { - Err(Error::CredentialAcquisition(kind.into())) + Err(Error::CredentialAcquisition( + litellm_auth_types::ErrorDetail::Unresolved { subject: kind }, + )) } }) } @@ -624,7 +634,11 @@ mod tests { ))) .expect_err("unsupported backend must fail validation"); - assert!(error.to_string().contains("unsupported OIDC reference")); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("unsupported OIDC reference") + ); } #[rstest::rstest] @@ -649,11 +663,11 @@ mod tests { }) .unwrap_err(); - assert_eq!( - error, - Error::InvalidConfiguration( - "request-controlled Azure credential references are not allowed".into() - ) + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("request-controlled Azure credential references are not allowed"), + "{error:?}" ); } @@ -715,7 +729,16 @@ mod tests { .await .unwrap_err(); - assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2)); + assert!(matches!(&error, Error::CredentialChain { earlier, .. } if earlier.len() == 1)); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.contains("managed-identity did not resolve"), + "{report}" + ); + assert!( + report.contains("developer-tools did not resolve"), + "{report}" + ); } #[derive(Debug)] @@ -763,7 +786,9 @@ mod tests { assert_eq!( error, - Error::EmptyCallerCredential("Azure AD token provider returned an empty token") + Error::EmptyCallerCredential { + caller: "Azure AD token provider" + } ); } } diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index bb4c022580a..7844b35ecc2 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -24,5 +24,6 @@ google-cloud-auth = { workspace = true, optional = true } http = { workspace = true, optional = true } [dev-dependencies] +litellm-error.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index e09a350c012..8ac1e0f54ba 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -73,13 +73,17 @@ impl VertexAuth { .or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV)); let project_id = get_vertex_ai_project(config, env_lookup); - if !has_authorization && static_token.is_none() { - let access = self.get_access_token(config, env_lookup).await?; - return Ok(VertexEnvironment { - headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, - project_id: project_id.unwrap_or(access.project_id), - }); - } + let static_token = match (has_authorization, static_token) { + (true, _) => None, + (false, Some(token)) => Some(token), + (false, None) => { + let access = self.get_access_token(config, env_lookup).await?; + return Ok(VertexEnvironment { + headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, + project_id: project_id.unwrap_or(access.project_id), + }); + } + }; let project_id = match project_id { Some(project_id) => project_id, @@ -90,14 +94,9 @@ impl VertexAuth { .await? } }; - let headers = if has_authorization { - headers - } else { - apply_credential( - headers, - static_token.as_deref().expect("static token was checked"), - CredentialPlacement::Bearer, - )? + let headers = match static_token { + Some(token) => apply_credential(headers, &token, CredentialPlacement::Bearer)?, + None => headers, }; Ok(VertexEnvironment { headers, @@ -234,14 +233,13 @@ enum CredentialCacheKey { Adc, } -const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \ - Set it to a readable file path, or to the credentials JSON itself."; - fn load_trusted_service_account(configured: &str) -> Result { if is_base64_encoded_json(configured) { - return Err(Error::InvalidConfiguration(ErrorDetail::Message( - BASE64_CREDENTIALS.into(), - ))); + return Err(Error::InvalidConfiguration( + ErrorDetail::Base64EncodedJson { + field: "vertex_credentials", + }, + )); } if names_a_credentials_file(configured) { CustomServiceAccount::from_file(configured) @@ -299,7 +297,8 @@ mod tests { else { panic!("{configured:?} should not load"); }; - assert!(error.to_string().contains(expected), "{error}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(expected), "{report}"); } #[rstest::rstest] @@ -321,7 +320,7 @@ mod tests { .await { Ok(_) => panic!("{configured:?} should not load"), - Err(error) => error.to_string(), + Err(error) => litellm_error::Report::new(&error).to_string(), } } diff --git a/litellm-rust/crates/auth-gcp/src/sdk.rs b/litellm-rust/crates/auth-gcp/src/sdk.rs index 566df291b7d..6cfeee64d20 100644 --- a/litellm-rust/crates/auth-gcp/src/sdk.rs +++ b/litellm-rust/crates/auth-gcp/src/sdk.rs @@ -55,7 +55,7 @@ impl CredentialsProvider for GoogleCredentials { entity_tag: EntityTag::new(), data, }) - .map_err(|_| CredentialsError::from_msg(false, "Google authentication failed")) + .map_err(|error| CredentialsError::from_source(false, error)) } async fn universe_domain(&self) -> Option { @@ -102,5 +102,9 @@ mod tests { ); let error = credentials.headers(Extensions::new()).await.unwrap_err(); assert!(!format!("{error:?}").contains("private")); + assert_eq!( + std::error::Error::source(&error).and_then(|cause| cause.downcast_ref::()), + Some(&Error::InvalidHeader) + ); } } diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml index b54c04f0229..400f2640221 100644 --- a/litellm-rust/crates/auth-types/Cargo.toml +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true serde.workspace = true serde_json.workspace = true subtle.workspace = true diff --git a/litellm-rust/crates/auth-types/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs index 94fb8a58467..4f69c6fffc8 100644 --- a/litellm-rust/crates/auth-types/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -86,9 +86,9 @@ impl CredentialPlan { Self::Caller(caller) => { let credential = caller.acquire().await?; if credential.secret().expose().is_empty() { - return Err(Error::EmptyCallerCredential( - "credential caller returned an empty credential", - )); + return Err(Error::EmptyCallerCredential { + caller: "credential caller", + }); } Ok(CredentialPlanResolution::Resolved(credential)) } @@ -104,8 +104,8 @@ mod tests { CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; - use crate::Error; use crate::SecretValue; + use crate::{Error, ErrorDetail}; #[derive(Debug)] struct HostResolver; @@ -150,9 +150,10 @@ mod tests { impl CredentialResolver for FailingResolver { fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { Box::pin(async { - Err(Error::CredentialAcquisition( - "host credential lookup failed".into(), - )) + Err(Error::CredentialAcquisition(ErrorDetail::failed( + "host credential lookup", + std::io::Error::from(std::io::ErrorKind::PermissionDenied), + ))) }) } } @@ -168,9 +169,14 @@ mod tests { .await .expect_err("acquisition errors cannot become fallback"); - assert_eq!( - error, - Error::CredentialAcquisition("host credential lookup failed".into()) + assert!( + matches!(error, Error::CredentialAcquisition(_)), + "{error:?}" + ); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("host credential lookup failed") ); } } diff --git a/litellm-rust/crates/auth-types/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs index ff17dcf3810..a8ba7b64e03 100644 --- a/litellm-rust/crates/auth-types/src/error.rs +++ b/litellm-rust/crates/auth-types/src/error.rs @@ -1,17 +1,23 @@ -use thiserror::Error as ThisError; +use litellm_error::{Report, Source}; -#[derive(Clone, Debug, ThisError, PartialEq, Eq)] +use crate::{ParamKind, ParamSpec}; + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { - #[error("invalid authentication configuration: {0}")] + #[error("invalid authentication configuration")] InvalidConfiguration(#[source] ErrorDetail), - #[error("credential acquisition failed: {0}")] + #[error("credential acquisition failed")] CredentialAcquisition(#[source] ErrorDetail), - #[error("credential caller failed: {0}")] - EmptyCallerCredential(&'static str), - #[error("{0}")] - ProviderAuthentication(String), - #[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::>().join("; "))] - CredentialChain(Vec), + #[error("{caller} returned an empty credential")] + EmptyCallerCredential { caller: &'static str }, + #[error("provider authentication failed")] + Provider(#[source] Source), + #[error("no credential in the chain succeeded{}", earlier_attempts(earlier))] + CredentialChain { + #[source] + last: Box, + earlier: Vec, + }, #[error( "Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable" )] @@ -22,7 +28,7 @@ pub enum Error { #[error("Missing {provider} {} - {}", .spec.setting, .spec.guidance())] MissingParam { provider: &'static str, - spec: &'static crate::ParamSpec, + spec: &'static ParamSpec, }, #[error("Missing {provider} API Base - {guidance}")] MissingApiBase { @@ -33,78 +39,282 @@ pub enum Error { InvalidHeader, } +fn earlier_attempts(earlier: &[Error]) -> String { + if earlier.is_empty() { + return String::new(); + } + let attempts = earlier + .iter() + .map(|failure| Report::new(failure).to_string()) + .collect::>() + .join("; "); + format!(" (earlier attempts: {attempts})") +} + +impl Error { + pub fn provider(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::Provider(Source::new(cause)) + } + + pub fn is_missing_credential(&self) -> bool { + match self { + Self::MissingApiKey { .. } => true, + Self::MissingParam { spec, .. } => match spec.kind { + ParamKind::Credential => true, + ParamKind::Setting => false, + }, + Self::InvalidConfiguration(_) + | Self::CredentialAcquisition(_) + | Self::EmptyCallerCredential { .. } + | Self::Provider(_) + | Self::CredentialChain { .. } + | Self::MissingApiBase { .. } + | Self::InvalidHeader => false, + } + } + + pub fn is_request(&self) -> bool { + match self { + Self::InvalidConfiguration(_) + | Self::EmptyCallerCredential { .. } + | Self::InvalidHeader => true, + Self::CredentialAcquisition(_) | Self::Provider(_) | Self::CredentialChain { .. } => { + false + } + Self::MissingParam { spec, .. } => match spec.kind { + ParamKind::Setting => true, + ParamKind::Credential => false, + }, + Self::MissingApiKey { .. } | Self::MissingApiBase { .. } => false, + } + } +} + #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ErrorDetail { - #[error("{0}")] - Message(String), #[error("{field} must be {expected}")] InvalidType { field: String, expected: &'static str, }, #[error("{subject} cannot be empty")] - Empty { subject: String }, + Empty { subject: &'static str }, #[error("credential header {0} already exists")] DuplicateHeader(&'static str), - #[error("{operation} failed: {source}")] + #[error("credential plan is not allowed by the provider auth policy")] + PlanNotAllowed, + #[error("{credential} requires {settings}")] + MissingSettings { + credential: &'static str, + settings: &'static str, + }, + #[error("credentials cannot be sent to a request-controlled {provider} endpoint")] + RequestControlledEndpoint { provider: &'static str }, + #[error("request-controlled {provider} credential references are not allowed")] + RequestControlledReference { provider: &'static str }, + #[error("request-controlled {provider} auth inputs cannot be combined with host credentials")] + MixedInputSources { provider: &'static str }, + #[error("unsupported OIDC reference")] + UnsupportedReference, + #[error("credential reference requires a host credential resolver")] + MissingResolver, + #[error("{subject} did not resolve to a value")] + Unresolved { subject: &'static str }, + #[error( + "The `{field}` value is base64-encoded JSON, not the JSON itself. Set it to a readable file path, or to the credentials JSON itself." + )] + Base64EncodedJson { field: &'static str }, + #[error("{operation} failed")] Failed { operation: &'static str, #[source] - source: ErrorSource, + cause: Source, }, } impl ErrorDetail { pub fn failed( operation: &'static str, - source: impl std::error::Error + Send + Sync + 'static, + cause: impl std::error::Error + Send + Sync + 'static, ) -> Self { Self::Failed { operation, - source: ErrorSource::new(source), + cause: Source::new(cause), } } } -impl From for ErrorDetail { - fn from(message: String) -> Self { - Self::Message(message) +#[cfg(test)] +mod tests { + use std::io; + + use litellm_error::{Report, Source}; + use rstest::rstest; + + use super::{Error, ErrorDetail}; + use crate::{AwsParams, VertexParams}; + + #[rstest] + #[case::api_key( + Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, + "Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable" + )] + #[case::another_api_key( + Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" }, + "Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable" + )] + #[case::api_base( + Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, + "Missing Example API Base - Pass api_base" + )] + #[case::another_api_base( + Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" }, + "Missing Custom API Base - Set CUSTOM_ENDPOINT" + )] + fn missing_inputs_keep_the_python_text(#[case] error: Error, #[case] expected: &str) { + assert_eq!(error.to_string(), expected); + } + + #[rstest] + #[case::configuration( + Error::InvalidConfiguration(ErrorDetail::Empty { subject: "credential selector" }), + "invalid authentication configuration: credential selector cannot be empty" + )] + #[case::acquisition( + Error::CredentialAcquisition(ErrorDetail::Unresolved { subject: "OIDC reference" }), + "credential acquisition failed: OIDC reference did not resolve to a value" + )] + #[case::empty_caller( + Error::EmptyCallerCredential { caller: "token provider" }, + "token provider returned an empty credential" + )] + fn the_report_names_every_layer_in_order(#[case] error: Error, #[case] expected: &str) { + assert!( + Report::new(&error).to_string().contains(expected), + "{}", + Report::new(&error) + ); + } + + #[rstest] + #[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed( + "test credential", + denied() + )))] + #[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed( + "test credential", + denied() + )))] + fn a_failure_keeps_its_cause_as_the_source_and_renders_it_once(#[case] error: Error) { + let cause = + std::iter::successors(std::error::Error::source(&error), |cause| cause.source()) + .find_map(|cause| cause.downcast_ref::()?.downcast_ref::()) + .expect("the io cause is in the source chain"); + assert_eq!(cause.kind(), io::ErrorKind::PermissionDenied); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches("access denied").count(), 1, "{report}"); + } + + #[rstest] + #[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed( + "test credential", + denied() + )))] + #[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed( + "test credential", + denied() + )))] + fn a_contextual_failure_leaves_its_cause_out_of_its_own_message(#[case] error: Error) { + assert!(!error.to_string().contains("access denied"), "{error}"); + assert!( + Report::new(&error) + .to_string() + .contains("test credential failed: access denied") + ); + } + + #[test] + fn a_provider_failure_keeps_the_provider_error_as_its_source() { + let error = Error::provider(denied()); + assert_eq!( + litellm_error::find::(&error).map(io::Error::kind), + Some(io::ErrorKind::PermissionDenied) + ); + assert!(!error.to_string().contains("access denied"), "{error}"); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches("access denied").count(), 1, "{report}"); + } + + #[test] + fn a_chain_keeps_its_last_failure_as_the_source_and_renders_every_attempt_once() { + let error = Error::CredentialChain { + last: Box::new(Error::provider(denied())), + earlier: vec![ + Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "OIDC reference", + }), + Error::EmptyCallerCredential { + caller: "token provider", + }, + ], + }; + + assert_eq!( + litellm_error::find::(&error).map(io::Error::kind), + Some(io::ErrorKind::PermissionDenied) + ); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches("access denied").count(), 1, "{report}"); + assert_eq!( + report.matches("OIDC reference did not resolve").count(), + 1, + "{report}" + ); + assert_eq!( + report + .matches("token provider returned an empty credential") + .count(), + 1, + "{report}" + ); + } + + fn denied() -> io::Error { + io::Error::new(io::ErrorKind::PermissionDenied, "access denied") + } + + #[rstest] + #[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, true)] + #[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), true)] + #[case::aws_secret_access_key(AwsParams::SECRET_ACCESS_KEY.missing("AWS"), true)] + #[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), false)] + #[case::aws_region(AwsParams::REGION.missing("AWS"), false)] + #[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)] + #[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), false)] + fn only_an_absent_credential_is_a_missing_credential( + #[case] error: Error, + #[case] expected: bool, + ) { + assert_eq!(error.is_missing_credential(), expected); + } + + #[rstest] + #[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), true)] + #[case::aws_region(AwsParams::REGION.missing("AWS"), true)] + #[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)] + #[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), false)] + #[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, false)] + #[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), true)] + #[case::invalid_header(Error::InvalidHeader, true)] + #[case::identity_provider_outage( + Error::CredentialAcquisition(ErrorDetail::MissingResolver), + false + )] + #[case::provider_sdk_failure(Error::provider(denied()), false)] + fn a_setting_the_caller_can_pass_is_theirs_to_fix_but_credentials_and_api_bases_are_not( + #[case] error: Error, + #[case] expected: bool, + ) { + assert_eq!(error.is_request(), expected); } } - -impl From<&str> for ErrorDetail { - fn from(message: &str) -> Self { - Self::Message(message.into()) - } -} - -#[derive(Clone, Debug)] -pub struct ErrorSource(std::sync::Arc); - -impl std::ops::Deref for ErrorSource { - type Target = dyn std::error::Error + Send + Sync; - - fn deref(&self) -> &Self::Target { - self.0.as_ref() - } -} - -impl std::fmt::Display for ErrorSource { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(&self.0, formatter) - } -} - -impl ErrorSource { - pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self { - Self(std::sync::Arc::new(error)) - } -} - -impl PartialEq for ErrorSource { - fn eq(&self, other: &Self) -> bool { - std::sync::Arc::ptr_eq(&self.0, &other.0) - } -} - -impl Eq for ErrorSource {} diff --git a/litellm-rust/crates/auth-types/src/http.rs b/litellm-rust/crates/auth-types/src/http.rs index 2de3bad9ecf..2ad07f7d541 100644 --- a/litellm-rust/crates/auth-types/src/http.rs +++ b/litellm-rust/crates/auth-types/src/http.rs @@ -21,9 +21,9 @@ pub fn apply_credential( placement: CredentialPlacement, ) -> Result, Error> { if credential.trim().is_empty() { - return Err(Error::InvalidConfiguration( - "credential cannot be empty".into(), - )); + return Err(Error::InvalidConfiguration(crate::ErrorDetail::Empty { + subject: "credential", + })); } if headers .iter() @@ -71,6 +71,10 @@ mod tests { ) .expect_err("provider policy must handle existing credentials"); - assert!(error.to_string().contains("already exists")); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("already exists") + ); } } diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index c875bda7979..99651977867 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -51,9 +51,9 @@ pub use credential::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; -pub use error::{Error, ErrorDetail, ErrorSource}; +pub use error::{Error, ErrorDetail}; pub use http::CredentialPlacement; -pub use params::{AwsParams, ParamSpec, VertexParams}; +pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth-types/src/params/aws.rs b/litellm-rust/crates/auth-types/src/params/aws.rs index 25fe52b2ce1..8075fc15e9f 100644 --- a/litellm-rust/crates/auth-types/src/params/aws.rs +++ b/litellm-rust/crates/auth-types/src/params/aws.rs @@ -4,15 +4,17 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use veil::Redact; -use super::ParamSpec; +use super::{ParamKind, ParamSpec}; const fn spec( setting: &'static str, + kind: ParamKind, wire: &'static [&'static str], env: &'static [&'static str], ) -> ParamSpec { ParamSpec { setting, + kind, wire, module_global: None, env, @@ -53,40 +55,67 @@ pub struct AwsParams { impl AwsParams { pub const ACCESS_KEY_ID: ParamSpec = spec( "access key id", + ParamKind::Credential, &["aws_access_key_id"], &["AWS_ACCESS_KEY_ID"], ); pub const SECRET_ACCESS_KEY: ParamSpec = spec( "secret access key", + ParamKind::Credential, &["aws_secret_access_key"], &["AWS_SECRET_ACCESS_KEY"], ); pub const SESSION_TOKEN: ParamSpec = spec( "session token", + ParamKind::Credential, &["aws_session_token"], &["AWS_SESSION_TOKEN"], ); pub const REGION: ParamSpec = spec( "region", + ParamKind::Setting, &["aws_region_name"], &["AWS_REGION_NAME", "AWS_REGION"], ); - pub const SESSION_NAME: ParamSpec = - spec("session name", &["aws_session_name"], &["AWS_SESSION_NAME"]); - pub const PROFILE_NAME: ParamSpec = - spec("profile name", &["aws_profile_name"], &["AWS_PROFILE_NAME"]); - pub const ROLE_NAME: ParamSpec = spec("role name", &["aws_role_name"], &["AWS_ROLE_NAME"]); + pub const SESSION_NAME: ParamSpec = spec( + "session name", + ParamKind::Setting, + &["aws_session_name"], + &["AWS_SESSION_NAME"], + ); + pub const PROFILE_NAME: ParamSpec = spec( + "profile name", + ParamKind::Setting, + &["aws_profile_name"], + &["AWS_PROFILE_NAME"], + ); + pub const ROLE_NAME: ParamSpec = spec( + "role name", + ParamKind::Setting, + &["aws_role_name"], + &["AWS_ROLE_NAME"], + ); pub const WEB_IDENTITY_TOKEN: ParamSpec = spec( "web identity token", + ParamKind::Credential, &["aws_web_identity_token"], &["AWS_WEB_IDENTITY_TOKEN"], ); - pub const STS_ENDPOINT: ParamSpec = - spec("STS endpoint", &["aws_sts_endpoint"], &["AWS_STS_ENDPOINT"]); - pub const EXTERNAL_ID: ParamSpec = - spec("external id", &["aws_external_id"], &["AWS_EXTERNAL_ID"]); + pub const STS_ENDPOINT: ParamSpec = spec( + "STS endpoint", + ParamKind::Setting, + &["aws_sts_endpoint"], + &["AWS_STS_ENDPOINT"], + ); + pub const EXTERNAL_ID: ParamSpec = spec( + "external id", + ParamKind::Setting, + &["aws_external_id"], + &["AWS_EXTERNAL_ID"], + ); pub const BEDROCK_RUNTIME_ENDPOINT: ParamSpec = spec( "Bedrock runtime endpoint", + ParamKind::Setting, &["aws_bedrock_runtime_endpoint"], &["AWS_BEDROCK_RUNTIME_ENDPOINT"], ); diff --git a/litellm-rust/crates/auth-types/src/params/mod.rs b/litellm-rust/crates/auth-types/src/params/mod.rs index 88bc64cc2fc..daad1aff8a8 100644 --- a/litellm-rust/crates/auth-types/src/params/mod.rs +++ b/litellm-rust/crates/auth-types/src/params/mod.rs @@ -10,11 +10,18 @@ pub use vertex::VertexParams; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct ParamSpec { pub setting: &'static str, + pub kind: ParamKind, pub wire: &'static [&'static str], pub module_global: Option<&'static str>, pub env: &'static [&'static str], } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ParamKind { + Credential, + Setting, +} + impl ParamSpec { /// The first non-blank value: each wire name through `params`, then each environment /// name through `env`. diff --git a/litellm-rust/crates/auth-types/src/params/vertex.rs b/litellm-rust/crates/auth-types/src/params/vertex.rs index 06ec3008dfb..2a6b8bc0ca6 100644 --- a/litellm-rust/crates/auth-types/src/params/vertex.rs +++ b/litellm-rust/crates/auth-types/src/params/vertex.rs @@ -2,7 +2,7 @@ use serde::{Deserialize, Deserializer, Serialize, de::Unexpected}; use serde_json::Value; use veil::Redact; -use super::ParamSpec; +use super::{ParamKind, ParamSpec}; /// The Vertex AI fields of Python's `GenericLiteLLMParams`, both the current names and the /// `vertex_ai_*` spellings `VertexBase.safe_get_vertex_ai_*` still read. @@ -27,18 +27,21 @@ pub struct VertexParams { impl VertexParams { pub const CREDENTIALS: ParamSpec = ParamSpec { setting: "credentials", + kind: ParamKind::Credential, wire: &["vertex_credentials", "vertex_ai_credentials"], module_global: None, env: &["VERTEXAI_CREDENTIALS"], }; pub const PROJECT: ParamSpec = ParamSpec { setting: "project", + kind: ParamKind::Setting, wire: &["vertex_project", "vertex_ai_project"], module_global: Some("vertex_project"), env: &["VERTEXAI_PROJECT"], }; pub const LOCATION: ParamSpec = ParamSpec { setting: "location", + kind: ParamKind::Setting, wire: &["vertex_location", "vertex_ai_location"], module_global: Some("vertex_location"), env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"], diff --git a/litellm-rust/crates/auth-types/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs index ee2bf08bede..8fff59e284a 100644 --- a/litellm-rust/crates/auth-types/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -1,5 +1,5 @@ use crate::http::apply_credential; -use crate::{CredentialPlacement, Error, ResolvedCredential}; +use crate::{CredentialPlacement, Error, ErrorDetail, ResolvedCredential}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CredentialPlanKind { @@ -31,11 +31,18 @@ pub struct ProviderAuthPolicy { impl ProviderAuthPolicy { pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool { - headers.iter().any(|(name, _)| { - self.accepted_existing_headers - .iter() - .any(|accepted| name.eq_ignore_ascii_case(accepted)) - }) + self.existing_credential_header(headers).is_some() + } + + fn existing_credential_header(&self, headers: &[(String, String)]) -> Option<&'static str> { + self.accepted_existing_headers + .iter() + .copied() + .find(|accepted| { + headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case(accepted)) + }) } pub fn apply( @@ -44,11 +51,11 @@ impl ProviderAuthPolicy { kind: CredentialPlanKind, credential: &ResolvedCredential, ) -> Result, Error> { - if self.has_existing_credential(&headers) { + if let Some(existing) = self.existing_credential_header(&headers) { return match self.existing_header_behavior { ExistingHeaderBehavior::Preserve => Ok(headers), ExistingHeaderBehavior::Reject => Err(Error::InvalidConfiguration( - "credential header already exists".into(), + ErrorDetail::DuplicateHeader(existing), )), }; } @@ -56,11 +63,7 @@ impl ProviderAuthPolicy { .rules .iter() .find(|rule| rule.kind == kind) - .ok_or_else(|| { - Error::InvalidConfiguration( - "credential plan is not allowed by the provider auth policy".into(), - ) - })?; + .ok_or(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed))?; apply_credential(headers, credential.secret().expose(), rule.placement) } } @@ -108,6 +111,10 @@ mod tests { ) .unwrap_err(); - assert!(error.to_string().contains("not allowed")); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("not allowed") + ); } } diff --git a/litellm-rust/crates/auth-types/tests/error.rs b/litellm-rust/crates/auth-types/tests/error.rs deleted file mode 100644 index aaedc3814f0..00000000000 --- a/litellm-rust/crates/auth-types/tests/error.rs +++ /dev/null @@ -1,74 +0,0 @@ -use litellm_auth_types::Error; -use rstest::rstest; - -#[rstest] -#[case::api_key( - Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, - "Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable" -)] -#[case::another_api_key( - Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" }, - "Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable" -)] -#[case::api_base( - Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, - "Missing Example API Base - Pass api_base" -)] -#[case::another_api_base( - Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" }, - "Missing Custom API Base - Set CUSTOM_ENDPOINT" -)] -#[case::configuration( - Error::InvalidConfiguration("credential selector is invalid".into()), - "invalid authentication configuration: credential selector is invalid" -)] -#[case::acquisition( - Error::CredentialAcquisition("token expired".into()), - "credential acquisition failed: token expired" -)] -#[case::caller( - Error::EmptyCallerCredential("empty token"), - "credential caller failed: empty token" -)] -#[case::provider( - Error::ProviderAuthentication("provider rejected credentials".into()), - "provider rejected credentials" -)] -#[case::chain( - Error::CredentialChain(vec![ - Error::CredentialAcquisition("token expired".into()), - Error::EmptyCallerCredential("empty token"), - ]), - "credential acquisition failed: credential acquisition failed: token expired; credential caller failed: empty token" -)] -fn display_preserves_failure_phase_and_caller_context( - #[case] error: Error, - #[case] expected: &str, -) { - assert_eq!(error.to_string(), expected); -} - -#[rstest] -#[case::configuration(true)] -#[case::acquisition(false)] -fn contextual_failures_keep_the_original_source(#[case] configuration: bool) { - use litellm_auth_types::ErrorDetail; - - let detail = ErrorDetail::failed( - "test credential", - std::io::Error::from(std::io::ErrorKind::PermissionDenied), - ); - let error = if configuration { - Error::InvalidConfiguration(detail) - } else { - Error::CredentialAcquisition(detail) - }; - let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| { - error.source() - }) - .find_map(|error| error.downcast_ref::()) - .expect("the original credential error remains available"); - assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied); - assert!(error.to_string().contains("test credential failed:")); - assert!(error.to_string().ends_with(&source.to_string())); -} diff --git a/litellm-rust/crates/auth-types/tests/params.rs b/litellm-rust/crates/auth-types/tests/params.rs index 9d488b170e1..38890edb39d 100644 --- a/litellm-rust/crates/auth-types/tests/params.rs +++ b/litellm-rust/crates/auth-types/tests/params.rs @@ -1,8 +1,9 @@ -use litellm_auth_types::{AwsParams, ParamSpec}; +use litellm_auth_types::{AwsParams, ParamKind, ParamSpec}; use rstest::rstest; const LOCATION: ParamSpec = ParamSpec { setting: "location", + kind: ParamKind::Setting, wire: &["vertex_location", "vertex_ai_location"], module_global: Some("vertex_location"), env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"], diff --git a/litellm-rust/crates/auth-types/tests/token.rs b/litellm-rust/crates/auth-types/tests/token.rs index d7b23531596..7f7027942b0 100644 --- a/litellm-rust/crates/auth-types/tests/token.rs +++ b/litellm-rust/crates/auth-types/tests/token.rs @@ -1,5 +1,4 @@ use std::{ - error::Error as StdError, future::{Future, poll_fn}, sync::{ Arc, @@ -53,7 +52,7 @@ async fn callbacks_acquire_fresh_credentials_on_demand(#[case] access_token: boo #[rstest] #[tokio::test] -async fn callback_errors_preserve_the_original_source() { +async fn callback_errors_keep_their_typed_cause() { let provider = TokenProviderHandle::from_callback(|| async { Err(Error::CredentialAcquisition(ErrorDetail::failed( "caller credential", @@ -62,13 +61,15 @@ async fn callback_errors_preserve_the_original_source() { }); let error = provider.acquire().await.unwrap_err(); - assert!(matches!(error, Error::CredentialAcquisition(_))); - let source = std::iter::successors(Some(&error as &(dyn StdError + 'static)), |error| { - (*error).source() - }) - .find_map(|error| error.downcast_ref::()) - .unwrap(); - assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied); + let Error::CredentialAcquisition(ErrorDetail::Failed { cause, .. }) = &error else { + panic!("the callback's error passes through unchanged: {error:?}"); + }; + assert_eq!( + cause + .downcast_ref::() + .map(std::io::Error::kind), + Some(std::io::ErrorKind::PermissionDenied) + ); } struct Release(Arc); diff --git a/litellm-rust/crates/cache-azure-blob/Cargo.toml b/litellm-rust/crates/cache-azure-blob/Cargo.toml index c28cb90d84a..202300c7ab4 100644 --- a/litellm-rust/crates/cache-azure-blob/Cargo.toml +++ b/litellm-rust/crates/cache-azure-blob/Cargo.toml @@ -16,6 +16,7 @@ azure_core = "1.1.0" azure_storage_blob = "1.1.0" futures-util.workspace = true reqwest.workspace = true +thiserror.workspace = true tokio.workspace = true url.workspace = true diff --git a/litellm-rust/crates/cache-azure-blob/src/cache.rs b/litellm-rust/crates/cache-azure-blob/src/cache.rs index c5c1fdd8ab9..e3b013eba4f 100644 --- a/litellm-rust/crates/cache-azure-blob/src/cache.rs +++ b/litellm-rust/crates/cache-azure-blob/src/cache.rs @@ -18,6 +18,10 @@ use url::Url; use crate::{credential::AzureBlobCredential, transport::ReqwestTransport}; +#[derive(Debug, thiserror::Error)] +#[error("account URL cannot hold a container path")] +struct NotABaseUrl; + pub struct AzureBlobCache { container: BlobContainerClient, codec: C, @@ -57,12 +61,12 @@ impl AzureBlobCache { codec: C, runtime: Handle, ) -> Result { - let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?; + let parsed = Url::parse(account_url).map_err(Error::unavailable)?; let account_url = parsed.as_str().trim_end_matches('/').to_string(); let container_url = { let mut url = parsed; url.path_segments_mut() - .map_err(|()| Error::Unavailable)? + .map_err(|()| Error::unavailable(NotABaseUrl))? .pop_if_empty() .push(container); url @@ -75,7 +79,7 @@ impl AzureBlobCache { ..BlobContainerClientOptions::default() }), ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let cache = Self { container: client, codec, @@ -101,7 +105,7 @@ impl AzureBlobCache { Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => { Ok(()) } - Err(_) => Err(Error::Unavailable), + Err(error) => Err(Error::unavailable(error)), } } @@ -116,7 +120,7 @@ impl AzureBlobCache { { Ok(_) => Ok(()), Err(error) if !overwrite && is_already_present(&error) => Ok(()), - Err(_) => Err(Error::Unavailable), + Err(error) => Err(Error::unavailable(error)), } } @@ -126,13 +130,9 @@ impl AzureBlobCache { Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => { return Ok(None); } - Err(_) => return Err(Error::Unavailable), + Err(error) => return Err(Error::unavailable(error)), }; - let bytes = response - .body - .collect() - .await - .map_err(|_| Error::Unavailable)?; + let bytes = response.body.collect().await.map_err(Error::unavailable)?; self.codec.decode(&bytes).map(Some) } @@ -140,16 +140,16 @@ impl AzureBlobCache { let mut pages = self .container .list_blobs(None) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .into_pages(); - while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? { - let page = page.into_model().map_err(|_| Error::Unavailable)?; + while let Some(page) = pages.try_next().await.map_err(Error::unavailable)? { + let page = page.into_model().map_err(Error::unavailable)?; for name in page.blob_items.into_iter().filter_map(|item| item.name) { self.container .blob_client(&name) .delete(None) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } } Ok(()) diff --git a/litellm-rust/crates/cache-azure-blob/src/credential.rs b/litellm-rust/crates/cache-azure-blob/src/credential.rs index d1a3d0e44ec..f58a599f8e3 100644 --- a/litellm-rust/crates/cache-azure-blob/src/credential.rs +++ b/litellm-rust/crates/cache-azure-blob/src/credential.rs @@ -17,6 +17,10 @@ const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; type EnvLookup = Arc Option + Send + Sync>; +#[derive(Debug, thiserror::Error)] +#[error("no Azure credential is available for blob storage")] +struct NoCredential; + pub struct AzureBlobCredential { service: AzureAuthService, env_lookup: EnvLookup, @@ -62,15 +66,8 @@ impl TokenCredential for AzureBlobCredential { &lookup, ) .await - .map_err(|error| { - azure_core::Error::with_message(ErrorKind::Credential, error.to_string()) - })? - .ok_or_else(|| { - azure_core::Error::with_message( - ErrorKind::Credential, - "no Azure credential is available for blob storage", - ) - })?; + .map_err(|error| azure_core::Error::new(ErrorKind::Credential, error))? + .ok_or_else(|| azure_core::Error::new(ErrorKind::Credential, NoCredential))?; let (token, expires_on) = match credential.into_value() { ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on), ResolvedCredential::Static(token) => (token, None), diff --git a/litellm-rust/crates/cache-azure-blob/src/transport.rs b/litellm-rust/crates/cache-azure-blob/src/transport.rs index 3914b92365c..62feee8fae0 100644 --- a/litellm-rust/crates/cache-azure-blob/src/transport.rs +++ b/litellm-rust/crates/cache-azure-blob/src/transport.rs @@ -24,6 +24,7 @@ impl HttpClient for ReqwestTransport { Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)), }; let response = outgoing.send().await.map_err(|error| { + let error = litellm_http::Error::from(error); let kind = if error.is_connect() { ErrorKind::Connection } else { @@ -41,9 +42,9 @@ impl HttpClient for ReqwestTransport { ); } } - let body = response - .bytes_stream() - .map_err(|error| azure_core::Error::new(ErrorKind::Io, error)); + let body = response.bytes_stream().map_err(|error| { + azure_core::Error::new(ErrorKind::Io, litellm_http::Error::from(error)) + }); Ok(AsyncRawResponse::new(status, headers, Box::pin(body))) } } diff --git a/litellm-rust/crates/cache-azure-blob/tests/cache.rs b/litellm-rust/crates/cache-azure-blob/tests/cache.rs index b9f50b824b3..3b014f53e82 100644 --- a/litellm-rust/crates/cache-azure-blob/tests/cache.rs +++ b/litellm-rust/crates/cache-azure-blob/tests/cache.rs @@ -128,7 +128,7 @@ fn connect_surfaces_service_failures() { ResponseCacheCodec, runtime.handle().clone(), )); - assert!(matches!(result, Err(Error::Unavailable))); + assert!(matches!(result, Err(Error::Unavailable(_)))); } #[rstest] @@ -263,7 +263,7 @@ fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[ca fixture.service.seed_blob(key, bytes); assert!(matches!( fixture.cache.get_cache(key, &no_ttl()), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(_)) )); } @@ -378,15 +378,15 @@ fn service_failures_map_to_unavailable(fixture: Fixture) { fixture.service.set_failing(true); assert!(matches!( fixture.cache.get_cache("key", &no_ttl()), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) )); assert!(matches!( fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) )); assert!(matches!( fixture.cache.flush_cache(), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) )); assert!(matches!( fixture.runtime.block_on( @@ -394,7 +394,7 @@ fn service_failures_map_to_unavailable(fixture: Fixture) { .cache .async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl()) ), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) )); } diff --git a/litellm-rust/crates/cache-azure-blob/tests/transport.rs b/litellm-rust/crates/cache-azure-blob/tests/transport.rs index c8b14cd8543..e600b15af0f 100644 --- a/litellm-rust/crates/cache-azure-blob/tests/transport.rs +++ b/litellm-rust/crates/cache-azure-blob/tests/transport.rs @@ -31,7 +31,7 @@ async fn connect(server: &MockServer) -> AzureBlobCache> { None, ClientOptions { transport: Some(Transport::new(Arc::new(ReqwestTransport( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), )))), ..ClientOptions::default() }, diff --git a/litellm-rust/crates/cache-disk/Cargo.toml b/litellm-rust/crates/cache-disk/Cargo.toml index 825b23899e7..b2f7004910f 100644 --- a/litellm-rust/crates/cache-disk/Cargo.toml +++ b/litellm-rust/crates/cache-disk/Cargo.toml @@ -12,6 +12,7 @@ rand.workspace = true rusqlite = { version = "0.39", features = ["bundled"] } serde-pickle = "1.2" serde_json.workspace = true +thiserror.workspace = true tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-disk/src/cache.rs b/litellm-rust/crates/cache-disk/src/cache.rs index b8a2c6ed1e4..880a1f86e71 100644 --- a/litellm-rust/crates/cache-disk/src/cache.rs +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -64,7 +64,7 @@ impl DiskCache { { tokio::task::spawn_blocking(move || operation(&store)) .await - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? } } @@ -162,7 +162,7 @@ impl BatchCache for DiskCache Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), }) .collect() @@ -186,7 +186,7 @@ impl BatchCache for DiskCache match self.decode_stored(value) { Ok(Some(value)) => Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), }, }) @@ -272,7 +272,7 @@ fn increment( Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64()))) }; store.update(key, unix_now(), &mut apply)?; - result.ok_or(Error::InvalidEntry) + result.ok_or(Error::InvalidEntry(None)) } fn unix_now() -> f64 { diff --git a/litellm-rust/crates/cache-disk/src/python/value.rs b/litellm-rust/crates/cache-disk/src/python/value.rs index 645eba7b757..6e7c2ee0473 100644 --- a/litellm-rust/crates/cache-disk/src/python/value.rs +++ b/litellm-rust/crates/cache-disk/src/python/value.rs @@ -3,8 +3,8 @@ use py_literal::Value; use serde_json::{Map, Number}; pub(crate) fn from_pickle(bytes: &[u8]) -> Result { - let value = serde_pickle::value_from_slice(bytes, Default::default()) - .map_err(|_| Error::InvalidEntry)?; + let value = + serde_pickle::value_from_slice(bytes, Default::default()).map_err(Error::invalid_entry)?; from_pickle_value(value) } @@ -70,7 +70,7 @@ fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result Result { - value.parse().map_err(|_| Error::InvalidEntry) + value.parse().map_err(Error::invalid_entry) } pub(crate) fn from_json(value: serde_json::Value) -> Value { @@ -101,7 +101,7 @@ pub(crate) fn from_json(value: serde_json::Value) -> Value { pub(crate) fn from_json_text(value: &str) -> Result { serde_json::from_str(value) .map(from_json) - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } pub(crate) fn is_truthy(value: &Value) -> bool { @@ -131,7 +131,7 @@ pub(crate) fn to_f64(value: &Value) -> Option { } pub(crate) fn to_json(value: &Value) -> Result, Error> { - serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry) + serde_json::to_vec(&to_json_value(value)?).map_err(Error::invalid_entry) } fn to_json_value(value: &Value) -> Result { @@ -142,12 +142,12 @@ fn to_json_value(value: &Value) -> Result { value .to_string() .parse::() - .map_err(|_| Error::InvalidEntry)?, + .map_err(Error::invalid_entry)?, ), Value::Float(value) => { - serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?) + serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry(None))?) } - Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry), + Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry(None)), Value::String(value) => serde_json::Value::String(value.clone()), Value::Tuple(values) | Value::List(values) | Value::Set(values) => { serde_json::Value::Array( @@ -162,7 +162,7 @@ fn to_json_value(value: &Value) -> Result { .iter() .map(|(key, value)| { let Value::String(key) = key else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; Ok((key.clone(), to_json_value(value)?)) }) diff --git a/litellm-rust/crates/cache-disk/src/sqlite.rs b/litellm-rust/crates/cache-disk/src/sqlite.rs index 24c2be2e1c5..14cbf18b2ff 100644 --- a/litellm-rust/crates/cache-disk/src/sqlite.rs +++ b/litellm-rust/crates/cache-disk/src/sqlite.rs @@ -48,13 +48,13 @@ struct Row { impl DiskcacheSqliteStore { pub fn open(directory: impl AsRef) -> Result { let directory = directory.as_ref().to_path_buf(); - fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?; - let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?; + fs::create_dir_all(&directory).map_err(Error::unavailable)?; + let directory = std::path::absolute(&directory).map_err(Error::unavailable)?; let database = directory.join("cache.db"); - let connection = Connection::open(database).map_err(|_| Error::Unavailable)?; + let connection = Connection::open(database).map_err(Error::unavailable)?; connection .busy_timeout(std::time::Duration::from_secs(60)) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let mut settings = read_settings(&connection)?; for (key, value) in default_settings() { @@ -74,7 +74,7 @@ impl DiskcacheSqliteStore { value )", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; for (key, value) in &settings { if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") { connection @@ -82,7 +82,7 @@ impl DiskcacheSqliteStore { "INSERT OR REPLACE INTO Settings VALUES (?, ?)", params![key, value], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } } for (key, value) in [ @@ -96,7 +96,7 @@ impl DiskcacheSqliteStore { "INSERT OR IGNORE INTO Settings VALUES (?, ?)", params![key, value], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } connection .execute_batch( @@ -117,7 +117,7 @@ impl DiskcacheSqliteStore { CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw); CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let eviction_policy = setting_string(&settings, "eviction_policy") .unwrap_or_else(|| "least-recently-stored".to_string()); @@ -128,23 +128,27 @@ impl DiskcacheSqliteStore { .execute_batch( "CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } "least-recently-used" => { connection .execute_batch( "CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } "least-frequently-used" => { connection .execute_batch( "CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; + } + _ => { + return Err(Error::unavailable(SettingFault::EvictionPolicy( + eviction_policy, + ))); } - _ => return Err(Error::Unavailable), } connection .execute_batch( @@ -170,12 +174,12 @@ impl DiskcacheSqliteStore { UPDATE Settings SET value = value - OLD.size WHERE key = \"size\"; END;", ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let min_file_size = setting_i64(&settings, "disk_min_file_size") .unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE) .try_into() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT); let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT); let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0; @@ -207,7 +211,7 @@ impl DiskcacheSqliteStore { |row| row.get::<_, Option>(0), ) .optional() - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .flatten() { cleanup.push(old_filename); @@ -221,7 +225,7 @@ impl DiskcacheSqliteStore { |row| row.get::<_, i64>(0), ) .optional() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if let Some(rowid) = rowid { connection .execute( @@ -230,7 +234,7 @@ impl DiskcacheSqliteStore { WHERE rowid = ?", params![now, expire_time, now, size, mode, filename, value, rowid], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } else { connection .execute( @@ -240,7 +244,7 @@ impl DiskcacheSqliteStore { ) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)", params![key, now, expire_time, now, size, mode, filename, value], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } cleanup.extend(self.cull(connection, now)?); Ok(cleanup) @@ -257,13 +261,13 @@ impl DiskcacheSqliteStore { WHERE expire_time IS NOT NULL AND expire_time < ? ORDER BY expire_time LIMIT ?", ) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .query_map(params![now, self.cull_limit], |row| { Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) }) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .collect::, _>>() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; for (_, filename) in &expired { if let Some(filename) = filename { cleanup.push(filename.clone()); @@ -272,7 +276,7 @@ impl DiskcacheSqliteStore { for (rowid, _) in &expired { connection .execute("DELETE FROM Cache WHERE rowid = ?", params![rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit); if remaining <= 0 || self.volume(connection)? < self.size_limit { @@ -283,19 +287,23 @@ impl DiskcacheSqliteStore { "least-recently-stored" => "store_time", "least-recently-used" => "access_time", "least-frequently-used" => "access_count", - _ => return Err(Error::Unavailable), + policy => { + return Err(Error::unavailable(SettingFault::EvictionPolicy( + policy.to_owned(), + ))); + } }; let rows = connection .prepare(&format!( "SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?" )) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .query_map(params![remaining], |row| { Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) }) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .collect::, _>>() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; for (_, filename) in &rows { if let Some(filename) = filename { cleanup.push(filename.clone()); @@ -304,7 +312,7 @@ impl DiskcacheSqliteStore { for (rowid, _) in rows { connection .execute("DELETE FROM Cache WHERE rowid = ?", params![rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } Ok(cleanup) } @@ -312,15 +320,15 @@ impl DiskcacheSqliteStore { fn volume(&self, connection: &Connection) -> Result { let page_count: i64 = connection .query_row("PRAGMA page_count", [], |row| row.get(0)) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let page_size: i64 = connection .query_row("PRAGMA page_size", [], |row| row.get(0)) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let size: i64 = connection .query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| { row.get(0) }) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(page_count.saturating_mul(page_size).saturating_add(size)) } } @@ -331,13 +339,13 @@ impl DiskStore for DiskcacheSqliteStore { } fn get(&self, key: &str, now: f64) -> Result, Error> { - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)"; let row = connection .query_row(select, params![key, now], row_from_query) .optional() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if !self.statistics && !has_get_update(&self.eviction_policy) { return row .map(|row| fetch_row(&self.directory, row)) @@ -348,7 +356,7 @@ impl DiskStore for DiskcacheSqliteStore { let row = connection .query_row(select, params![key, now], row_from_query) .optional() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(row) = row else { if self.statistics { connection @@ -356,7 +364,7 @@ impl DiskStore for DiskcacheSqliteStore { "UPDATE Settings SET value = value + 1 WHERE key = 'misses'", [], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } return Ok(None); }; @@ -369,14 +377,14 @@ impl DiskStore for DiskcacheSqliteStore { "UPDATE Settings SET value = value + 1 WHERE key = 'hits'", [], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } else if !hit && self.statistics { connection .execute( "UPDATE Settings SET value = value + 1 WHERE key = 'misses'", [], ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } if has_get_update(&self.eviction_policy) && hit { let update = match self.eviction_policy.as_str() { @@ -384,16 +392,20 @@ impl DiskStore for DiskcacheSqliteStore { "least-frequently-used" => { "UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?" } - _ => return Err(Error::Unavailable), + policy => { + return Err(Error::unavailable(SettingFault::EvictionPolicy( + policy.to_owned(), + ))); + } }; if self.eviction_policy == "least-recently-used" { connection .execute(update, params![now, rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } else { connection .execute(update, params![rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } } value @@ -409,7 +421,7 @@ impl DiskStore for DiskcacheSqliteStore { ) -> Result<(), Error> { let columns = store_value(&self.directory, self.min_file_size, value)?; let new_filename = columns.filename.clone(); - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let result = transactional(&connection, |connection| { self.set_locked(connection, key, columns, expire_time, now) }); @@ -428,7 +440,7 @@ impl DiskStore for DiskcacheSqliteStore { } fn pop(&self, key: &str, now: f64) -> Result, Error> { - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let selected = transactional(&connection, |connection| { let row = connection .query_row( @@ -439,13 +451,13 @@ impl DiskStore for DiskcacheSqliteStore { row_from_query, ) .optional() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(row) = row else { return Ok(None); }; connection .execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid]) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(Some(row)) })?; let Some(row) = selected else { @@ -460,7 +472,7 @@ impl DiskStore for DiskcacheSqliteStore { } fn clear(&self) -> Result<(), Error> { - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let mut last_rowid = 0_i64; loop { let batch = transactional(&connection, |connection| { @@ -469,13 +481,13 @@ impl DiskStore for DiskcacheSqliteStore { "SELECT rowid, filename FROM Cache WHERE rowid > ? ORDER BY rowid LIMIT 100", ) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .query_map(params![last_rowid], |row| { Ok((row.get::<_, i64>(0)?, row.get::<_, Option>(1)?)) }) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .collect::, _>>() - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if rows.is_empty() { return Ok(rows); } @@ -486,7 +498,7 @@ impl DiskStore for DiskcacheSqliteStore { .join(","); connection .execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), []) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(rows) })?; if batch.is_empty() { @@ -509,7 +521,7 @@ impl DiskStore for DiskcacheSqliteStore { now: f64, apply: &mut dyn FnMut(Option) -> Result<(StoredValue, Option), Error>, ) -> Result<(), Error> { - let connection = self.connection.lock().map_err(|_| Error::Unavailable)?; + let connection = self.connection.lock()?; let mut created_filename = None; let result = transactional(&connection, |connection| { let current = connection @@ -521,7 +533,7 @@ impl DiskStore for DiskcacheSqliteStore { row_from_query, ) .optional() - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .map(|row| fetch_row(&self.directory, row)) .transpose()? .flatten(); @@ -582,21 +594,35 @@ fn read_settings(connection: &Connection) -> Result, Erro }; statement .query_map([], |row| Ok((row.get(0)?, row.get(1)?))) - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .collect::, _>>() - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) +} + +#[derive(Debug, thiserror::Error)] +enum SettingFault { + #[error("unknown eviction policy {0:?}")] + EvictionPolicy(String), + #[error("setting {0:?} is not an SQLite pragma")] + NotPragma(String), + #[error("setting {0:?} has an unsupported value type")] + PragmaValue(String), } fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> { - let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?; + let pragma = key + .strip_prefix("sqlite_") + .ok_or_else(|| Error::unavailable(SettingFault::NotPragma(key.to_owned())))?; match value { Value::Integer(value) => connection .pragma_update(None, pragma, value) - .map_err(|_| Error::Unavailable), + .map_err(Error::unavailable), Value::Text(value) => connection .pragma_update(None, pragma, value) - .map_err(|_| Error::Unavailable), - _ => Err(Error::Unavailable), + .map_err(Error::unavailable), + _ => Err(Error::unavailable(SettingFault::PragmaValue( + key.to_owned(), + ))), } } @@ -634,7 +660,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result, Error> { Value::Text(value) => Ok(Some(StoredValue::Text(value))), Value::Integer(value) => Ok(Some(StoredValue::Integer(value))), Value::Real(value) => Ok(Some(StoredValue::Float(value))), - Value::Null => Err(Error::InvalidEntry), + Value::Null => Err(Error::InvalidEntry(None)), }, MODE_BINARY | MODE_PICKLE => { let bytes = match row.value { @@ -645,7 +671,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result, Error> { }; value } - _ => return Err(Error::InvalidEntry), + _ => return Err(Error::InvalidEntry(None)), }; Ok(Some(if row.mode == MODE_BINARY { StoredValue::Bytes(bytes) @@ -663,24 +689,24 @@ fn fetch_row(directory: &Path, row: Row) -> Result, Error> { } Value::Blob(value) => value, Value::Text(value) => value.into_bytes(), - _ => return Err(Error::InvalidEntry), + _ => return Err(Error::InvalidEntry(None)), }; Ok(Some(StoredValue::Text( - String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?, + String::from_utf8(bytes).map_err(Error::invalid_entry)?, ))) } - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } fn read_file(directory: &Path, filename: Option<&str>) -> Result>, Error> { let Some(filename) = filename else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; match fs::read(directory.join(filename)) { Ok(value) => Ok(Some(value)), Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), - Err(_) => Err(Error::Unavailable), + Err(error) => Err(Error::unavailable(error)), } } @@ -712,7 +738,7 @@ fn store_value( let bytes = value.into_bytes(); let filename = write_file(directory, &bytes)?; Ok(StoredColumns { - size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?, + size: i64::try_from(bytes.len()).map_err(Error::unavailable)?, mode: MODE_TEXT, filename: Some(filename), value: None, @@ -727,7 +753,7 @@ fn store_value( StoredValue::Bytes(value) => { let filename = write_file(directory, &value)?; Ok(StoredColumns { - size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?, + size: i64::try_from(value.len()).map_err(Error::unavailable)?, mode: MODE_BINARY, filename: Some(filename), value: None, @@ -742,7 +768,7 @@ fn store_value( StoredValue::Pickle(value) => { let filename = write_file(directory, &value)?; Ok(StoredColumns { - size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?, + size: i64::try_from(value.len()).map_err(Error::unavailable)?, mode: MODE_PICKLE, filename: Some(filename), value: None, @@ -761,14 +787,14 @@ fn write_file(directory: &Path, bytes: &[u8]) -> Result { let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]); let path = directory.join(&filename); if let Some(parent) = path.parent() { - fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?; + fs::create_dir_all(parent).map_err(Error::unavailable)?; } let mut file = OpenOptions::new() .write(true) .create_new(true) .open(path) - .map_err(|_| Error::Unavailable)?; - file.write_all(bytes).map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; + file.write_all(bytes).map_err(Error::unavailable)?; Ok(filename) } @@ -789,12 +815,12 @@ fn transactional( ) -> Result { connection .execute_batch("BEGIN IMMEDIATE") - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; match operation(connection) { Ok(value) => { connection .execute_batch("COMMIT") - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(value) } Err(error) => { diff --git a/litellm-rust/crates/cache-disk/tests/python_compat.rs b/litellm-rust/crates/cache-disk/tests/python_compat.rs index 9cbef8573bd..e167b0adb37 100644 --- a/litellm-rust/crates/cache-disk/tests/python_compat.rs +++ b/litellm-rust/crates/cache-disk/tests/python_compat.rs @@ -84,7 +84,7 @@ fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation match expected { ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected), ReadExpectation::Miss => assert_eq!(result.unwrap(), None), - ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))), + ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry(_)))), } } diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index 91630879cbe..dbd8c9ac908 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -13,9 +13,11 @@ litellm-auth-types.workspace = true litellm-cache.workspace = true percent-encoding.workspace = true reqwest.workspace = true +thiserror.workspace = true tokio.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-testing.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index bad81573cb4..a6b08445bc8 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -25,6 +25,14 @@ pub fn key_prefix(gcs_path: Option<&str>) -> String { } } +#[derive(Debug, thiserror::Error)] +#[error("GCS answered {0}")] +struct UnexpectedStatus(reqwest::StatusCode); + +#[derive(Debug, thiserror::Error)] +#[error("GCS cache worker thread panicked")] +struct WorkerPanicked; + #[derive(Clone, Debug)] pub struct GcsConfig { pub bucket_name: String, @@ -114,9 +122,9 @@ impl GcsCache { .body(payload) .send() .await - .map_err(|_| Error::Unavailable)?; + .map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?; if !response.status().is_success() { - return Err(Error::Unavailable); + return Err(Error::unavailable(UnexpectedStatus(response.status()))); } Ok(()) } @@ -135,18 +143,18 @@ impl GcsCache { .header(reqwest::header::CONTENT_TYPE, "application/json") .send() .await - .map_err(|_| Error::Unavailable)?; + .map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?; if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(None); } if !response.status().is_success() { - return Err(Error::Unavailable); + return Err(Error::unavailable(UnexpectedStatus(response.status()))); } - let body = response.bytes().await.map_err(|_| Error::Unavailable)?; - self.codec - .decode(&body) - .map(Some) - .map_err(|_| Error::InvalidEntry) + let body = response + .bytes() + .await + .map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?; + self.codec.decode(&body).map(Some) } fn run_sync(future: F) -> Result @@ -158,7 +166,7 @@ impl GcsCache { tokio::runtime::Builder::new_current_thread() .enable_all() .build() - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) .and_then(|runtime| runtime.block_on(future)) }; match tokio::runtime::Handle::try_current() { @@ -169,7 +177,7 @@ impl GcsCache { scope .spawn(|| run(future)) .join() - .map_err(|_| Error::Unavailable) + .map_err(|_| Error::unavailable(WorkerPanicked)) .and_then(|result| result) }), Err(_) => run(future), @@ -242,7 +250,7 @@ impl BatchCache for GcsCache { match self.async_get_cache(&key, &context).await { Ok(Some(value)) => Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), } } diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index adb601c276c..079ec44dbce 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -30,7 +30,7 @@ impl TokenSource for GcpTokenSource { self.auth .access_token(&self.config, &|name| std::env::var(name).ok()) .await - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) } } diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs index 096b691ea16..f8dba94ea8e 100644 --- a/litellm-rust/crates/cache-gcs/tests/cache.rs +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -55,19 +55,27 @@ async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer Ok(Some(json!({"value": "entry"}))) )] #[case::missing("missing", ResponseTemplate::new(404), Ok(None))] -#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))] -#[case::unauthorized("unauthorized", ResponseTemplate::new(401), Err(Error::Unavailable))] +#[case::server_error( + "server-error", + ResponseTemplate::new(500), + Err("cache is unavailable: GCS answered 500 Internal Server Error") +)] +#[case::unauthorized( + "unauthorized", + ResponseTemplate::new(401), + Err("cache is unavailable: GCS answered 401 Unauthorized") +)] #[case::invalid( "invalid", ResponseTemplate::new(200).set_body_string("not json"), - Err(Error::InvalidEntry) + Err("invalid cache entry: ") )] #[tokio::test] async fn get_maps_statuses_and_decode_failures( #[future(awt)] server: MockServer, #[case] key: &str, #[case] response: ResponseTemplate, - #[case] expected: Result, Error>, + #[case] expected: Result, &str>, ) { Mock::given(method("GET")) .and(path(format!("/storage/v1/b/bucket/o/{key}"))) @@ -76,8 +84,21 @@ async fn get_maps_statuses_and_decode_failures( .mount(&server) .await; let cache = support::cache(&server, None); - assert_eq!(cache.get_cache(key, &context()), expected); - assert_eq!(cache.async_get_cache(key, &context()).await, expected); + let outcome = |result: Result, Error>| { + result.map_err(|error| litellm_error::Report::new(&error).to_string()) + }; + let sync = outcome(cache.get_cache(key, &context())); + let asynchronous = outcome(cache.async_get_cache(key, &context()).await); + match expected { + Ok(value) => { + assert_eq!(sync, Ok(value.clone())); + assert_eq!(asynchronous, Ok(value)); + } + Err(prefix) => { + assert!(sync.unwrap_err().starts_with(prefix)); + assert!(asynchronous.unwrap_err().starts_with(prefix)); + } + } } #[rstest] @@ -97,7 +118,7 @@ async fn cache_exposes_its_configuration(#[future(awt)] server: MockServer) { path_service_account: Some("/secrets/sa.json".into()), ..support::config(&server, Some("folder")) }, - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), litellm_cache::JsonCodec::::new(), ); assert_eq!(cache.bucket_name(), "bucket"); @@ -257,7 +278,7 @@ struct FailingTokenSource; impl TokenSource for FailingTokenSource { fn bearer_token(&self) -> Pin> + Send + '_>> { - Box::pin(async { Err(Error::Unavailable) }) + Box::pin(async { Err(Error::Unavailable(None)) }) } } @@ -267,14 +288,14 @@ async fn token_source_failure_skips_http(#[future(awt)] server: MockServer) { let cache = support::cache_with_token(&server, None, Arc::new(FailingTokenSource)); assert_eq!( cache.get_cache("key", &context()).unwrap_err(), - Error::Unavailable + Error::Unavailable(None) ); assert_eq!( cache .async_set_cache("key", json!(1), context()) .await .unwrap_err(), - Error::Unavailable + Error::Unavailable(None) ); assert_eq!(server.received_requests().await.unwrap().len(), 0); } diff --git a/litellm-rust/crates/cache-gcs/tests/support/mod.rs b/litellm-rust/crates/cache-gcs/tests/support/mod.rs index beb9aa39d9c..99fb0270a74 100644 --- a/litellm-rust/crates/cache-gcs/tests/support/mod.rs +++ b/litellm-rust/crates/cache-gcs/tests/support/mod.rs @@ -29,7 +29,7 @@ pub fn cache_with_token( ) -> JsonGcsCache { GcsCache::with_token_source( config(server, gcs_path), - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), JsonCodec::new(), token, ) diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 74476559cbe..73fbf14b07a 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -2,7 +2,7 @@ use std::{ cmp::Reverse, collections::{BinaryHeap, HashMap, HashSet}, hash::Hash, - sync::{Arc, Mutex}, + sync::{Arc, Mutex, MutexGuard}, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -84,6 +84,10 @@ impl InMemoryCache { } } + fn lock(&self) -> Result>, Error> { + Ok(self.state.lock()?) + } + pub fn set_cache( &self, key: impl Into, @@ -94,13 +98,13 @@ impl InMemoryCache { return Ok(CacheWrite::Disabled); } let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; self.store(&mut state, key.into(), value, ttl, now) } pub fn get_cache(&self, key: &str) -> Result, Error> { let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; if state .expirations .get(key) @@ -124,7 +128,7 @@ impl InMemoryCache { /// fits. pub fn evict_cache(&self) -> Result<(), Error> { let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; Self::evict(&mut state, self.max_size_in_memory, now, None); Ok(()) } @@ -132,7 +136,7 @@ impl InMemoryCache { /// `evict_element_if_expired`: `true` when `key` had expired and was removed. pub fn evict_element_if_expired(&self, key: &str) -> Result { let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; let expired = state .expirations .get(key) @@ -153,12 +157,7 @@ impl InMemoryCache { /// The number of stored entries, expired ones included until they are evicted. pub fn len(&self) -> Result { - Ok(self - .state - .lock() - .map_err(|_| Error::Unavailable)? - .values - .len()) + Ok(self.lock()?.values.len()) } pub fn is_empty(&self) -> Result { @@ -167,12 +166,7 @@ impl InMemoryCache { /// Entries in the expiration heap, stale ones included; bounded by eviction. pub fn expiration_heap_len(&self) -> Result { - Ok(self - .state - .lock() - .map_err(|_| Error::Unavailable)? - .expiration_heap - .len()) + Ok(self.lock()?.expiration_heap.len()) } pub fn max_size_in_memory(&self) -> usize { @@ -184,13 +178,7 @@ impl InMemoryCache { } pub fn expires_at(&self, key: &str) -> Result, Error> { - Ok(self - .state - .lock() - .map_err(|_| Error::Unavailable)? - .expirations - .get(key) - .copied()) + Ok(self.lock()?.expirations.get(key).copied()) } pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { @@ -198,7 +186,7 @@ impl InMemoryCache { } pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result, Error> { - let state = self.state.lock().map_err(|_| Error::Unavailable)?; + let state = self.lock()?; let mut expirations = state .expirations .iter() @@ -213,13 +201,13 @@ impl InMemoryCache { } pub fn delete_cache(&self, key: &str) -> Result<(), Error> { - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; Self::remove(&mut state, key); Ok(()) } pub fn flush_cache(&self) -> Result<(), Error> { - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; state.values.clear(); state.expirations.clear(); state.expiration_heap.clear(); @@ -316,7 +304,7 @@ where return Ok(candidate); } let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; Self::evict(&mut state, self.max_size_in_memory, now, Some(key)); let existing = state .values @@ -351,7 +339,7 @@ impl CounterCache for InMemoryCache { return Ok(amount); } let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; let value = Self::live(&mut state, key, now).unwrap_or_default() + amount; self.store(&mut state, key.into(), value, self.get_ttl(&context), now)?; Ok(value) @@ -424,7 +412,7 @@ where return Ok(values); } let now = (self.now)(); - let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut state = self.lock()?; let mut stored = Self::live(&mut state, key, now).unwrap_or_default(); stored.extend(values.iter().cloned()); self.store(&mut state, key.into(), stored, ttl, now)?; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index ce7d6ac8a98..e6c74e5ccda 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -52,7 +52,7 @@ fn measured(capacity: usize) -> InMemoryCache { Some(4), Some(Arc::new(|value: &String| { if value.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } Ok(value.len()) })), @@ -275,7 +275,7 @@ fn evict_element_if_expired_reports_removal(clock: Clock) { #[case::fits("ok", Ok(true))] #[case::at_limit("four", Ok(true))] #[case::too_large("oversized", Ok(false))] -#[case::measure_error("", Err(Error::InvalidEntry))] +#[case::measure_error("", Err(Error::InvalidEntry(None)))] fn check_value_size_applies_the_entry_limit( #[case] value: &str, #[case] expected: Result, @@ -310,7 +310,7 @@ fn disabled_size_limited_and_validated_writes_are_observable() { assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into())); assert_eq!( cache.set_cache("invalid", String::new(), None), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!(cache.get_cache("invalid").unwrap(), None); cache.delete_cache("small").unwrap(); diff --git a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml index a44bef0a731..8a4caa8def7 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -18,6 +18,7 @@ tokio.workspace = true uuid.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } futures-executor = "0.3" litellm-cache-testing.workspace = true diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs b/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs index a140e0af174..9c93cc45386 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs @@ -36,7 +36,7 @@ impl QdrantSemanticCache { let exists = client .collection_exists(config.collection_name.clone()) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if !exists { client .create_collection( @@ -48,7 +48,7 @@ impl QdrantSemanticCache { .quantization_config(quantization(&config.quantization)), ) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } let _ = client .create_field_index(CreateFieldIndexCollectionBuilder::new( @@ -99,13 +99,13 @@ impl QdrantSemanticCache { .async_embed(&prompt, context.metadata.as_ref()) .await?; let response = - String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?; + String::from_utf8(self.codec.encode(&value)?).map_err(Error::invalid_entry)?; let payload = Payload::try_from(json!({ "litellm_cache_key": key, "text": prompt, "response": response, })) - .map_err(|_| Error::InvalidEntry)?; + .map_err(Error::invalid_entry)?; self.client .upsert_points( UpsertPointsBuilder::new( @@ -119,7 +119,7 @@ impl QdrantSemanticCache { .wait(true), ) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; Ok(()) } @@ -152,7 +152,7 @@ impl QdrantSemanticCache { ), ) .await - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(point) = result.result.into_iter().next() else { return Ok(SemanticLookup::miss(Some(0.0))); }; @@ -170,7 +170,7 @@ impl QdrantSemanticCache { let response = payload .get("response") .and_then(Value::as_str) - .ok_or(Error::InvalidEntry)?; + .ok_or(Error::InvalidEntry(None))?; Ok(SemanticLookup { value: Some(self.codec.decode(response.as_bytes())?), similarity: Some(similarity), diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs index b7fbcd9b02d..bad780bc57f 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -55,10 +55,9 @@ impl Embedder for OpenAiEmbedder { } .send() .await - .map_err(|_| Error::Unavailable)? - .error_for_status() - .map_err(|_| Error::Unavailable)?; - let body: Value = response.json().await.map_err(|_| Error::Unavailable)?; + .and_then(reqwest::Response::error_for_status) + .map_err(http_error)?; + let body: Value = response.json().await.map_err(http_error)?; body.get("data") .and_then(Value::as_array) .and_then(|data| data.first()) @@ -70,6 +69,14 @@ impl Embedder for OpenAiEmbedder { .map(|value| value.as_f64().map(|value| value as f32)) .collect::>>() }) - .ok_or(Error::Unavailable) + .ok_or(Error::unavailable(MissingEmbedding)) } } + +fn http_error(error: reqwest::Error) -> Error { + Error::unavailable(litellm_http::Error::from(error)) +} + +#[derive(Debug, thiserror::Error)] +#[error("embeddings response has no embedding")] +struct MissingEmbedding; diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs index 24e6e5eba3e..42261d21524 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -108,7 +108,7 @@ fn config(base: String, timeout: Option) -> OpenAiEmbedderConfig { async fn posts_embeddings_request_and_parses_vector() { let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; let embedder = OpenAiEmbedder::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), config( format!("{}/", server.base_url()), Some(Duration::from_secs(1)), @@ -134,13 +134,19 @@ async fn posts_embeddings_request_and_parses_vector() { } #[rstest] -#[case::error_status("500 Internal Server Error", "{}", 0, None, Err(Error::Unavailable))] +#[case::error_status( + "500 Internal Server Error", + "{}", + 0, + None, + Err(Error::Unavailable(None)) +)] #[case::timed_out( "200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#, 500, Some(Duration::from_millis(200)), - Err(Error::Unavailable) + Err(Error::Unavailable(None)) )] #[case::within_timeout( "200 OK", @@ -149,7 +155,7 @@ async fn posts_embeddings_request_and_parses_vector() { Some(Duration::from_secs(1)), Ok(vec![0.1, 0.2]) )] -#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable))] +#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable(None)))] #[tokio::test] async fn status_timeout_and_body_errors_are_unavailable( #[case] status: &str, @@ -161,16 +167,23 @@ async fn status_timeout_and_body_errors_are_unavailable( let server = TestHttpServer::response_after(status, body, Duration::from_millis(delay_ms)).await; let embedder = OpenAiEmbedder::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), config(server.base_url(), timeout), ); - assert_eq!(embedder.async_embed("hello", None).await, expected); + let result = embedder.async_embed("hello", None).await; + if let Err(error) = &result { + let report = litellm_error::Report::new(error).to_string(); + assert!(!report.contains(&server.base_url()), "{report}"); + } + let kind = + |result: Result, Error>| result.map_err(|error| std::mem::discriminant(&error)); + assert_eq!(kind(result), kind(expected)); } #[rstest] fn sync_embedding_is_unsupported() { let embedder = OpenAiEmbedder::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), config("http://127.0.0.1:9".to_owned(), None), ); assert_eq!( diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs index fe25c503a36..b0764b4470b 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -53,7 +53,10 @@ impl Embedder for FixedEmbedder { .lock() .unwrap() .push((input.to_owned(), metadata.cloned())); - self.vectors.get(input).cloned().ok_or(Error::Unavailable) + self.vectors + .get(input) + .cloned() + .ok_or(Error::Unavailable(None)) } } @@ -396,9 +399,9 @@ async fn lookup_reports_python_semantic_similarity( #[rstest] #[case::codec_decodes_the_payload(Some(json!("{\"a\":1}")), Ok(Some(json!({"a": 1}))))] -#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry))] -#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry))] -#[case::missing_response(None, Err(Error::InvalidEntry))] +#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry(None)))] +#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry(None)))] +#[case::missing_response(None, Err(Error::InvalidEntry(None)))] #[tokio::test(flavor = "multi_thread")] async fn stored_responses_go_through_the_codec( #[case] response: Option, @@ -419,9 +422,12 @@ async fn stored_responses_go_through_the_codec( .into(), }); + let kind = |result: Result, Error>| { + result.map_err(|error| std::mem::discriminant(&error)) + }; assert_eq!( - cache.async_get_cache("key", &context("hello")).await, - expected + kind(cache.async_get_cache("key", &context("hello")).await), + kind(expected) ); server.stop(); } @@ -434,7 +440,7 @@ async fn embedding_failures_propagate() { assert_eq!( cache.async_get_cache("key", &context("unknown")).await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); server.stop(); } @@ -493,10 +499,10 @@ async fn stopped_qdrant_server_maps_to_unavailable() { let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; server.stop(); tokio::time::sleep(Duration::from_millis(50)).await; - assert_eq!( + assert!(matches!( cache.async_get_cache("key", &context("hello")).await, - Err(Error::Unavailable) - ); + Err(Error::Unavailable(Some(_))) + )); } /// `_payload_matches_cache_key` compares `str(cached_key) == str(key)`, so a point whose stored diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index 02feb3cbdf4..42fc6ae585f 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -63,13 +63,13 @@ impl Inner { .arg(CACHE_KEY_FIELD) .arg(tag) .query::<()>(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if let Some(ttl) = ttl { redis::cmd("EXPIRE") .arg(&hash_key) .arg(ttl_seconds(ttl)) .query::<()>(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } Ok(()) } @@ -111,7 +111,7 @@ impl Inner { .arg("vector") .arg(vector_buffer(vector)) .query::(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(fields) = first_document(&result) else { return Ok(SemanticLookup::miss(Some(0.0))); }; diff --git a/litellm-rust/crates/cache-redis-semantic/src/index.rs b/litellm-rust/crates/cache-redis-semantic/src/index.rs index c141e47003a..23401938a0e 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/index.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/index.rs @@ -40,10 +40,10 @@ impl Index { Some(false) => self.isolated(connection, dims)?, None => match create_index(connection, &self.name, dims) { Ok(()) => self.name.clone(), - Err(_) => match index_compatible(connection, &self.name, dims)? { + Err(error) => match index_compatible(connection, &self.name, dims)? { Some(true) => self.name.clone(), Some(false) => self.isolated(connection, dims)?, - None => return Err(Error::Unavailable), + None => return Err(error), }, }, }; @@ -59,7 +59,7 @@ impl Index { redis::cmd("FT.DROPINDEX") .arg(&name) .query::<()>(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; create_index(connection, &name, dims)?; Ok(name) } @@ -109,7 +109,7 @@ fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> .arg("SEPARATOR") .arg(",") .query::<()>(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } fn index_compatible( @@ -123,14 +123,18 @@ fn index_compatible( { Ok(info) => info, Err(error) if unknown_index(&error) => return Ok(None), - Err(_) => return Err(Error::Unavailable), + Err(error) => return Err(Error::unavailable(error)), }; Ok(Some(schema_compatible(&info, dims))) } fn unknown_index(error: &redis::RedisError) -> bool { - let message = error.to_string().to_lowercase(); - message.contains("unknown") && message.contains("index") + error + .code() + .is_some_and(|code| code.eq_ignore_ascii_case("unknown")) + && error + .detail() + .is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index")) } struct Attribute { diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs index fa4ac00e767..cc42f4c3924 100644 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -79,7 +79,7 @@ fn s(value: &str) -> redis::Value { } fn unknown_index_error() -> redis::RedisError { - redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) + support::server_error("Unknown index name") } fn info_missing(index: &str) -> MockCmd { @@ -320,7 +320,7 @@ fn store_writes_the_redisvl_hash( #[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))] #[case::above_distance_threshold("key1", "0.5", None, Ok(None))] #[case::other_cache_key("other", "0.05", None, Ok(None))] -#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry))] +#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))] fn lookup_applies_threshold_scope_and_codec( #[case] stored_tag: &str, #[case] distance: &str, @@ -343,7 +343,10 @@ fn lookup_applies_threshold_scope_and_codec( FakeEmbedder::new(&[]), ); - assert_eq!(cache.get_cache("key1", &context), expected); + let kind = |result: Result, Error>| { + result.map_err(|error| std::mem::discriminant(&error)) + }; + assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected)); } #[rstest] @@ -451,10 +454,7 @@ fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticC info_missing(INDEX), MockCmd::new( create_index_command(INDEX, 3), - Err::<&str, _>(redis::RedisError::from(( - redis::ErrorKind::Extension, - "Index already exists", - ))), + Err::<&str, _>(support::server_error("Index already exists")), ), info(INDEX, compatible_info(3)), hset(INDEX, PROMPT, "key1", &VECTOR, &entry), @@ -673,10 +673,10 @@ fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) .set_cache("key1", entry.clone(), &context()) .unwrap(); - assert_eq!( - worker_b.get_cache("key1", &context()).unwrap_err(), - Error::Unavailable - ); + assert!(matches!( + worker_b.get_cache("key1", &context()), + Err(Error::Unavailable(Some(_))) + )); } #[fixture] @@ -732,10 +732,10 @@ fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option { let name = text(1); if state.indexes.contains_key(&name) { - return Err(error("Index already exists")); + return Err(server_error("Index already exists")); } let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes()); let prefix = args[position("PREFIX").unwrap() + 2].clone(); @@ -94,7 +94,7 @@ impl FakeSearch { let index = state .indexes .get(&text(1)) - .ok_or_else(|| error("Unknown index name"))?; + .ok_or_else(|| server_error("Unknown index name"))?; Ok(index_info(index)) } "FT.DROPINDEX" => { @@ -118,7 +118,7 @@ impl FakeSearch { let index = state .indexes .get(&text(1)) - .ok_or_else(|| error("no such index"))?; + .ok_or_else(|| server_error("no such index"))?; let query = text(2); let tag = query_tag(&query); let params = args.iter().position(|arg| arg == b"PARAMS").unwrap(); @@ -155,7 +155,7 @@ impl FakeSearch { ])) } "PING" => Ok(redis::Value::SimpleString("PONG".into())), - _ => Err(error("unsupported command")), + _ => Err(server_error("unsupported command")), } } } @@ -192,8 +192,11 @@ impl redis::ConnectionLike for FakeSearch { } } -fn error(message: &'static str) -> redis::RedisError { - redis::RedisError::from((redis::ErrorKind::Extension, message)) +pub fn server_error(reply: &str) -> redis::RedisError { + redis::parse_redis_value(format!("-{reply}\r\n").as_bytes()) + .expect("error reply") + .extract_error() + .expect_err("server error") } fn bulk(bytes: &[u8]) -> redis::Value { diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index bebfcafd98a..f7621299506 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,8 +7,10 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true +litellm-error.workspace = true redis = { workspace = true, features = ["cluster", "tls-rustls"] } r2d2 = "0.8.10" +thiserror.workspace = true tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index ebcf0b6916b..b873a2c0282 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -141,7 +141,7 @@ where redis::Value::Nil => Ok(None), redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some), redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } @@ -152,7 +152,7 @@ where match self.decode_response(value) { Ok(Some(value)) => Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), } } diff --git a/litellm-rust/crates/cache-redis/src/claim.rs b/litellm-rust/crates/cache-redis/src/claim.rs index e4cbe8c8d5c..c57a1b03395 100644 --- a/litellm-rust/crates/cache-redis/src/claim.rs +++ b/litellm-rust/crates/cache-redis/src/claim.rs @@ -17,7 +17,7 @@ fn stored_bytes(value: redis::Value) -> Result>, Error> { redis::Value::Nil => Ok(None), redis::Value::BulkString(bytes) => Ok(Some(bytes)), redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } @@ -37,13 +37,13 @@ where { let payload = codec.encode(&candidate)?; if payload.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } for _ in 0..CLAIM_ATTEMPTS { let current = stored_bytes( connection .get::<_, redis::Value>(key) - .map_err(|_| Error::Unavailable)?, + .map_err(Error::unavailable)?, )? .filter(|bytes| !bytes.is_empty()); let existing = current @@ -63,14 +63,18 @@ where .arg(write) .arg(u8::from(refresh)) .query::(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; if applied { return Ok(existing.unwrap_or(candidate)); } } - Err(Error::Unavailable) + Err(Error::unavailable(ClaimContended)) } +#[derive(Debug, thiserror::Error)] +#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")] +struct ClaimContended; + impl ClaimCache for RedisCache where S: CacheCodec + Clone + 'static, diff --git a/litellm-rust/crates/cache-redis/src/connection.rs b/litellm-rust/crates/cache-redis/src/connection.rs index 2f58e2a9b80..99befa16ffa 100644 --- a/litellm-rust/crates/cache-redis/src/connection.rs +++ b/litellm-rust/crates/cache-redis/src/connection.rs @@ -36,19 +36,19 @@ where ) -> Result { match self { Self::Pool(pool) => { - let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; + let mut pooled = pool.get().map_err(Error::unavailable)?; let result = operation(&mut ConnectionRef::Node(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable)); + pooled.failed = matches!(result, Err(Error::Unavailable(_))); result } Self::Cluster(pool) => { - let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; + let mut pooled = pool.get().map_err(Error::unavailable)?; let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable)); + pooled.failed = matches!(result, Err(Error::Unavailable(_))); result } Self::Fixed(connection) => { - let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + let mut connection = connection.lock()?; operation(&mut ConnectionRef::Node(&mut *connection)) } } @@ -61,7 +61,7 @@ where { tokio::task::spawn_blocking(move || connections.execute(operation)) .await - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? } pub fn fixed(connection: C) -> Self { @@ -95,7 +95,7 @@ fn pool(manager: M) -> Result, Error> { .connection_timeout(REDIS_TIMEOUT) .test_on_check_out(false) .build(manager) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } fn close_idle(pool: &r2d2::Pool) @@ -123,7 +123,7 @@ impl ConnectionManager { fn open(url: &str) -> Result { redis::Client::open(url) .map(Self) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } } @@ -151,14 +151,26 @@ impl r2d2::ManageConnection for ConnectionManager { } } +#[derive(Debug, thiserror::Error)] +enum TopologyFault { + #[error("Redis Cluster needs at least one startup node")] + NoStartupNodes, + #[error("Redis Cluster nodes need a TCP address")] + NonTcpAddress, + #[error("Redis Cluster answered a fan-out command without a per-node map")] + FanOutReply, + #[error("scan target does not match the connection topology")] + ScanTarget, +} + pub(crate) struct ClusterConnectionManager(ClusterClient); impl ClusterConnectionManager { fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { if startup_nodes.is_empty() { - return Err(Error::Unavailable); + return Err(Error::unavailable(TopologyFault::NoStartupNodes)); } - let info = url.into_connection_info().map_err(|_| Error::Unavailable)?; + let info = url.into_connection_info().map_err(Error::unavailable)?; let nodes = startup_nodes .iter() .map(|node| node_info(&info, node)) @@ -168,7 +180,7 @@ impl ClusterConnectionManager { .response_timeout(REDIS_TIMEOUT) .build() .map(Self) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } } @@ -185,7 +197,7 @@ fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result return Err(Error::Unavailable), + _ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)), }; Ok(info.clone().set_addr(addr)) } @@ -298,7 +310,7 @@ impl ConnectionRef<'_> { .and_then(|replies| pipeline.query(&mut Replies(Some(replies)))) } } - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } pub(crate) fn scan( @@ -311,7 +323,7 @@ impl ConnectionRef<'_> { Self::Node(connection) => { let page = scan_command(0, pattern, count) .query::(*connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; vec![(None, page)] } Self::Cluster(connection) => connection @@ -322,7 +334,7 @@ impl ConnectionRef<'_> { Some(ResponsePolicy::Special), )), ) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) .and_then(primary_pages)? .into_iter() .map(|(node, page)| (Some(node), page)) @@ -362,7 +374,7 @@ impl ConnectionRef<'_> { pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result { match self { - Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable), + Self::Node(connection) => command.query(*connection).map_err(Error::unavailable), Self::Cluster(connection) => { let value = connection .route_command( @@ -372,18 +384,16 @@ impl ConnectionRef<'_> { Some(ResponsePolicy::Special), )), ) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let redis::Value::Map(entries) = value else { - return Err(Error::Unavailable); + return Err(Error::unavailable(TopologyFault::FanOutReply)); }; let mut replies = entries .into_iter() .map(|(node, reply)| { Ok(( - redis::from_redis_value::(node) - .map_err(|_| Error::Unavailable)?, - redis::from_redis_value::(reply) - .map_err(|_| Error::Unavailable)?, + redis::from_redis_value::(node).map_err(Error::unavailable)?, + redis::from_redis_value::(reply).map_err(Error::unavailable)?, )) }) .collect::, Error>>()?; @@ -400,7 +410,7 @@ impl ConnectionRef<'_> { pub(crate) fn flushall(&mut self) -> Result<(), Error> { let command = redis::cmd("FLUSHALL"); match self { - Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable), + Self::Node(connection) => command.query(*connection).map_err(Error::unavailable), Self::Cluster(connection) => connection .route_command( &command, @@ -410,7 +420,7 @@ impl ConnectionRef<'_> { )), ) .map(|_| ()) - .map_err(|_| Error::Unavailable), + .map_err(Error::unavailable), } } @@ -424,7 +434,7 @@ impl ConnectionRef<'_> { let command = scan_command(cursor, pattern, count); match (self, node) { (Self::Node(connection), None) => { - command.query(*connection).map_err(|_| Error::Unavailable) + command.query(*connection).map_err(Error::unavailable) } (Self::Cluster(connection), Some(node)) => connection .route_command( @@ -434,9 +444,9 @@ impl ConnectionRef<'_> { port: node.port(), }), ) - .map_err(|_| Error::Unavailable) - .and_then(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable)), - _ => Err(Error::Unavailable), + .map_err(Error::unavailable) + .and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)), + _ => Err(Error::unavailable(TopologyFault::ScanTarget)), } } } @@ -445,14 +455,14 @@ type ScanPage = (u64, Vec); fn primary_pages(value: redis::Value) -> Result, Error> { let redis::Value::Map(entries) = value else { - return Err(Error::Unavailable); + return Err(Error::unavailable(TopologyFault::FanOutReply)); }; entries .into_iter() .map(|(node, page)| { - let node = redis::from_redis_value::(node).map_err(|_| Error::Unavailable)?; - let node = NodeAddress::try_from(node.as_str()).map_err(|_| Error::Unavailable)?; - let page = redis::from_redis_value::(page).map_err(|_| Error::Unavailable)?; + let node = redis::from_redis_value::(node).map_err(Error::unavailable)?; + let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?; + let page = redis::from_redis_value::(page).map_err(Error::unavailable)?; Ok((node, page)) }) .collect() diff --git a/litellm-rust/crates/cache-redis/src/counter.rs b/litellm-rust/crates/cache-redis/src/counter.rs index adcd597b2f4..a4beb9f5712 100644 --- a/litellm-rust/crates/cache-redis/src/counter.rs +++ b/litellm-rust/crates/cache-redis/src/counter.rs @@ -100,7 +100,7 @@ fn increment( .arg(amount) .arg(ttl) .query(connection) - .map_err(|_| Error::Unavailable); + .map_err(Error::unavailable); } connection .query_pipeline( @@ -140,7 +140,7 @@ where } fn count(value: redis::Value) -> Result, Error> { - redis::from_redis_value(value).map_err(|_| Error::InvalidEntry) + redis::from_redis_value(value).map_err(Error::invalid_entry) } impl BoundedCounterCache for RedisCache @@ -182,7 +182,7 @@ where .arg(value) .arg(ttl) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await } @@ -201,5 +201,5 @@ fn increment_with_floor( .arg(amount) .arg(ttl) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } diff --git a/litellm-rust/crates/cache-redis/src/keys.rs b/litellm-rust/crates/cache-redis/src/keys.rs index b9d7bb20ff3..2e8a9e60e65 100644 --- a/litellm-rust/crates/cache-redis/src/keys.rs +++ b/litellm-rust/crates/cache-redis/src/keys.rs @@ -16,7 +16,7 @@ where redis::cmd("TTL") .arg(key) .query::(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await?; Ok(u64::try_from(ttl).ok().map(Duration::from_secs)) @@ -36,7 +36,7 @@ where .arg(key) .arg(ttl) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await } diff --git a/litellm-rust/crates/cache-redis/src/lifecycle.rs b/litellm-rust/crates/cache-redis/src/lifecycle.rs index 6ab8c31509e..54fa8ecc3bf 100644 --- a/litellm-rust/crates/cache-redis/src/lifecycle.rs +++ b/litellm-rust/crates/cache-redis/src/lifecycle.rs @@ -11,11 +11,11 @@ where C: redis::ConnectionLike + Send + 'static, { fn sync_ping(&self) -> Result { - self.execute(|connection| connection.ping().map_err(|_| Error::Unavailable)) + self.execute(|connection| connection.ping().map_err(Error::unavailable)) } async fn ping(&self) -> Result { - self.run(|connection| connection.ping().map_err(|_| Error::Unavailable)) + self.run(|connection| connection.ping().map_err(Error::unavailable)) .await } } @@ -33,9 +33,10 @@ where RedisTopology::Cluster { .. } => "Redis Cluster", }; let ping = self - .run(|connection| Ok(connection.ping().map_err(|error| error.to_string()))) + .run(|connection| Ok(connection.ping().map_err(Error::unavailable))) .await - .unwrap_or_else(|error| Err(error.to_string())); + .and_then(|ping| ping) + .map_err(|error| litellm_error::Report::new(&error).to_string()); Ok(match ping { Ok(true) => CacheConnectionResult { status: CacheConnectionStatus::Success, diff --git a/litellm-rust/crates/cache-redis/src/queue.rs b/litellm-rust/crates/cache-redis/src/queue.rs index 623bce23640..d55e4f49187 100644 --- a/litellm-rust/crates/cache-redis/src/queue.rs +++ b/litellm-rust/crates/cache-redis/src/queue.rs @@ -29,7 +29,7 @@ where ttl: Option, ) -> Result { if values.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } let key = self.namespaced_key(key); let ttl = self.ttl_or_default(ttl); @@ -58,7 +58,7 @@ where async fn async_rpush(&self, key: &str, values: Vec) -> Result { if values.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } let key = self.namespaced_key(key); self.run(move |connection| { @@ -66,7 +66,7 @@ where .arg(key) .arg(values) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await } @@ -78,7 +78,7 @@ where max_len: usize, ) -> Result { if values.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } let key = self.namespaced_key(key); let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len); @@ -108,7 +108,7 @@ where let mut pipeline = redis::pipe(); for operation in operations { if operation.values.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } pipeline .cmd("RPUSH") @@ -130,7 +130,7 @@ where .run(move |connection| { command .query::(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await?; lpop_result(value, count.is_some()) @@ -182,7 +182,7 @@ fn redis_bytes(value: redis::Value) -> Result, Error> { match value { redis::Value::BulkString(bytes) => Ok(bytes), redis::Value::SimpleString(text) => Ok(text.into_bytes()), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } @@ -221,6 +221,6 @@ fn lpop_result(value: redis::Value, multiple: bool) -> Result, _>>() .map(RedisLpopResult::Values), value if !multiple => redis_bytes(value).map(RedisLpopResult::Value), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } diff --git a/litellm-rust/crates/cache-redis/src/script.rs b/litellm-rust/crates/cache-redis/src/script.rs index e8a16bcc2b5..09987437ed4 100644 --- a/litellm-rust/crates/cache-redis/src/script.rs +++ b/litellm-rust/crates/cache-redis/src/script.rs @@ -116,7 +116,7 @@ fn eval( .arg(keys) .arg(arguments) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } impl ScriptCache for RedisCache diff --git a/litellm-rust/crates/cache-redis/src/store.rs b/litellm-rust/crates/cache-redis/src/store.rs index 111d481ad02..daadc940ca9 100644 --- a/litellm-rust/crates/cache-redis/src/store.rs +++ b/litellm-rust/crates/cache-redis/src/store.rs @@ -32,7 +32,7 @@ where self.execute(|connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) } @@ -41,7 +41,7 @@ where let value = self.execute(|connection| { connection .get::<_, redis::Value>(key) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) })?; self.decode_response(value) } @@ -58,7 +58,7 @@ where self.run(move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await } @@ -73,7 +73,7 @@ where .run(move |connection| { connection .get::<_, redis::Value>(key) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) }) .await?; self.decode_response(value) @@ -158,7 +158,7 @@ pub(crate) fn mget( redis::cmd("MGET") .arg(keys) .query(connection) - .map_err(|_| Error::Unavailable) + .map_err(Error::unavailable) } impl DeleteCache for RedisCache @@ -168,12 +168,12 @@ where { fn delete_cache(&self, key: &str) -> Result<(), Error> { let key = self.namespaced_key(key); - self.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)) + self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable)) } async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { let key = self.namespaced_key(key); - self.run(move |connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)) + self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable)) .await } } @@ -188,7 +188,7 @@ where return Ok(0); } let keys = self.namespaced_keys(&keys); - self.run(move |connection| connection.del(keys).map_err(|_| Error::Unavailable)) + self.run(move |connection| connection.del(keys).map_err(Error::unavailable)) .await } } @@ -225,7 +225,7 @@ fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<( if !keys.is_empty() { connection .del::<_, usize>(keys) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; } Ok(true) }) diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 70baeffd572..9ec85251733 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -178,11 +178,11 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa ]); assert_eq!( cache.set_cache("invalid", 255, &context), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache.async_set_cache("invalid", 255, context.clone()).await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache @@ -191,7 +191,7 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa context.clone(), ) .await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache @@ -200,15 +200,15 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa ("invalid".into(), 255, None), ]) .await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache.get_cache("invalid", &context), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache.async_get_cache("invalid", &context).await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); } @@ -516,7 +516,7 @@ async fn sets_add_members_and_arm_the_default_ttl() { cache .async_set_cache_sadd("members", Vec::new(), None) .await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); } @@ -547,7 +547,7 @@ async fn queues_push_and_pop_namespaced_lists() { ); assert_eq!( cache.async_rpush("queue", Vec::new()).await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); assert_eq!( cache.async_lpop("queue", Some(2)).await.unwrap(), @@ -673,15 +673,15 @@ async fn rpush_and_trim_raises_when_a_queued_command_fails() { )]) .with_namespace(Some("ns".into())); - assert_eq!( + assert!(matches!( cache .async_rpush_and_trim("buf", vec!["c".into(), "d".into()], 3) .await, - Err(Error::Unavailable) - ); + Err(Error::Unavailable(Some(_))) + )); assert_eq!( cache.async_rpush_and_trim("buf", Vec::new(), 3).await, - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); } @@ -1122,24 +1122,30 @@ async fn test_connection_reports_success_with_the_python_message() { } #[rstest] -#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", false)] +#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", None)] #[case::connection_refused( Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))), "Redis connection failed:", - true + Some("connection refused") )] #[tokio::test] async fn test_connection_failures_use_the_python_result_contract( #[case] reply: redis::RedisResult<&'static str>, #[case] message: &str, - #[case] has_error: bool, + #[case] cause: Option<&str>, ) { let cache = json_cache(vec![MockCmd::new(redis::cmd("PING"), reply)]); let result = cache.test_connection().await.unwrap(); assert_eq!(result.status, CacheConnectionStatus::Failed); assert!(result.message.starts_with(message), "{}", result.message); - assert_eq!(result.error.is_some(), has_error); + assert_eq!(result.error.is_some(), cause.is_some()); + assert!(cause.is_none_or(|cause| { + result + .error + .as_deref() + .is_some_and(|error| error.contains(cause)) + })); } #[rstest] diff --git a/litellm-rust/crates/cache-redis/tests/cluster.rs b/litellm-rust/crates/cache-redis/tests/cluster.rs index c1a9a70a5d2..e1ba01bcb04 100644 --- a/litellm-rust/crates/cache-redis/tests/cluster.rs +++ b/litellm-rust/crates/cache-redis/tests/cluster.rs @@ -66,7 +66,7 @@ fn constructor_rejects_unusable_cluster_configs( JsonCodec::new(), ) .err(); - assert!(matches!(error, Some(Error::Unavailable))); + assert!(matches!(error, Some(Error::Unavailable(_)))); } #[rstest] diff --git a/litellm-rust/crates/cache-redis/tests/support/mod.rs b/litellm-rust/crates/cache-redis/tests/support/mod.rs index bff1e618582..b04e1c29f3b 100644 --- a/litellm-rust/crates/cache-redis/tests/support/mod.rs +++ b/litellm-rust/crates/cache-redis/tests/support/mod.rs @@ -16,7 +16,7 @@ impl CacheCodec for TaggedByteCodec { fn encode(&self, value: &u8) -> Result, Error> { if *value > 127 { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } Ok(vec![self.0, *value]) } @@ -24,7 +24,7 @@ impl CacheCodec for TaggedByteCodec { fn decode(&self, bytes: &[u8]) -> Result { match bytes { [tag, value] if *tag == self.0 => Ok(*value), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } } diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs index 68f2348e28b..deb1a56b503 100644 --- a/litellm-rust/crates/cache-response/src/buffer.rs +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -26,7 +26,7 @@ impl WriteBuffer { now: Duration, ) -> Result<(), Error> { let pending = { - let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?; + let mut entries = self.entries.lock()?; entries.push((request.clone(), response, now)); (entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries)) }; @@ -39,7 +39,7 @@ impl WriteBuffer { } pub fn clear(&self) -> Result<(), Error> { - self.entries.lock().map_err(|_| Error::Unavailable)?.clear(); + self.entries.lock()?.clear(); Ok(()) } } diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index 6b0f29e0a58..1adebe65bc3 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -14,24 +14,24 @@ impl CacheCodec for ResponseCacheCodec { .timestamp .is_some_and(|timestamp| !timestamp.is_finite()) { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } // Python reads a `response` that is either a dict or a serialized string, so every // other shape is written serialized. A string on the wire is therefore always a // serialized response, which keeps string-valued responses unambiguous. if value.timestamp.is_none() || value.response.is_object() { - return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry); + return serde_json::to_vec(value).map_err(Error::invalid_entry); } - let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?; + let response = serde_json::to_string(&value.response).map_err(Error::invalid_entry)?; serde_json::to_vec(&CacheEntry { timestamp: value.timestamp, response: Value::String(response), }) - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } fn decode(&self, bytes: &[u8]) -> Result { - let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?; + let text = std::str::from_utf8(bytes).map_err(Error::invalid_entry)?; let value = decode_value(text)?; let Some(timestamp) = value.get("timestamp") else { return Ok(CacheEntry { @@ -40,9 +40,9 @@ impl CacheCodec for ResponseCacheCodec { }); }; let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; - let response = match value.get("response").ok_or(Error::InvalidEntry)? { + let response = match value.get("response").ok_or(Error::InvalidEntry(None))? { Value::String(text) => decode_value(text)?, response => response.clone(), }; @@ -58,25 +58,25 @@ fn decode_value(text: &str) -> Result { return Ok(value); } check_literal_depth(text)?; - let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?; + let literal: py_literal::Value = text.parse().map_err(Error::invalid_entry)?; literal_value(literal, 0) } fn literal_value(value: py_literal::Value, depth: usize) -> Result { use py_literal::Value as Literal; if depth > 128 { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } match value { Literal::String(text) => Ok(Value::String(text)), Literal::Boolean(value) => Ok(Value::Bool(value)), Literal::None => Ok(Value::Null), Literal::Integer(value) => { - serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry) + serde_json::from_str(&value.to_string()).map_err(Error::invalid_entry) } Literal::Float(value) => serde_json::Number::from_f64(value) .map(Value::Number) - .ok_or(Error::InvalidEntry), + .ok_or(Error::InvalidEntry(None)), Literal::List(values) | Literal::Tuple(values) => values .into_iter() .map(|value| literal_value(value, depth + 1)) @@ -86,13 +86,13 @@ fn literal_value(value: py_literal::Value, depth: usize) -> Result .into_iter() .map(|(key, value)| { let Literal::String(key) = key else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; Ok((key, literal_value(value, depth + 1)?)) }) .collect::, _>>() .map(Value::Object), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } @@ -118,7 +118,7 @@ fn check_literal_depth(text: &str) -> Result<(), Error> { '[' | '{' | '(' => { depth += 1; if depth > 128 { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } } ']' | '}' | ')' => depth = depth.saturating_sub(1), diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index a5bbef99a3a..1505297315f 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ConnectionCache, Error, - FlushCache, + FlushCache, ShortBatch, semantic::{SemanticCache, SemanticLookup}, }; use serde_json::Value; @@ -114,7 +114,7 @@ where .get_cache(&cache_key(&request.key), &request.context) { Ok(entry) => entry, - Err(Error::InvalidEntry) => None, + Err(Error::InvalidEntry(_)) => None, Err(error) => return Err(error), }; Ok(Self::fresh_or_miss(entry, now, request.max_age)) @@ -134,7 +134,7 @@ where .await { Ok(entry) => entry, - Err(Error::InvalidEntry) => None, + Err(Error::InvalidEntry(_)) => None, Err(error) => return Err(error), }; Ok(Self::fresh_or_miss(entry, now, request.max_age)) @@ -338,7 +338,10 @@ where now: Duration, ) -> Result { if readable.len() != entries.len() { - return Err(Error::Unavailable); + return Err(Error::unavailable(ShortBatch { + requested: readable.len(), + returned: entries.len(), + })); } let mut values = vec![None; requests.len()]; for ((index, request), entry) in readable.into_iter().zip(entries) { @@ -361,7 +364,7 @@ where value: Self::fresh_or_miss(lookup.value, now, max_age), similarity: lookup.similarity, }), - Err(Error::InvalidEntry) => Ok(SemanticLookup::miss(None)), + Err(Error::InvalidEntry(_)) => Ok(SemanticLookup::miss(None)), Err(error) => Err(error), } } diff --git a/litellm-rust/crates/cache-response/tests/codec.rs b/litellm-rust/crates/cache-response/tests/codec.rs index 4fb5a5094cb..289465fcda9 100644 --- a/litellm-rust/crates/cache-response/tests/codec.rs +++ b/litellm-rust/crates/cache-response/tests/codec.rs @@ -44,25 +44,23 @@ fn decode_reads_every_python_envelope(#[case] bytes: &[u8], #[case] response: Va #[case::non_utf8(vec![0xff, 0xfe])] #[case::deep_nesting(format!("{}None{}", "[".repeat(1000), "]".repeat(1000)).into_bytes())] fn decode_rejects_invalid_entries(#[case] bytes: Vec) { - assert_eq!( - ResponseCacheCodec.decode(&bytes).unwrap_err(), - Error::InvalidEntry - ); + assert!(matches!( + ResponseCacheCodec.decode(&bytes), + Err(Error::InvalidEntry(_)) + )); } #[rstest] #[case::nan(f64::NAN)] #[case::infinity(f64::INFINITY)] fn encode_rejects_non_finite_timestamps(#[case] timestamp: f64) { - assert_eq!( - ResponseCacheCodec - .encode(&CacheEntry { - timestamp: Some(timestamp), - response: json!({}), - }) - .unwrap_err(), - Error::InvalidEntry - ); + assert!(matches!( + ResponseCacheCodec.encode(&CacheEntry { + timestamp: Some(timestamp), + response: json!({}), + }), + Err(Error::InvalidEntry(_)) + )); } #[rstest] diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 655fcb8a46a..70ae259d33a 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -158,8 +158,16 @@ fn scored(timestamp: f64, similarity: f64) -> Result, )] #[case::no_search(Ok(SemanticLookup::miss(None)), true, Ok(SemanticLookup::miss(None)))] #[case::disabled_reads_skip_the_backend(scored(95.0, 0.95), false, Ok(SemanticLookup::miss(None)))] -#[case::invalid_entry_is_a_miss(Err(Error::InvalidEntry), true, Ok(SemanticLookup::miss(None)))] -#[case::backend_errors_propagate(Err(Error::Unavailable), true, Err(Error::Unavailable))] +#[case::invalid_entry_is_a_miss( + Err(Error::InvalidEntry(None)), + true, + Ok(SemanticLookup::miss(None)) +)] +#[case::backend_errors_propagate( + Err(Error::Unavailable(None)), + true, + Err(Error::Unavailable(None)) +)] #[tokio::test] async fn semantic_lookup_applies_freshness_to_the_value_only( #[case] backend: Result, Error>, @@ -768,7 +776,7 @@ async fn gcs() -> (wiremock::MockServer, Gcs) { path_service_account: None, endpoint: server.uri(), }, - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), litellm_cache_response::ResponseCacheCodec, Arc::new(litellm_cache_gcs::StaticTokenSource("token".into())), ))); diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml index eb3a2fff1ac..d51c282d4d7 100644 --- a/litellm-rust/crates/cache-s3/Cargo.toml +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -20,6 +20,7 @@ reqwest.workspace = true tokio.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-testing.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs index f4f06e5371d..659390a5659 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -46,7 +46,7 @@ impl ProvideCredentials for S3Credentials { self.auth .resolve_credentials(self.config.clone(), &self.env) .await - .map_err(|_| CredentialsError::provider_error("S3 cache authentication failed")) + .map_err(CredentialsError::provider_error) }) } } diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index ced948e80c6..5369bca0f18 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -125,7 +125,7 @@ impl S3Cache { } None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"), }; - request.send().await.map_err(|_| Error::Unavailable)?; + request.send().await.map_err(Error::unavailable)?; Ok(()) } @@ -152,7 +152,7 @@ impl S3Cache { return Ok(None); } } - return Err(Error::Unavailable); + return Err(Error::unavailable(error)); } }; if let Some(expires) = output.expires_string() @@ -165,7 +165,7 @@ impl S3Cache { .body .collect() .await - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? .into_bytes(); self.codec.decode(&bytes).map(Some) } diff --git a/litellm-rust/crates/cache-s3/src/transport.rs b/litellm-rust/crates/cache-s3/src/transport.rs index eabc54ac9e2..bf1711276f2 100644 --- a/litellm-rust/crates/cache-s3/src/transport.rs +++ b/litellm-rust/crates/cache-s3/src/transport.rs @@ -29,9 +29,10 @@ impl HttpConnector for ReqwestHttpClient { .try_into_http1x() .map_err(|error| ConnectorError::other(error.into(), None))? .map(reqwest::Body::wrap); - let request = reqwest::Request::try_from(request) + let request = litellm_http::outbound::from_http(request) .map_err(|error| ConnectorError::other(error.into(), None))?; let response = client.execute(request).await.map_err(|error| { + let error = litellm_http::Error::from(error); if error.is_timeout() { ConnectorError::timeout(error.into()) } else { diff --git a/litellm-rust/crates/cache-s3/tests/cache.rs b/litellm-rust/crates/cache-s3/tests/cache.rs index 53a51f8a289..d0469620493 100644 --- a/litellm-rust/crates/cache-s3/tests/cache.rs +++ b/litellm-rust/crates/cache-s3/tests/cache.rs @@ -147,15 +147,15 @@ async fn async_set_signs_with_the_configured_keys(#[future(awt)] server: MockSer #[case::malformed( "malformed", ResponseTemplate::new(200).set_body_string("not a cache entry"), - Err(Error::InvalidEntry) + Err("invalid cache entry: ") )] -#[case::server_error("broken", ResponseTemplate::new(500), Err(Error::Unavailable))] +#[case::server_error("broken", ResponseTemplate::new(500), Err("cache is unavailable: "))] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn get_maps_s3_responses( #[future(awt)] server: MockServer, #[case] key: &str, #[case] response: ResponseTemplate, - #[case] expected: Result, Error>, + #[case] expected: Result, &str>, ) { Mock::given(method("GET")) .and(path(format!("/cache-bucket/team/{key}"))) @@ -164,9 +164,21 @@ async fn get_maps_s3_responses( .await; let cache = support::cache(&server.uri()); let context = ExactCacheContext::default(); - - assert_eq!(cache.get_cache(key, &context), expected); - assert_eq!(cache.async_get_cache(key, &context).await, expected); + let outcome = |result: Result, Error>| { + result.map_err(|error| litellm_error::Report::new(&error).to_string()) + }; + let sync = outcome(cache.get_cache(key, &context)); + let asynchronous = outcome(cache.async_get_cache(key, &context).await); + match expected { + Ok(value) => { + assert_eq!(sync, Ok(value.clone())); + assert_eq!(asynchronous, Ok(value)); + } + Err(prefix) => { + assert!(sync.unwrap_err().starts_with(prefix)); + assert!(asynchronous.unwrap_err().starts_with(prefix)); + } + } } #[rstest] diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs index b628c1df431..7378dde21bf 100644 --- a/litellm-rust/crates/cache-s3/tests/support/mod.rs +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -34,7 +34,7 @@ pub fn config(endpoint: &str) -> S3CacheConfig { pub fn cache_with(config: S3CacheConfig, runtime: Handle) -> JsonS3Cache { S3Cache::new( config, - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), JsonCodec::new(), runtime, ) diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml index ec755e61b8c..5b4e99de382 100644 --- a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -10,9 +10,11 @@ litellm-cache.workspace = true litellm-cache-redis.workspace = true redis = { workspace = true, features = ["tls-rustls"] } sha2.workspace = true +thiserror.workspace = true uuid = { version = "1", features = ["v4"] } [dev-dependencies] +litellm-error.workspace = true litellm-cache-testing.workspace = true redis-test.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/src/index.rs b/litellm-rust/crates/cache-valkey-semantic/src/index.rs index a0b7be4ca7b..3dfff32ff7c 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/index.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/index.rs @@ -23,8 +23,7 @@ pub(crate) fn ensure_index( ) -> Result<(), Error> { if index .dimension - .lock() - .map_err(|_| Error::Unavailable)? + .lock()? .is_some_and(|existing| existing == dimension) { return Ok(()); @@ -50,24 +49,42 @@ pub(crate) fn ensure_index( .arg("DISTANCE_METRIC") .arg("COSINE") .query::(connection) - .map(|_| ()) - .map_err(|error| error.to_string()); - if let Err(message) = create { - if !message.to_ascii_lowercase().contains("already exists") { - return Err(Error::Unavailable); + .map(|_| ()); + if let Err(error) = create { + if !index_exists(&error) { + return Err(Error::unavailable(error)); } let info = redis::cmd("FT.INFO") .arg(&index.name) .query::(connection) - .map_err(|_| Error::Unavailable)?; - if index_dimension_from_info(&info).is_some_and(|existing| existing != dimension) { - return Err(Error::Unavailable); + .map_err(Error::unavailable)?; + if let Some(existing) = + index_dimension_from_info(&info).filter(|existing| *existing != dimension) + { + return Err(Error::unavailable(DimensionMismatch { + existing, + dimension, + })); } } - *index.dimension.lock().map_err(|_| Error::Unavailable)? = Some(dimension); + *index.dimension.lock()? = Some(dimension); Ok(()) } +fn index_exists(error: &redis::RedisError) -> bool { + error.code().is_some() + && error + .detail() + .is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists")) +} + +#[derive(Debug, thiserror::Error)] +#[error("index exists with dimension {existing}, not {dimension}")] +struct DimensionMismatch { + existing: usize, + dimension: usize, +} + /// `_extract_index_dim`: flatten each attribute one level and read the value after /// `dimensions`. fn index_dimension_from_info(value: &redis::Value) -> Option { diff --git a/litellm-rust/crates/cache-valkey-semantic/src/search.rs b/litellm-rust/crates/cache-valkey-semantic/src/search.rs index 3b4fae84317..c5e4e748b58 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/search.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/search.rs @@ -53,9 +53,7 @@ pub(crate) fn write_document( .arg(ttl.as_secs()) .ignore(); } - pipeline - .query::<()>(connection) - .map_err(|_| Error::Unavailable) + pipeline.query::<()>(connection).map_err(Error::unavailable) } /// The KNN-1 search within `scope`: the closest document's similarity, and its stored response @@ -83,7 +81,7 @@ pub(crate) fn search_document( .arg("DIALECT") .arg(2) .query::(connection) - .map_err(|_| Error::Unavailable)?; + .map_err(Error::unavailable)?; let Some(fields) = search_fields(response)? else { return Ok(SemanticLookup::miss(Some(0.0))); }; @@ -91,7 +89,7 @@ pub(crate) fn search_document( fields .iter() .find_map(|(field, value)| (field == name).then(|| value.clone())) - .ok_or(Error::InvalidEntry) + .ok_or(Error::InvalidEntry(None)) }; let response = field("response")?; let similarity = 1.0 - parse_f64(&field("vector_distance")?)?; @@ -105,24 +103,24 @@ type SearchFields = Vec<(String, Vec)>; fn search_fields(value: redis::Value) -> Result, Error> { let redis::Value::Array(values) = value else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; - let total = parse_i64(values.first().ok_or(Error::InvalidEntry)?)?; + let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?; if total <= 0 || values.len() < 3 { return Ok(None); } let redis::Value::Array(fields) = &values[2] else { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); }; let (pairs, remainder) = fields.as_chunks::<2>(); if !remainder.is_empty() { - return Err(Error::InvalidEntry); + return Err(Error::InvalidEntry(None)); } pairs .iter() .map(|pair| { Ok(( - value_text(&pair[0]).ok_or(Error::InvalidEntry)?, + value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?, value_bytes(&pair[1])?, )) }) @@ -132,16 +130,16 @@ fn search_fields(value: redis::Value) -> Result, Error> { fn parse_i64(value: &redis::Value) -> Result { value_text(value) - .ok_or(Error::InvalidEntry)? + .ok_or(Error::InvalidEntry(None))? .parse() - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } fn parse_f64(value: &[u8]) -> Result { std::str::from_utf8(value) - .map_err(|_| Error::InvalidEntry)? + .map_err(Error::invalid_entry)? .parse() - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } pub(crate) fn value_text(value: &redis::Value) -> Option { @@ -158,6 +156,6 @@ fn value_bytes(value: &redis::Value) -> Result, Error> { redis::Value::BulkString(bytes) => Ok(bytes.clone()), redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()), redis::Value::Int(value) => Ok(value.to_string().into_bytes()), - _ => Err(Error::InvalidEntry), + _ => Err(Error::InvalidEntry(None)), } } diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs b/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs index 1c693d0875d..10e32c0183b 100644 --- a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs @@ -83,10 +83,7 @@ fn ok() -> redis::RedisResult { } fn already_exists() -> redis::RedisResult { - Err(redis::RedisError::from(( - redis::ErrorKind::Io, - "already exists", - ))) + Err(support::server_error("Index already exists")) } fn bulk(value: &[u8]) -> redis::Value { @@ -250,19 +247,29 @@ fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: Sema #[rstest] #[case::nested_matching(nested_dimension_info(3), Ok(()))] -#[case::nested_mismatch(nested_dimension_info(2), Err(Error::Unavailable))] +#[case::nested_mismatch( + nested_dimension_info(2), + Err("cache is unavailable: index exists with dimension 2, not 3") +)] #[case::flat_matching(flat_dimension_info("3"), Ok(()))] -#[case::flat_mismatch(flat_dimension_info("2"), Err(Error::Unavailable))] +#[case::flat_mismatch( + flat_dimension_info("2"), + Err("cache is unavailable: index exists with dimension 2, not 3") +)] #[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))] fn existing_index_dimension_must_match_embedding( #[case] info: redis::Value, - #[case] expected: Result<(), Error>, + #[case] expected: Result<(), &str>, entry: Value, context: SemanticCacheContext, ) { let recording = recording([already_exists(), Ok(info)]); - assert_eq!(recording.cache.set_cache("key", entry, &context), expected); + let result = recording.cache.set_cache("key", entry, &context); + assert_eq!( + result.map_err(|error| litellm_error::Report::new(&error).to_string()), + expected.map_err(str::to_owned) + ); } #[rstest] @@ -274,10 +281,10 @@ fn index_creation_failures_are_unavailable( ) { let recording = recording([reply]); - assert_eq!( + assert!(matches!( recording.cache.set_cache("key", entry, &context), - Err(Error::Unavailable) - ); + Err(Error::Unavailable(Some(_))) + )); } #[rstest] @@ -291,10 +298,10 @@ fn index_creation_failures_are_unavailable( bulk(b"document"), redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]), ]), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) )] -#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry))] -#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry))] +#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))] +#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))] fn get_applies_threshold_and_decodes_entry( #[case] reply: redis::Value, #[case] expected: Result, Error>, @@ -302,7 +309,13 @@ fn get_applies_threshold_and_decodes_entry( ) { let recording = recording([ok(), Ok(reply)]); - assert_eq!(recording.cache.get_cache("key", &context), expected); + let kind = |result: Result, Error>| { + result.map_err(|error| std::mem::discriminant(&error)) + }; + assert_eq!( + kind(recording.cache.get_cache("key", &context)), + kind(expected) + ); } #[rstest] diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs index 2c6ff06b134..2af05943c98 100644 --- a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs +++ b/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs @@ -74,7 +74,7 @@ impl FakeSearch { "FT.CREATE" => { let name = text(1); if state.indexes.contains_key(&name) { - return Err(error("Index already exists")); + return Err(server_error("Index already exists")); } let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes()); let prefix = args[position("PREFIX").unwrap() + 2].clone(); @@ -94,7 +94,7 @@ impl FakeSearch { let index = state .indexes .get(&text(1)) - .ok_or_else(|| error("Unknown index name"))?; + .ok_or_else(|| server_error("Unknown index name"))?; Ok(index_info(index)) } "FT.DROPINDEX" => { @@ -118,7 +118,7 @@ impl FakeSearch { let index = state .indexes .get(&text(1)) - .ok_or_else(|| error("no such index"))?; + .ok_or_else(|| server_error("no such index"))?; let query = text(2); let tag = query_tag(&query); let params = args.iter().position(|arg| arg == b"PARAMS").unwrap(); @@ -155,7 +155,7 @@ impl FakeSearch { ])) } "PING" => Ok(redis::Value::SimpleString("PONG".into())), - _ => Err(error("unsupported command")), + _ => Err(server_error("unsupported command")), } } } @@ -192,8 +192,11 @@ impl redis::ConnectionLike for FakeSearch { } } -fn error(message: &'static str) -> redis::RedisError { - redis::RedisError::from((redis::ErrorKind::Extension, message)) +pub fn server_error(reply: &str) -> redis::RedisError { + redis::parse_redis_value(format!("-{reply}\r\n").as_bytes()) + .expect("error reply") + .extract_error() + .expect_err("server error") } fn bulk(bytes: &[u8]) -> redis::Value { diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index 51bbe12a85e..282c34e6b02 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } strum.workspace = true diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs index ac9d8fc8764..b30f9e162f2 100644 --- a/litellm-rust/crates/cache/src/capabilities.rs +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -48,7 +48,7 @@ pub trait BatchCache: BaseCache { .map(|key| match self.get_cache(key, context) { Ok(Some(value)) => Ok(BatchEntry::Hit(value)), Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), Err(error) => Err(error), }) .collect() @@ -65,7 +65,7 @@ pub trait BatchCache: BaseCache { entries.push(match self.async_get_cache(&key, &context).await { Ok(Some(value)) => BatchEntry::Hit(value), Ok(None) => BatchEntry::Miss, - Err(Error::InvalidEntry) => BatchEntry::Invalid, + Err(Error::InvalidEntry(_)) => BatchEntry::Invalid, Err(error) => return Err(error), }); } diff --git a/litellm-rust/crates/cache/src/codec.rs b/litellm-rust/crates/cache/src/codec.rs index 6d47c682406..b802e40a8b3 100644 --- a/litellm-rust/crates/cache/src/codec.rs +++ b/litellm-rust/crates/cache/src/codec.rs @@ -41,10 +41,10 @@ where type Value = V; fn encode(&self, value: &Self::Value) -> Result, Error> { - serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + serde_json::to_vec(value).map_err(Error::invalid_entry) } fn decode(&self, bytes: &[u8]) -> Result { - serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry) + serde_json::from_slice(bytes).map_err(Error::invalid_entry) } } diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs index c4bc953239d..3deee2560f4 100644 --- a/litellm-rust/crates/cache/src/dual.rs +++ b/litellm-rust/crates/cache/src/dual.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use crate::{ BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheContext, ClaimCache, CounterCache, - DeleteCache, Error, FlushCache, IncrementOperation, SetCache, TtlCache, + DeleteCache, Error, FlushCache, IncrementOperation, SetCache, ShortBatch, TtlCache, }; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] @@ -98,7 +98,7 @@ impl DualCache { fn remote(&self, result: Result) -> Result, Error> { match result { Ok(value) => Ok(Some(value)), - Err(Error::Unavailable) + Err(Error::Unavailable(_)) if self.remote_failure_policy == RemoteFailurePolicy::UseLocal => { Ok(None) @@ -136,7 +136,10 @@ where remote: Vec>, ) -> Result>, Error> { if missing.len() != remote.len() { - return Err(Error::Unavailable); + return Err(Error::unavailable(ShortBatch { + requested: missing.len(), + returned: remote.len(), + })); } for (index, entry) in missing.into_iter().zip(remote) { if let BatchEntry::Hit(value) = &entry { @@ -389,7 +392,10 @@ where } let values = self.l2.async_increment_pipeline(operations.clone()).await?; if values.len() != operations.len() { - return Err(Error::Unavailable); + return Err(Error::unavailable(ShortBatch { + requested: operations.len(), + returned: values.len(), + })); } for (operation, value) in operations.iter().zip(&values) { self.l1 diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index 79ef9cd18b1..8861651a1c8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -1,9 +1,11 @@ +use litellm_error::Source; + #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum Error { #[error("cache is unavailable")] - Unavailable, + Unavailable(#[source] Option), #[error("invalid cache entry")] - InvalidEntry, + InvalidEntry(#[source] Option), #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, #[error("operation is not supported by this cache")] @@ -11,3 +13,59 @@ pub enum Error { #[error("semantic cache requires request messages")] MissingPrompt, } + +#[derive(Debug, thiserror::Error)] +#[error("cache backend returned {returned} results for {requested} keys")] +pub struct ShortBatch { + pub requested: usize, + pub returned: usize, +} + +#[derive(Debug, thiserror::Error)] +#[error("cache lock is poisoned")] +struct Poisoned; + +impl From> for Error { + fn from(_: std::sync::PoisonError) -> Self { + Self::unavailable(Poisoned) + } +} + +impl Error { + pub fn unavailable(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::Unavailable(Some(Source::new(cause))) + } + + pub fn invalid_entry(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::InvalidEntry(Some(Source::new(cause))) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::Error; + + #[rstest] + #[case::unavailable( + Error::unavailable(std::io::Error::other("connection refused")), + "cache is unavailable" + )] + #[case::invalid_entry( + Error::invalid_entry(std::io::Error::other("connection refused")), + "invalid cache entry" + )] + fn the_backend_cause_is_the_source_and_rendered_once( + #[case] error: Error, + #[case] message: &str, + ) { + let report = litellm_error::Report::new(&error).to_string(); + assert_eq!(report, format!("{message}: connection refused")); + assert!(!error.to_string().contains("connection refused")); + let cause = std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .expect("cause"); + assert!(cause.downcast_ref::().is_some()); + } +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index 55d8a5fa28b..4e07f3b4e8c 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -23,4 +23,4 @@ pub use codec::{CacheCodec, JsonCodec}; pub use dual::{ DEFAULT_DELETE_BATCH_SIZE, DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy, }; -pub use error::Error; +pub use error::{Error, ShortBatch}; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index baf968f4659..84a3b8c4d2a 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -67,7 +67,7 @@ impl BaseCache for TestCache { } fn set_cache(&self, _: &str, _: Self::Value, _: &ExactCacheContext) -> Result<(), Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } async fn async_set_cache( @@ -77,7 +77,7 @@ impl BaseCache for TestCache { context: ExactCacheContext, ) -> Result<(), Error> { if key == "unavailable" { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } self.writes .lock() @@ -118,7 +118,7 @@ impl BaseCache for RecordingCounter { impl CounterCache for RecordingCounter { fn increment_cache(&self, key: &str, amount: f64, _: ExactCacheContext) -> Result { if key == "unavailable" { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } let mut total = self.total.lock().unwrap(); *total += amount; @@ -226,7 +226,7 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure(cache: Te context.clone(), ) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert_eq!( *cache.writes.lock().unwrap(), @@ -323,7 +323,7 @@ async fn default_increment_pipeline_stops_at_the_first_failure() { operation("skipped", 1.0, None), ]) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); let keys = cache .increments diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs index 5320545fd49..ff4aeed87e6 100644 --- a/litellm-rust/crates/cache/tests/codec.rs +++ b/litellm-rust/crates/cache/tests/codec.rs @@ -31,12 +31,18 @@ fn json_codec_round_trips_typed_domain_values() { #[case::wrongly_typed(br#"{"deployment":12}"#)] fn json_codec_rejects_malformed_and_wrongly_typed_entries(#[case] bytes: &[u8]) { let codec = JsonCodec::::new(); - assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry); + assert!(matches!( + codec.decode(bytes), + Err(Error::InvalidEntry(Some(_))) + )); } #[rstest] fn json_codec_propagates_encoding_errors() { let codec = JsonCodec::>::new(); let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]); - assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry); + assert!(matches!( + codec.encode(&value), + Err(Error::InvalidEntry(Some(_))) + )); } diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs index 87dd76388ef..86af5e80abb 100644 --- a/litellm-rust/crates/cache/tests/dual.rs +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -71,7 +71,7 @@ where impl CounterCache for TestCache { fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result { if self.fail { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } let mut value = self.value.lock().unwrap(); let incremented = value.unwrap_or_default() + amount; @@ -92,7 +92,7 @@ where _: ExactCacheContext, ) -> Result { if self.fail { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } let mut value = self.value.lock().unwrap(); let winner = match value.as_ref() { @@ -113,7 +113,7 @@ fn failed_l2_increment_leaves_l1_unchanged() { assert_eq!( cache.increment_cache("counter", 2.0, ExactCacheContext::default()), - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert_eq!( l1.get_cache("counter", &ExactCacheContext::default()) @@ -276,11 +276,11 @@ impl BaseCache for Unavailable { } fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } } @@ -288,13 +288,13 @@ impl BatchCache for Unavailable {} impl DeleteCache for Unavailable { fn delete_cache(&self, _: &str) -> Result<(), Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } } impl FlushCache for Unavailable { fn flush_cache(&self) -> Result<(), Error> { - Err(Error::Unavailable) + Err(Error::Unavailable(None)) } } @@ -306,7 +306,7 @@ impl ClaimCache for Unavailable { _: &[String], _: ExactCacheContext, ) -> Result { - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) } } @@ -316,9 +316,12 @@ fn remote_failure_policy_selects_propagation_or_the_local_tier() { let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable)); assert_eq!( strict.set_cache("key", "value".into(), &context), - Err(Error::Unavailable) + Err(Error::Unavailable(None)) + ); + assert_eq!( + strict.get_cache("key", &context), + Err(Error::Unavailable(None)) ); - assert_eq!(strict.get_cache("key", &context), Err(Error::Unavailable)); let l1 = Arc::new(TestCache::new(None, false)); let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable)) @@ -347,7 +350,7 @@ fn claim_fallback_does_not_hide_non_availability_errors() { &[], ExactCacheContext::default() ), - Err(Error::InvalidEntry) + Err(Error::InvalidEntry(None)) ); } @@ -419,7 +422,7 @@ impl Tier { fn check(&self) -> Result<(), Error> { if self.fail { - return Err(Error::Unavailable); + return Err(Error::Unavailable(None)); } Ok(()) } @@ -598,7 +601,7 @@ async fn failed_async_l2_increment_leaves_l1_unchanged(log: Log) { cache .async_increment("counter", 2.0, seconds(30), true) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert_eq!( l1.counter("counter"), @@ -669,7 +672,7 @@ async fn failed_l2_increment_pipeline_leaves_l1_unchanged(log: Log) { ttl: None, }]) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert_eq!(events(&log), ["l2 pipeline a"]); assert_eq!(l1.counter("a"), Some((1.0, ExactCacheContext::default()))); @@ -713,7 +716,7 @@ async fn failed_local_set_add_never_reaches_l2(log: Log) { cache .async_set_cache_sadd("members", vec!["a".into()], None) .await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert!(events(&log).is_empty()); } @@ -768,7 +771,7 @@ async fn bulk_delete_stops_before_l2_when_the_local_delete_fails(log: Log) { ); assert_eq!( cache.delete_cache_keys(vec!["a".into()]).await, - Err(Error::Unavailable) + Err(Error::Unavailable(None)) ); assert!(events(&log).is_empty()); } diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index ed5e0fb9691..1cd05710366 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -6,12 +6,14 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true litellm-host.workspace = true litellm-host-python.workspace = true pyo3.workspace = true strum.workspace = true serde_json.workspace = true +thiserror.workspace = true [dev-dependencies] litellm-auth.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index d891b5a1f0c..493e91bb5ea 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -134,9 +134,10 @@ impl LegacyLogging { } fn logger(&self) -> PyResult<&PythonLogger> { - self.logger.as_ref().ok_or_else(|| { - pyo3::exceptions::PyRuntimeError::new_err("call logging is not initialized") - }) + Ok(self + .logger + .as_ref() + .ok_or(crate::error::Error::LoggingNotInitialized)?) } /// The keyword view the rest of the call reads: a copy, so the deployment hook's own diff --git a/litellm-rust/crates/callbacks-legacy-python/src/error.rs b/litellm-rust/crates/callbacks-legacy-python/src/error.rs new file mode 100644 index 00000000000..97ce62ae99c --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/error.rs @@ -0,0 +1,17 @@ +use litellm_error::Report; +use pyo3::{PyErr, exceptions::PyRuntimeError}; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Error { + #[error("call logging is not initialized")] + LoggingNotInitialized, +} + +impl From for PyErr { + fn from(error: Error) -> Self { + let message = Report::new(&error).to_string(); + match error { + Error::LoggingNotInitialized => PyRuntimeError::new_err(message), + } + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 38c6b1aedbd..f453237b469 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -11,6 +11,7 @@ mod adapter; mod call; mod callbacks; mod deferred; +mod error; mod logger; mod mapping; mod python; diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 9746a002bb9..5b0cd22332f 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -16,6 +16,7 @@ thiserror.workspace = true tracing.workspace = true [dev-dependencies] +litellm-error.workspace = true rstest.workspace = true tempfile.workspace = true tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs index 95c33967cf0..e5219ac992c 100644 --- a/litellm-rust/crates/config/src/error.rs +++ b/litellm-rust/crates/config/src/error.rs @@ -1,9 +1,16 @@ +use std::{io, path::PathBuf}; + #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("could not read config")] - Read(#[from] std::io::Error), + #[error("could not read config {}", path.display())] + Read { path: PathBuf, source: io::Error }, #[error("invalid YAML config")] Parse(#[from] serde_yaml_ng::Error), + #[error("invalid YAML config {}", path.display())] + ParseFile { + path: PathBuf, + source: serde_yaml_ng::Error, + }, #[error("Redis host does not yet support configured setting {0}")] RedisSetting(String), #[error("invalid Redis setting {setting}: {reason}")] diff --git a/litellm-rust/crates/config/src/includes.rs b/litellm-rust/crates/config/src/includes.rs index 78df8e9a33a..27b908650df 100644 --- a/litellm-rust/crates/config/src/includes.rs +++ b/litellm-rust/crates/config/src/includes.rs @@ -15,11 +15,31 @@ struct Includes { } fn read(path: &Path) -> Result { - Ok(serde_yaml_ng::from_str(&std::fs::read_to_string(path)?)?) + let text = std::fs::read_to_string(path).map_err(|source| Error::Read { + path: path.to_owned(), + source, + })?; + serde_yaml_ng::from_str(&text).map_err(|source| Error::ParseFile { + path: path.to_owned(), + source, + }) +} + +fn canonical(path: &Path) -> Result { + path.canonicalize().map_err(|source| Error::Read { + path: path.to_owned(), + source, + }) } fn entries(config: &Mapping, path: &Path) -> Result, Error> { - let includes: Includes = serde_yaml_ng::from_value(Value::Mapping(config.clone()))?; + let includes: Includes = + serde_yaml_ng::from_value(Value::Mapping(config.clone())).map_err(|source| { + Error::ParseFile { + path: path.to_owned(), + source, + } + })?; Ok(includes .include .into_iter() @@ -28,7 +48,7 @@ fn entries(config: &Mapping, path: &Path) -> Result, Erro } pub(super) fn load(path: &Path) -> Result { - let root = path.canonicalize()?; + let root = canonical(path)?; let mut merged = read(&root)?; let mut pending: VecDeque<_> = entries(&merged, &root)?.into(); let mut loaded = BTreeSet::from([root.clone()]); @@ -36,12 +56,11 @@ pub(super) fn load(path: &Path) -> Result { while let Some((entry, declaring)) = pending.pop_front() { let declared = declaring.parent().unwrap_or(Path::new(".")).join(&entry); let fallback = root.parent().unwrap_or(Path::new(".")).join(&entry); - let location = if declared.exists() { - declared + let location = canonical(if declared.exists() { + &declared } else { - fallback - } - .canonicalize()?; + &fallback + })?; if !loaded.insert(location.clone()) { continue; } diff --git a/litellm-rust/crates/config/tests/config.rs b/litellm-rust/crates/config/tests/config.rs index d425b62efc2..3c7d02203e5 100644 --- a/litellm-rust/crates/config/tests/config.rs +++ b/litellm-rust/crates/config/tests/config.rs @@ -84,9 +84,15 @@ fn rejects_malformed_and_incomplete_config(#[case] yaml: &str) { #[rstest] fn distinguishes_read_errors_from_parse_errors(directory: TempDir) { + let error = Config::load(directory.path().join("missing.yaml")).unwrap_err(); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("missing.yaml") + ); assert!(matches!( - Config::load(directory.path().join("missing.yaml")), - Err(Error::Read(error)) if error.kind() == std::io::ErrorKind::NotFound + error, + Error::Read { source, .. } if source.kind() == std::io::ErrorKind::NotFound )); } diff --git a/litellm-rust/crates/core-utils/Cargo.toml b/litellm-rust/crates/core-utils/Cargo.toml index 1feea5fcc08..2e5857c79f3 100644 --- a/litellm-rust/crates/core-utils/Cargo.toml +++ b/litellm-rust/crates/core-utils/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] fancy-regex.workspace = true +litellm-error.workspace = true litellm-tracing.workspace = true litellm-llms-types.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/core-utils/src/call_arguments.rs b/litellm-rust/crates/core-utils/src/call_arguments.rs index 062dad0a92a..3277ca26c7f 100644 --- a/litellm-rust/crates/core-utils/src/call_arguments.rs +++ b/litellm-rust/crates/core-utils/src/call_arguments.rs @@ -1,5 +1,7 @@ use std::ops::Deref; +use litellm_error::Source; + use serde::{Deserialize, Serialize, de::DeserializeOwned}; use serde_json::{Map, Value}; @@ -17,9 +19,11 @@ impl CallArguments { } #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -#[error("invalid argument: {path}")] +#[error("invalid argument {path}")] pub struct ArgumentError { pub path: String, + #[source] + source: Source, } pub fn parse_options(arguments: &CallArguments) -> Result { @@ -28,6 +32,7 @@ pub fn parse_options(arguments: &CallArguments) -> Result( body: &B, consumed: &[&str], ) -> Result { - let Value::Object(fields) = - serde_json::to_value(body).map_err(|_| crate::params::Error::Body)? + let Value::Object(fields) = serde_json::to_value(body) + .map_err(|error| crate::params::Error::Serialize(Source::new(error)))? else { return Err(crate::params::Error::Body); }; @@ -183,9 +188,9 @@ mod tests { assert_eq!(arguments.get("enabled"), Some(&Value::Null)); assert_eq!(arguments.get("missing"), None); let invalid = serde_json::from_value(json!({"enabled":0})).unwrap(); - assert_eq!( - parse_options::(&invalid).err().unwrap().path, - "enabled" - ); + let error = parse_options::(&invalid).err().unwrap(); + assert_eq!(error.path, "enabled"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.starts_with("invalid argument enabled: invalid type")); } } diff --git a/litellm-rust/crates/core-utils/src/decline.rs b/litellm-rust/crates/core-utils/src/decline.rs new file mode 100644 index 00000000000..6f398b16a06 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/decline.rs @@ -0,0 +1,16 @@ +/// Rust core cannot serve this request, and nothing was sent to the provider yet. +/// +/// Every error on the call path carries it unchanged, so the host can hand the same request +/// to another implementation: the SDK runs legacy Python, the gateway forwards to the proxy. +/// Raise it only before the provider call. Anything after the send is terminal. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +#[error("declined by the rust path: {reason}")] +pub struct Decline { + pub reason: &'static str, +} + +impl Decline { + pub const fn new(reason: &'static str) -> Self { + Self { reason } + } +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs index 8b63fadf5a6..c99bbf78e02 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs @@ -337,15 +337,15 @@ mod tests { #[case::request_timed_out("Request timed out")] #[case::timed_out_generating("Timed out generating response")] #[case::read_operation("The read operation timed out")] - fn timeout_markers_win_over_every_family(#[case] marker: &str) { + fn timeout_markers_win_over_every_family( + #[case] marker: &str, + #[values("mistral", "vertex_ai", "cohere", "reducto")] provider: &str, + ) { let body = format!("rate limit invalid api token {marker}"); - for provider in ["mistral", "vertex_ai", "cohere", "reducto"] { - assert_eq!( - mapped(provider, &http(429, &body)).error, - PublicError::Timeout { status: 408 }, - "{provider}" - ); - } + assert_eq!( + mapped(provider, &http(429, &body)).error, + PublicError::Timeout { status: 408 } + ); } #[test] 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/core-utils/src/lib.rs b/litellm-rust/crates/core-utils/src/lib.rs index a937f55654e..cfbf451bf1b 100644 --- a/litellm-rust/crates/core-utils/src/lib.rs +++ b/litellm-rust/crates/core-utils/src/lib.rs @@ -1,5 +1,6 @@ pub mod call_arguments; pub mod core_helpers; +pub mod decline; pub mod dot_notation_indexing; pub mod exception_mapping_utils; pub mod get_llm_provider_logic; diff --git a/litellm-rust/crates/core-utils/src/params.rs b/litellm-rust/crates/core-utils/src/params.rs index 9545a3ef17b..f163399b878 100644 --- a/litellm-rust/crates/core-utils/src/params.rs +++ b/litellm-rust/crates/core-utils/src/params.rs @@ -4,6 +4,8 @@ pub enum Error { ExtraBody, #[error("invalid request: body must be a JSON object")] Body, + #[error("invalid request: body could not be serialized")] + Serialize(#[source] litellm_error::Source), } use std::ops::{Deref, DerefMut}; diff --git a/litellm-rust/crates/core-utils/src/url_utils.rs b/litellm-rust/crates/core-utils/src/url_utils.rs index ebbe7fca00b..63e4f041a72 100644 --- a/litellm-rust/crates/core-utils/src/url_utils.rs +++ b/litellm-rust/crates/core-utils/src/url_utils.rs @@ -4,7 +4,7 @@ use url::Url; #[derive(Debug, thiserror::Error)] pub enum ApiUrlError { - #[error("invalid URL: {0}")] + #[error("invalid URL")] Parse(#[from] url::ParseError), #[error("URL cannot be used as a base")] CannotBeBase, diff --git a/litellm-rust/crates/cost/Cargo.toml b/litellm-rust/crates/cost/Cargo.toml index 3ed27e82c9a..6921e6151b5 100644 --- a/litellm-rust/crates/cost/Cargo.toml +++ b/litellm-rust/crates/cost/Cargo.toml @@ -5,6 +5,9 @@ edition.workspace = true license.workspace = true repository.workspace = true +[dependencies] +thiserror.workspace = true + [dev-dependencies] criterion.workspace = true proptest.workspace = true diff --git a/litellm-rust/crates/cost/src/error.rs b/litellm-rust/crates/cost/src/error.rs new file mode 100644 index 00000000000..1753080617f --- /dev/null +++ b/litellm-rust/crates/cost/src/error.rs @@ -0,0 +1,27 @@ +#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)] +pub enum Error { + #[error("model has no input rate")] + MissingInputRate, + #[error("model has no output rate")] + MissingOutputRate, + #[error("rate must be finite and nonnegative")] + InvalidRate, + #[error("region multiplier must be finite and positive")] + InvalidRegionMultiplier, + #[error("off-peak pricing needs a billing minute within the day")] + InvalidBillingTime, + #[error("off-peak window must start before it ends within the day")] + InvalidOffPeakWindow, + #[error("cached tokens exceed prompt tokens")] + CacheExceedsPrompt, + #[error("cache write details do not add up to the cache write tokens")] + InvalidCacheWriteDetails, + #[error("token count or cost overflows")] + TokenCountOverflow, + #[error("service tier is repeated or cannot carry its own rates")] + DuplicateTier, + #[error("prompt token threshold is repeated")] + DuplicateThreshold, + #[error("service tier is repeated or cannot carry its own rates within a threshold")] + DuplicateThresholdTier, +} diff --git a/litellm-rust/crates/cost/src/lib.rs b/litellm-rust/crates/cost/src/lib.rs index 32cc755ce15..24c69afd9d4 100644 --- a/litellm-rust/crates/cost/src/lib.rs +++ b/litellm-rust/crates/cost/src/lib.rs @@ -1,3 +1,7 @@ +mod error; + +pub use error::Error; + #[derive(Clone, Copy, Debug, PartialEq)] pub enum Rate { Missing, @@ -156,22 +160,6 @@ pub struct EffectiveRates { pub cache_write_1h: f64, } -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum PricingError { - MissingInputRate, - MissingOutputRate, - InvalidRate, - InvalidRegionMultiplier, - InvalidBillingTime, - InvalidOffPeakWindow, - CacheExceedsPrompt, - InvalidCacheWriteDetails, - TokenCountOverflow, - DuplicateTier, - DuplicateThreshold, - DuplicateThresholdTier, -} - fn selected_tier(tier: ServiceTier) -> ServiceTier { if tier == ServiceTier::Fast { ServiceTier::Priority @@ -210,9 +198,9 @@ fn valid_rates(rates: Rates) -> bool { }) } -fn validate_tiers(tiers: &[TierRates], duplicate: PricingError) -> Result<(), PricingError> { +fn validate_tiers(tiers: &[TierRates], duplicate: Error) -> Result<(), Error> { if tiers.iter().any(|entry| !valid_rates(entry.rates)) { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } if tiers.iter().enumerate().any(|(index, entry)| { matches!( @@ -227,20 +215,20 @@ fn validate_tiers(tiers: &[TierRates], duplicate: PricingError) -> Result<(), Pr Ok(()) } -pub fn compile(pricing: &Pricing<'_>) -> Result { +pub fn compile(pricing: &Pricing<'_>) -> Result { if !valid_rates(pricing.standard) { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } - validate_tiers(pricing.tiers, PricingError::DuplicateTier)?; + validate_tiers(pricing.tiers, Error::DuplicateTier)?; if let Some(window) = pricing.off_peak { if window.start_utc_minute >= 1440 || window.end_utc_minute > 1440 || window.start_utc_minute >= window.end_utc_minute { - return Err(PricingError::InvalidOffPeakWindow); + return Err(Error::InvalidOffPeakWindow); } if !valid_rates(window.rates) { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } } let mut thresholds: Vec<_> = pricing @@ -248,9 +236,9 @@ pub fn compile(pricing: &Pricing<'_>) -> Result { .iter() .map(|entry| { if !valid_rates(entry.standard) { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } - validate_tiers(entry.tiers, PricingError::DuplicateThresholdTier)?; + validate_tiers(entry.tiers, Error::DuplicateThresholdTier)?; Ok(CompiledThreshold { above_prompt_tokens: entry.above_prompt_tokens, standard: entry.standard, @@ -263,7 +251,7 @@ pub fn compile(pricing: &Pricing<'_>) -> Result { .windows(2) .any(|pair| pair[0].above_prompt_tokens == pair[1].above_prompt_tokens) { - return Err(PricingError::DuplicateThreshold); + return Err(Error::DuplicateThreshold); } Ok(PricingPlan { standard: pricing.standard, @@ -274,11 +262,7 @@ pub fn compile(pricing: &Pricing<'_>) -> Result { } impl PricingPlan { - fn resolve_rates( - &self, - request: &Request, - threshold_tokens: u64, - ) -> Result { + fn resolve_rates(&self, request: &Request, threshold_tokens: u64) -> Result { let tier = selected_tier(request.service_tier); let base = self .tiers @@ -305,13 +289,13 @@ impl PricingPlan { || window.end_utc_minute > 1440 || window.start_utc_minute >= window.end_utc_minute { - return Err(PricingError::InvalidOffPeakWindow); + return Err(Error::InvalidOffPeakWindow); } let minute = request .billed_at_utc_minute - .ok_or(PricingError::InvalidBillingTime)?; + .ok_or(Error::InvalidBillingTime)?; if minute >= 1440 { - return Err(PricingError::InvalidBillingTime); + return Err(Error::InvalidBillingTime); } if (window.start_utc_minute..window.end_utc_minute).contains(&minute) { Ok(window.rates.overlay(selected)) @@ -322,26 +306,26 @@ impl PricingPlan { } } - fn checked_rate(rate: Rate, missing: PricingError) -> Result { + fn checked_rate(rate: Rate, missing: Error) -> Result { let value = rate.value().ok_or(missing)?; if !value.is_finite() || value < 0.0 { - return Err(PricingError::InvalidRate); + return Err(Error::InvalidRate); } Ok(value) } - pub fn calculate(&self, request: &Request) -> Result { + pub fn calculate(&self, request: &Request) -> Result { let usage = request.usage; let cached = usage .cache_read_tokens .checked_add(usage.cache_write_tokens) - .ok_or(PricingError::TokenCountOverflow)?; + .ok_or(Error::TokenCountOverflow)?; let (regular, threshold_tokens) = match usage.prompt_convention { PromptConvention::IncludesCache => ( usage .prompt_tokens .checked_sub(cached) - .ok_or(PricingError::CacheExceedsPrompt)?, + .ok_or(Error::CacheExceedsPrompt)?, usage.prompt_tokens, ), PromptConvention::ExcludesCache => ( @@ -349,7 +333,7 @@ impl PricingPlan { usage .prompt_tokens .checked_add(cached) - .ok_or(PricingError::TokenCountOverflow)?, + .ok_or(Error::TokenCountOverflow)?, ), }; let writes = match (usage.cache_write_5m_tokens, usage.cache_write_1h_tokens) { @@ -357,26 +341,20 @@ impl PricingPlan { (Some(five), Some(one)) if five.checked_add(one) == Some(usage.cache_write_tokens) => { (five, one) } - _ => return Err(PricingError::InvalidCacheWriteDetails), + _ => return Err(Error::InvalidCacheWriteDetails), }; let rates = self.resolve_rates(request, threshold_tokens)?; - let input = Self::checked_rate(rates.input, PricingError::MissingInputRate)?; - let output = Self::checked_rate(rates.output, PricingError::MissingOutputRate)?; - let read = Self::checked_rate( - rates.cache_read.or(rates.input), - PricingError::MissingInputRate, - )?; - let write = Self::checked_rate( - rates.cache_write.or(rates.input), - PricingError::MissingInputRate, - )?; + let input = Self::checked_rate(rates.input, Error::MissingInputRate)?; + let output = Self::checked_rate(rates.output, Error::MissingOutputRate)?; + let read = Self::checked_rate(rates.cache_read.or(rates.input), Error::MissingInputRate)?; + let write = Self::checked_rate(rates.cache_write.or(rates.input), Error::MissingInputRate)?; let write_1h = Self::checked_rate( rates.cache_write_1h.or(rates.cache_write).or(rates.input), - PricingError::MissingInputRate, + Error::MissingInputRate, )?; let multiplier = request.region_multiplier.unwrap_or(1.0); if !multiplier.is_finite() || multiplier <= 0.0 { - return Err(PricingError::InvalidRegionMultiplier); + return Err(Error::InvalidRegionMultiplier); } let cost = Cost { uncached_input: regular as f64 * input, @@ -394,12 +372,12 @@ impl PricingPlan { }, }; if !cost.total().is_finite() { - return Err(PricingError::TokenCountOverflow); + return Err(Error::TokenCountOverflow); } Ok(cost) } } -pub fn calculate(pricing: &Pricing<'_>, request: &Request) -> Result { +pub fn calculate(pricing: &Pricing<'_>, request: &Request) -> Result { compile(pricing)?.calculate(request) } diff --git a/litellm-rust/crates/cost/tests/calculation.rs b/litellm-rust/crates/cost/tests/calculation.rs index f8152483a7c..af1a73d8c61 100644 --- a/litellm-rust/crates/cost/tests/calculation.rs +++ b/litellm-rust/crates/cost/tests/calculation.rs @@ -1,5 +1,5 @@ use litellm_cost::{ - OffPeakRates, Pricing, PricingError, PromptConvention, Rate, Rates, Request, ServiceTier, + Error, OffPeakRates, Pricing, PromptConvention, Rate, Rates, Request, ServiceTier, ThresholdPolicy, ThresholdRates, TierRates, Usage, calculate, compile, }; use rstest::rstest; @@ -132,7 +132,7 @@ fn split_writes_and_invalid_accounting() { }; assert_eq!( calculate(&pricing(standard), &overlapping), - Err(PricingError::CacheExceedsPrompt) + Err(Error::CacheExceedsPrompt) ); let incomplete = Request { usage: Usage { @@ -143,7 +143,7 @@ fn split_writes_and_invalid_accounting() { }; assert_eq!( calculate(&pricing(standard), &incomplete), - Err(PricingError::InvalidCacheWriteDetails) + Err(Error::InvalidCacheWriteDetails) ); } @@ -297,10 +297,10 @@ fn compile_rejects_ambiguous_rates() { }; assert_eq!( compile(&specification).err(), - Some(PricingError::DuplicateThreshold) + Some(Error::DuplicateThreshold) ); let invalid = pricing(rates(Rate::Value(f64::NAN), Rate::Value(1.0))); - assert_eq!(compile(&invalid).err(), Some(PricingError::InvalidRate)); + assert_eq!(compile(&invalid).err(), Some(Error::InvalidRate)); } #[test] @@ -326,7 +326,7 @@ fn off_peak_is_one_non_wrapping_utc_window() { }; assert_eq!( calculate(&specification, &base), - Err(PricingError::InvalidBillingTime) + Err(Error::InvalidBillingTime) ); assert_eq!(calculate(&specification, &start).unwrap().input(), 100.0); assert_eq!(calculate(&specification, &end).unwrap().input(), 200.0); @@ -347,11 +347,11 @@ fn missing_rates_and_free_rates_remain_distinct() { }; assert_eq!( calculate(&pricing(Rates::EMPTY), &empty), - Err(PricingError::MissingInputRate) + Err(Error::MissingInputRate) ); assert_eq!( calculate(&pricing(rates(Rate::Value(0.0), Rate::Missing)), &empty), - Err(PricingError::MissingOutputRate) + Err(Error::MissingOutputRate) ); assert_eq!( calculate(&pricing(rates(Rate::Value(0.0), Rate::Value(0.0))), &empty) @@ -481,7 +481,7 @@ fn regional_multiplier_applies_after_input_components_are_summed() { }; assert_eq!( calculate(&pricing(standard), &invalid), - Err(PricingError::InvalidRegionMultiplier) + Err(Error::InvalidRegionMultiplier) ); } @@ -531,5 +531,5 @@ fn compilation_sorts_thresholds_and_rejects_duplicate_tiers() { thresholds: &[], ..specification }; - assert_eq!(compile(&invalid).err(), Some(PricingError::DuplicateTier)); + assert_eq!(compile(&invalid).err(), Some(Error::DuplicateTier)); } diff --git a/litellm-rust/crates/error/Cargo.toml b/litellm-rust/crates/error/Cargo.toml new file mode 100644 index 00000000000..fa2d26a1b99 --- /dev/null +++ b/litellm-rust/crates/error/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "litellm-error" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "Shared error rendering and type-erased causes for litellm-rust boundaries" + +[dependencies] +thiserror.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/error/src/lib.rs b/litellm-rust/crates/error/src/lib.rs new file mode 100644 index 00000000000..ec8d5c3e1f4 --- /dev/null +++ b/litellm-rust/crates/error/src/lib.rs @@ -0,0 +1,132 @@ +use std::{error::Error, fmt, ops::Deref, sync::Arc}; + +pub struct Report<'a>(&'a (dyn Error + 'static)); + +impl<'a> Report<'a> { + pub fn new(error: &'a (dyn Error + 'static)) -> Self { + Self(error) + } +} + +impl fmt::Display for Report<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "{}", self.0)?; + std::iter::successors(self.0.source(), |&cause| cause.source()) + .try_for_each(|cause| write!(formatter, ": {cause}")) + } +} + +pub fn find<'a, T: Error + 'static>(error: &'a (dyn Error + 'static)) -> Option<&'a T> { + std::iter::successors(Some(error), |&cause| cause.source()).find_map(|cause| { + cause + .downcast_ref::() + .or_else(|| cause.downcast_ref::()?.downcast_ref::()) + }) +} + +#[derive(Clone, Debug, thiserror::Error)] +#[error(transparent)] +pub struct Source(Arc); + +impl Source { + pub fn new(error: impl Error + Send + Sync + 'static) -> Self { + Self(Arc::new(error)) + } +} + +impl From> for Source { + fn from(error: Box) -> Self { + Self(Arc::from(error)) + } +} + +impl Deref for Source { + type Target = dyn Error + Send + Sync; + + fn deref(&self) -> &Self::Target { + self.0.as_ref() + } +} + +impl PartialEq for Source { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +impl Eq for Source {} + +#[cfg(test)] +mod tests { + use std::io; + + use rstest::rstest; + + use super::{Report, Source, find}; + + #[derive(Debug, thiserror::Error)] + enum Leaf { + #[error("could not read the config")] + Read(#[source] io::Error), + #[error("could not parse the config")] + Parse(#[source] Source), + #[error("config is empty")] + Empty, + } + + #[derive(Debug, thiserror::Error)] + enum Outer { + #[error("startup failed")] + Startup(#[from] Leaf), + #[error(transparent)] + Nested(Leaf), + } + + fn denied() -> io::Error { + io::Error::new(io::ErrorKind::PermissionDenied, "permission denied") + } + + #[rstest] + #[case::leaf(Outer::Nested(Leaf::Empty), "config is empty")] + #[case::context_over_a_cause( + Outer::Startup(Leaf::Read(denied())), + "startup failed: could not read the config: permission denied" + )] + #[case::transparent_adds_no_text( + Outer::Nested(Leaf::Read(denied())), + "could not read the config: permission denied" + )] + #[case::erased_cause_keeps_its_chain( + Outer::Nested(Leaf::Parse(Source::new(Leaf::Read(denied())))), + "could not parse the config: could not read the config: permission denied" + )] + fn report_renders_every_layer_once(#[case] error: Outer, #[case] expected: &str) { + assert_eq!(Report::new(&error).to_string(), expected); + } + + #[rstest] + #[case::direct_source(Outer::Startup(Leaf::Read(denied())))] + #[case::behind_an_erased_source(Outer::Startup(Leaf::Parse(Source::new(Leaf::Read( + denied() + )))))] + #[case::behind_a_boxed_source(Outer::Startup(Leaf::Parse(Source::from( + Box::::from(denied()) + ))))] + #[case::behind_a_transparent_layer(Outer::Nested(Leaf::Parse(Source::new(denied()))))] + fn find_reaches_the_typed_cause_anywhere_in_the_chain(#[case] error: Outer) { + let cause = find::(&error).expect("io cause"); + assert_eq!(cause.kind(), io::ErrorKind::PermissionDenied); + } + + #[test] + fn find_answers_none_when_the_chain_has_no_such_cause() { + assert!(find::(&Outer::Nested(Leaf::Empty)).is_none()); + } + + #[test] + fn a_clone_shares_its_cause() { + let source = Source::new(denied()); + assert_eq!(source.clone(), source); + assert_ne!(Source::new(denied()), source); + } +} diff --git a/litellm-rust/crates/framer/src/error.rs b/litellm-rust/crates/framer/src/error.rs index 879d7557671..dde421d71af 100644 --- a/litellm-rust/crates/framer/src/error.rs +++ b/litellm-rust/crates/framer/src/error.rs @@ -1,21 +1,21 @@ #[cfg(feature = "sse")] #[derive(Debug, thiserror::Error)] pub enum SseError { - #[error("body stream failed: {0}")] + #[error("body stream failed")] Body(#[from] std::io::Error), - #[error("SSE field is not UTF-8: {0}")] + #[error("SSE field is not UTF-8")] InvalidUtf8(#[from] std::str::Utf8Error), } #[cfg(feature = "aws")] #[derive(Debug, thiserror::Error)] pub enum EventStreamError { - #[error("body stream failed: {0}")] + #[error("body stream failed")] Body(#[from] std::io::Error), #[error("invalid AWS EventStream frame length: {0}")] InvalidLength(usize), #[error("truncated AWS EventStream frame")] Truncated, - #[error("malformed AWS EventStream frame: {0}")] + #[error("malformed AWS EventStream frame")] Malformed(#[from] aws_smithy_eventstream::error::Error), } diff --git a/litellm-rust/crates/framer/src/framer.rs b/litellm-rust/crates/framer/src/framer.rs deleted file mode 100644 index 507aed54700..00000000000 --- a/litellm-rust/crates/framer/src/framer.rs +++ /dev/null @@ -1,13 +0,0 @@ -use futures_util::Stream; - -use crate::Error; - -pub trait Framer: Send { - type Frame: Send; - - fn frame(self, input: S) -> impl Stream> + Send - where - S: Stream> + Send, - B: bytes::Buf + Send, - E: std::error::Error + Send + Sync + 'static; -} diff --git a/litellm-rust/crates/gateway-auth/AGENTS.md b/litellm-rust/crates/gateway-auth/AGENTS.md index 58457a3151b..0a4bc6450ae 100644 --- a/litellm-rust/crates/gateway-auth/AGENTS.md +++ b/litellm-rust/crates/gateway-auth/AGENTS.md @@ -20,7 +20,9 @@ Local UI login continues using axum-login and tower-sessions. Only after session Authentication evidence and principals contain no raw token, password, request body, or mutable accounting state. Session ownership is scoped by principal authority, subject, verifier, and credential ID, so separate credentials do not silently share an MCP session. Credential rotation may retain ownership when the verifier preserves a stable credential ID. Scope and expiry checks still run for each operation -Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. HTTP adapters map those outcomes to status codes. Inference keeps its API-specific error envelopes and MCP translates policy failures to its own protocol +Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. Bearer tokens, virtual keys and UI sessions all fail with the one `Error`, classified once by `status()`, and `is_denied()` tells a rejected caller from an auth failure. Inference keeps its API-specific error envelopes, the UI renders `Error` in its own envelope and MCP translates policy failures to its own protocol + +The plain-text rejection body is the top-level message only, never the `Report` chain, because the caller is unauthenticated and the sources carry secret-manager and backend detail. That stays reachable through `source()` for logs Virtual-key storage, JWT verification, OAuth2 introspection, trusted-proxy validation, SSO, and custom Python hook adapters are not implemented here yet. They should supply these contracts rather than bypassing the shared authorization boundary. Request-body-dependent custom hooks will need a bounded endpoint adapter after parsing diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml index cf040c6eef3..27f20203843 100644 --- a/litellm-rust/crates/gateway-auth/Cargo.toml +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -10,6 +10,7 @@ axum = { workspace = true, features = ["matched-path"] } axum-login.workspace = true litellm-auth-types.workspace = true litellm-config.workspace = true +litellm-error.workspace = true litellm-secrets.workspace = true sha2.workspace = true subtle.workspace = true diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs index 262a666df80..7bb6292b18c 100644 --- a/litellm-rust/crates/gateway-auth/src/error.rs +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -2,44 +2,59 @@ use axum::{ http::StatusCode, response::{IntoResponse, Response}, }; - -#[derive(Debug, thiserror::Error)] -pub enum KeyError { - #[error("missing or invalid virtual key")] - Invalid, - #[error("virtual key has expired")] - Expired, - #[error("virtual key lookup unavailable")] - Lookup(#[source] Box), -} +use litellm_error::Source; #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("operation is not permitted")] - Forbidden, - #[error("credential has expired")] - Expired, - #[error("authenticated request context is missing")] - MissingIdentity, - #[error("authentication service unavailable")] - Unavailable, - #[error("gateway auth not configured")] - Unconfigured, #[error("missing or invalid bearer token")] InvalidToken, + #[error("invalid UI credentials or session")] + InvalidSession, + #[error("credential has expired")] + Expired, + #[error("operation is not permitted")] + Forbidden, + #[error("authentication service unavailable")] + Unavailable(#[source] Source), #[error("gateway authentication unavailable")] Secret(#[from] litellm_secrets::Error), + #[error("UI session unavailable")] + Session(#[from] tower_sessions::session::Error), + #[error("gateway auth not configured")] + Unconfigured, + #[error("UI_USERNAME and UI_PASSWORD must be nonempty")] + UiUnconfigured, + #[error("authenticated request context is missing")] + MissingIdentity, + #[error("UI session layer is not installed")] + MissingSessionLayer, } impl Error { pub fn status(&self) -> StatusCode { match self { - Self::InvalidToken | Self::Expired => StatusCode::UNAUTHORIZED, + Self::InvalidToken | Self::InvalidSession | Self::Expired => StatusCode::UNAUTHORIZED, Self::Forbidden => StatusCode::FORBIDDEN, - Self::Unavailable => StatusCode::SERVICE_UNAVAILABLE, - Self::MissingIdentity | Self::Unconfigured | Self::Secret(_) => { - StatusCode::INTERNAL_SERVER_ERROR - } + Self::Unavailable(_) => StatusCode::SERVICE_UNAVAILABLE, + Self::Secret(_) + | Self::Session(_) + | Self::Unconfigured + | Self::UiUnconfigured + | Self::MissingIdentity + | Self::MissingSessionLayer => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + pub fn is_denied(&self) -> bool { + match self { + Self::InvalidToken | Self::InvalidSession | Self::Expired | Self::Forbidden => true, + Self::Unavailable(_) + | Self::Secret(_) + | Self::Session(_) + | Self::Unconfigured + | Self::UiUnconfigured + | Self::MissingIdentity + | Self::MissingSessionLayer => false, } } } @@ -50,24 +65,31 @@ impl IntoResponse for Error { } } -#[derive(Debug, thiserror::Error)] -pub enum UiAuthError { - #[error("UI_USERNAME and UI_PASSWORD must be nonempty")] - Unconfigured, - #[error("invalid UI credentials or session")] - Unauthorized, - #[error("UI authentication unavailable")] - Unavailable, - #[error("UI session unavailable")] - Session(#[from] tower_sessions::session::Error), -} +#[cfg(test)] +mod tests { + use rstest::rstest; -impl IntoResponse for UiAuthError { - fn into_response(self) -> Response { - let status = match self { - Self::Unauthorized => StatusCode::UNAUTHORIZED, - _ => StatusCode::INTERNAL_SERVER_ERROR, - }; - (status, self.to_string()).into_response() + use super::*; + + #[rstest] + #[tokio::test] + async fn rejection_body_hides_backend_causes() { + let error = + Error::Unavailable(Source::new(std::io::Error::other("private-backend-detail"))); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("private-backend-detail") + ); + let response = error.into_response(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + let body = axum::body::to_bytes(response.into_body(), 4096) + .await + .unwrap(); + assert!( + !std::str::from_utf8(&body) + .unwrap() + .contains("private-backend-detail") + ); } } diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 46d8c02253f..1fe8fa60e0e 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -2,7 +2,7 @@ use std::{future::Future, pin::Pin, time::SystemTime}; use litellm_auth_types::SecretValue; -use crate::{KeyError, hash_token}; +use crate::{Error, hash_token}; #[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] pub struct KeyHash(String); @@ -24,7 +24,7 @@ pub enum KeyStatus { } pub type LookupFuture<'a> = - Pin, KeyError>> + Send + 'a>>; + Pin, Error>> + Send + 'a>>; pub trait KeyLookup: Send + Sync { fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a>; @@ -34,15 +34,15 @@ pub async fn verify_key( lookup: &dyn KeyLookup, token: &SecretValue, now: impl FnOnce() -> SystemTime, -) -> Result { +) -> Result { let hash = KeyHash::from_token(token); match lookup.lookup(&hash).await? { Some(KeyStatus::Active { expires_at }) => { if expires_at.is_some_and(|expiry| expiry <= now()) { - return Err(KeyError::Expired); + return Err(Error::Expired); } Ok(hash) } - Some(KeyStatus::Revoked) | None => Err(KeyError::Invalid), + Some(KeyStatus::Revoked) | None => Err(Error::InvalidToken), } } diff --git a/litellm-rust/crates/gateway-auth/src/lib.rs b/litellm-rust/crates/gateway-auth/src/lib.rs index 1b400a94746..31ce17b3aa4 100644 --- a/litellm-rust/crates/gateway-auth/src/lib.rs +++ b/litellm-rust/crates/gateway-auth/src/lib.rs @@ -16,7 +16,7 @@ pub use authorization::{ AccessRequest, AuthenticatedRequest, AuthorizedOperation, Authorizer, McpAction, NoAdditionalPolicy, Permissions, UiAction, }; -pub use error::{Error, KeyError, UiAuthError}; +pub use error::Error; pub use http::{Bearer, CredentialExtractor, RequireMasterKey, authenticate}; pub use identity::{ AuthenticatedCaller, Authentication, AuthenticationMethod, Principal, PrincipalKind, diff --git a/litellm-rust/crates/gateway-auth/src/ui.rs b/litellm-rust/crates/gateway-auth/src/ui.rs index 9e6b77a17d0..3d655bde187 100644 --- a/litellm-rust/crates/gateway-auth/src/ui.rs +++ b/litellm-rust/crates/gateway-auth/src/ui.rs @@ -11,9 +11,9 @@ use tower_sessions::Session; use veil::Redact; use crate::{ - AccessRequest, AuthenticatedRequest, Authentication, AuthenticationMethod, LocalAdministrator, - NoAdditionalPolicy, Permissions, Principal, PrincipalKind, SystemClock, UiAction, UiAuthError, - VerifiedIdentity, + AccessRequest, AuthenticatedRequest, Authentication, AuthenticationMethod, Error, + LocalAdministrator, NoAdditionalPolicy, Permissions, Principal, PrincipalKind, SystemClock, + UiAction, VerifiedIdentity, }; pub const UI_CSRF_KEY: &str = "litellm.ui.csrf"; @@ -50,9 +50,9 @@ pub struct UiBackend { } impl UiBackend { - pub fn new(username: String, password: SecretValue) -> Result { + pub fn new(username: String, password: SecretValue) -> Result { if username.trim().is_empty() || password.expose().trim().is_empty() { - return Err(UiAuthError::Unconfigured); + return Err(Error::UiUnconfigured); } Ok(Self { user: UiUser { @@ -89,28 +89,28 @@ pub struct UiSession { } impl FromRequestParts for UiSession { - type Rejection = UiAuthError; + type Rejection = Error; async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { let auth = UiAuthSession::from_request_parts(parts, state) .await - .map_err(|_| UiAuthError::Unavailable)?; - let user = auth.user.ok_or(UiAuthError::Unauthorized)?; + .map_err(|_| Error::MissingSessionLayer)?; + let user = auth.user.ok_or(Error::InvalidSession)?; let session = Session::from_request_parts(parts, state) .await - .map_err(|_| UiAuthError::Unavailable)?; + .map_err(|_| Error::MissingSessionLayer)?; let expected = session .get::(UI_CSRF_KEY) .await? - .ok_or(UiAuthError::Unauthorized)?; + .ok_or(Error::InvalidSession)?; let provided = parts .headers .get(AUTHORIZATION) .and_then(|value| value.to_str().ok()) .and_then(|value| value.strip_prefix("Bearer ")) - .ok_or(UiAuthError::Unauthorized)?; + .ok_or(Error::InvalidSession)?; if !bool::from(Sha256::digest(expected).ct_eq(&Sha256::digest(provided))) { - return Err(UiAuthError::Unauthorized); + return Err(Error::InvalidSession); } let identity = session_identity(&user).await?; parts.extensions.insert(identity.clone()); @@ -118,7 +118,7 @@ impl FromRequestParts for UiSession { } } -async fn session_identity(user: &UiUser) -> Result { +async fn session_identity(user: &UiUser) -> Result { crate::authentication::resolve( VerifiedIdentity { principal: Principal::new( @@ -142,7 +142,6 @@ async fn session_identity(user: &UiUser) -> Result Ok(Some(SecretValue::new("resolved-key"))), "EMPTY" => Ok(Some(SecretValue::new(""))), - "ERROR" => Err(litellm_secrets::Error::ExternalRead(Box::new( - std::io::Error::other("private-backend-detail"), - ))), + "ERROR" => Err(litellm_secrets::Error::ExternalRead( + litellm_error::Source::new(std::io::Error::other("private-backend-detail")), + )), _ => Ok(None), } }) @@ -132,7 +132,11 @@ impl Authenticator for Verifier { Box::pin(async move { match self.failure { Some(401) => return Err(Error::InvalidToken), - Some(_) => return Err(Error::Unavailable), + Some(_) => { + return Err(Error::Unavailable(litellm_error::Source::new( + std::io::Error::other("verifier offline"), + ))); + } None => (), } let litellm_gateway_auth::Credential::Token(credential) = credential else { diff --git a/litellm-rust/crates/gateway-auth/tests/keys.rs b/litellm-rust/crates/gateway-auth/tests/keys.rs index 51ef23a859f..4d16d3573bc 100644 --- a/litellm-rust/crates/gateway-auth/tests/keys.rs +++ b/litellm-rust/crates/gateway-auth/tests/keys.rs @@ -2,7 +2,7 @@ use std::time::{Duration, SystemTime}; use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ - KeyError, + Error, keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key}, }; use rstest::{fixture, rstest}; @@ -59,7 +59,7 @@ async fn checks_expiration( match valid { true => assert_eq!(result.unwrap(), hash), - false => assert!(matches!(result, Err(KeyError::Expired))), + false => assert!(matches!(result, Err(Error::Expired))), } } @@ -79,7 +79,7 @@ async fn rejects_unusable_keys( assert!(matches!( verify_key(&lookup, &token, || now).await, - Err(KeyError::Invalid) + Err(Error::InvalidToken) )); } @@ -93,7 +93,7 @@ async fn rejects_a_different_token(token: SecretValue, now: SystemTime) { assert!(matches!( verify_key(&lookup, &SecretValue::new("sk-other-credential"), || now).await, - Err(KeyError::Invalid) + Err(Error::InvalidToken) )); } @@ -102,9 +102,9 @@ struct Unavailable; impl KeyLookup for Unavailable { fn lookup<'a>(&'a self, _: &'a KeyHash) -> LookupFuture<'a> { Box::pin(async { - Err(KeyError::Lookup(Box::new(std::io::Error::other( - "backend detail", - )))) + Err(Error::Unavailable(litellm_error::Source::new( + std::io::Error::other("backend detail"), + ))) }) } } @@ -114,10 +114,8 @@ impl KeyLookup for Unavailable { async fn lookup_failure_preserves_source_without_exposing_it(token: SecretValue, now: SystemTime) { let error = verify_key(&Unavailable, &token, || now).await.unwrap_err(); - assert!(matches!(&error, KeyError::Lookup(_))); - assert_eq!( - std::error::Error::source(&error).unwrap().to_string(), - "backend detail" - ); + assert_eq!(error.status(), axum::http::StatusCode::SERVICE_UNAVAILABLE); + let cause = litellm_error::find::(&error).expect("lookup cause"); + assert_eq!(cause.to_string(), "backend detail"); assert!(!error.to_string().contains("backend detail")); } diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 312d44b9a4b..034c3a9826e 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -11,6 +11,7 @@ axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true bytes.workspace = true litellm-auth.workspace = true +litellm-error.workspace = true litellm-gateway-auth.workspace = true litellm-inference.workspace = true litellm-inference-transcription.workspace = true @@ -34,5 +35,6 @@ litellm-cache-memory.workspace = true futures-util.workspace = true tokio = { workspace = true, features = ["io-util"] } rstest.workspace = true +tokio-tungstenite.workspace = true tower = { version = "0.5.3", features = ["util"] } wiremock.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs index 8d571eccc5b..d4cb467f6fd 100644 --- a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs +++ b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs @@ -27,7 +27,7 @@ async fn handle( upload, }: InferenceBody, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, "/audio/transcriptions", &body)?; request::authorize_model(identity, deployment, &body).await?; let audio = match upload { Some(upload) => { diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index 426b8b91182..72f07d5d621 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -5,7 +5,7 @@ use litellm_gateway_auth::AuthenticatedRequest; use serde::Deserialize; use serde_json::{Map, Value}; -use crate::Error; +use crate::BodyError; #[derive(Default, Deserialize)] #[serde(default, deny_unknown_fields)] @@ -24,10 +24,9 @@ type Prepared = (Map, CacheOptions); pub(crate) fn prepare( identity: &AuthenticatedRequest, body: Map, -) -> Result { +) -> Result { let controls: Controls = match body.get("cache").filter(|value| !value.is_null()) { - Some(value) => serde_json::from_value(value.clone()) - .map_err(|error| Error::InvalidBody(error.to_string()))?, + Some(value) => serde_json::from_value(value.clone()).map_err(BodyError::CacheControls)?, None => Controls::default(), }; let caching: Option = body @@ -35,7 +34,7 @@ pub(crate) fn prepare( .filter(|value| !value.is_null()) .map(|value| serde_json::from_value(value.clone())) .transpose() - .map_err(|error| Error::InvalidBody(error.to_string()))?; + .map_err(BodyError::Caching)?; let caller = identity.caller(); let options = CacheOptions { policy: CachePolicy { @@ -62,11 +61,12 @@ pub(crate) fn prepare( )) } -fn duration(seconds: f64) -> Result { - Duration::try_from_secs_f64(seconds) - .ok() - .filter(|duration| !duration.is_zero()) - .ok_or_else(|| Error::InvalidBody("cache durations must be finite and positive".into())) +fn duration(seconds: f64) -> Result { + match Duration::try_from_secs_f64(seconds) { + Ok(duration) if !duration.is_zero() => Ok(duration), + Ok(_) => Err(BodyError::CacheDuration(None)), + Err(error) => Err(BodyError::CacheDuration(Some(error))), + } } #[derive(Clone, Default)] diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 742a373b286..fb8a6f4609d 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -3,13 +3,16 @@ use std::sync::Arc; use axum::{ Json, - extract::{Path, State}, + extract::State, response::{IntoResponse, Response}, }; use litellm_inference_chat::types::ChatCompletionsCall; use serde_json::{Map, Value}; -use crate::{Error, Gateway, JsonObject, request}; +use crate::{ + Error, Gateway, JsonObject, + request::{self, Path}, +}; pub(crate) async fn create( State(gateway): State>, @@ -40,7 +43,7 @@ async fn handle( identity: &AuthenticatedRequest, body: Map, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, "/chat/completions", &body)?; request::authorize_model(identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(identity, body)?; let route = gateway.chat_completions.clone(); diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index ec82a6ede12..190561079f6 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -3,7 +3,6 @@ use axum::{ Json, response::{IntoResponse, Response}, }; -use litellm_http::transport::Error as TransportError; use litellm_inference::Error as InferenceError; use litellm_llms::base_llm::ocr::error::Error as OcrError; use serde_json::{Map, Value, json}; @@ -12,12 +11,14 @@ use serde_json::{Map, Value, json}; pub enum Error { #[error(transparent)] Auth(#[from] litellm_gateway_auth::Error), - #[error("invalid request body: {0}")] - InvalidBody(String), + #[error("invalid request body")] + InvalidBody(#[from] BodyError), + #[error("invalid path parameter")] + InvalidPath(#[source] axum::extract::rejection::PathRejection), #[error( - "/v1/messages: Invalid model name passed in model={0}. Call `/v1/models` to view available models for your key." + "{route}: Invalid model name passed in model={model}. Call `/v1/models` to view available models for your key." )] - UnknownModel(String), + UnknownModel { route: &'static str, model: String }, #[error(transparent)] Route(#[from] InferenceError), #[error(transparent)] @@ -26,8 +27,30 @@ pub enum Error { Unsupported(String), #[error("request body exceeds the size limit")] BodyTooLarge, - #[error("{0}")] - Internal(String), + #[error("unexpected HTTP host operation")] + Internal, +} + +#[derive(Debug, thiserror::Error)] +pub enum BodyError { + #[error("could not read the body")] + Read(#[source] axum::extract::rejection::BytesRejection), + #[error("invalid JSON")] + Json(#[source] serde_json::Error), + #[error("expected a JSON object")] + NotAnObject, + #[error("invalid multipart body")] + Multipart(#[source] axum::extract::multipart::MultipartRejection), + #[error("invalid multipart field")] + MultipartField(#[source] axum::extract::multipart::MultipartError), + #[error("model is required")] + MissingModel, + #[error("invalid cache controls")] + CacheControls(#[source] serde_json::Error), + #[error("invalid caching flag")] + Caching(#[source] serde_json::Error), + #[error("cache durations must be finite and positive")] + CacheDuration(#[source] Option), } impl IntoResponse for Error { @@ -40,7 +63,7 @@ impl From> for Error { fn from(error: litellm_host_http::Error) -> Self { match error { litellm_host_http::Error::Call(error) => Self::Route(error), - litellm_host_http::Error::Protocol => Self::Internal(error.to_string()), + litellm_host_http::Error::Protocol => Self::Internal, } } } @@ -49,58 +72,50 @@ impl Error { pub fn status(&self) -> StatusCode { match self { Self::Auth(error) => error.status(), - Self::Unsupported(_) - | Self::Route(InferenceError::Unsupported(_)) - | Self::Ocr(OcrError::Unsupported(_)) => StatusCode::NOT_IMPLEMENTED, + Self::InvalidPath(rejection) => rejection.status(), + Self::Unsupported(_) => StatusCode::NOT_IMPLEMENTED, Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE, - Self::Ocr( - OcrError::Auth(litellm_auth::Error::MissingApiKey { .. }) - | OcrError::MissingAzureAiCredentials - | OcrError::MissingAzureDocumentIntelligenceCredentials - | OcrError::MissingReductoApiKey, - ) => StatusCode::UNAUTHORIZED, - Self::Ocr(error) => error - .http_status_code() - .and_then(|status| StatusCode::from_u16(status).ok()) - .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), - Self::InvalidBody(_) | Self::UnknownModel(_) => StatusCode::BAD_REQUEST, - Self::Route(InferenceError::Transport(TransportError::Http { status, .. })) => { - StatusCode::from_u16(*status).unwrap_or(StatusCode::BAD_GATEWAY) - } - Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => { - StatusCode::UNAUTHORIZED - } - Self::Route(error) if error.is_request() => StatusCode::BAD_REQUEST, - Self::Route(_) | Self::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR, + Self::InvalidBody(_) | Self::UnknownModel { .. } => StatusCode::BAD_REQUEST, + Self::Internal => StatusCode::INTERNAL_SERVER_ERROR, + Self::Ocr(error) => ocr_status(error), + Self::Route(error) => route_status(error), } } pub fn openai_response(self) -> Response { + (self.status(), Json(self.openai_body())).into_response() + } + + /// The OpenAI envelope as an `event: error` frame, for a stream that fails after its + /// headers went out. + pub fn responses_sse_frame(&self) -> String { + let event = json!({ + "type": "error", + "code": self.status().as_u16().to_string(), + "message": self.message(), + "param": null, + }); + format!("event: error\ndata: {event}\n\n") + } + + fn openai_body(&self) -> Value { let status = self.status(); - let message = match &self { - Self::UnknownModel(model) => format!("Invalid model name passed in model={model}"), - _ => self.to_string(), - }; - ( - status, - Json(json!({"error": { - "message": message, - "type": error_type(status), - "param": null, - "code": status.as_u16(), - }})), - ) - .into_response() + json!({"error": { + "message": self.message(), + "type": error_type(status), + "param": null, + "code": status.as_u16(), + }}) } /// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream /// body already in that shape passes through, any other has its message extracted. pub fn body(&self, request_id: Option<&str>) -> Value { let raw = match self { - Self::Route(InferenceError::Transport(TransportError::Http { body, .. })) => { - body.clone() + Self::Route(error) if let Some((_, body)) = error.upstream_response() => { + body.to_owned() } - other => other.to_string(), + other => other.message(), }; let parsed = serde_json::from_str::(&raw).ok(); let envelope = match parsed { @@ -117,6 +132,49 @@ impl Error { pub fn sse_frame(&self) -> String { format!("event: error\ndata: {}\n\n", self.body(None)) } + + fn message(&self) -> String { + match self { + Self::Auth(error) => error.to_string(), + Self::InvalidBody(_) + | Self::InvalidPath(_) + | Self::UnknownModel { .. } + | Self::Route(_) + | Self::Ocr(_) + | Self::Unsupported(_) + | Self::BodyTooLarge + | Self::Internal => litellm_error::Report::new(self).to_string(), + } + } +} + +fn route_status(error: &InferenceError) -> StatusCode { + if error.declined().is_some() { + return StatusCode::NOT_IMPLEMENTED; + } + if let Some((status, _)) = error.upstream_response() { + return StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY); + } + if error.is_missing_credential() { + return StatusCode::UNAUTHORIZED; + } + if error.is_request() { + return StatusCode::BAD_REQUEST; + } + StatusCode::INTERNAL_SERVER_ERROR +} + +fn ocr_status(error: &OcrError) -> StatusCode { + if let Some((status, ..)) = error.upstream_response() { + return StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY); + } + if error.is_missing_credential() { + return StatusCode::UNAUTHORIZED; + } + error + .http_status_code() + .and_then(|status| StatusCode::from_u16(status).ok()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR) } fn error_type(status: StatusCode) -> &'static str { @@ -133,13 +191,13 @@ fn error_type(status: StatusCode) -> &'static str { } fn envelope(status: StatusCode, message: &str) -> Map { - let Value::Object(envelope) = json!({ - "type": "error", - "error": {"type": error_type(status), "message": message}, - }) else { - unreachable!("a json object literal is an object") - }; - envelope + Map::from_iter([ + ("type".to_owned(), Value::from("error")), + ( + "error".to_owned(), + json!({"type": error_type(status), "message": message}), + ), + ]) } fn is_anthropic_error(object: &Map) -> bool { @@ -177,14 +235,22 @@ mod tests { use super::*; fn upstream(status: u16, body: &str) -> Error { - Error::Route(InferenceError::Transport(TransportError::Http { - status, - body: body.into(), - })) + Error::Route(litellm_http::Error::upstream(status, body.into()).into()) + } + + fn invalid_json() -> Error { + BodyError::Json(serde_json::from_str::("{").unwrap_err()).into() + } + + fn unknown_model() -> Error { + Error::UnknownModel { + route: "/v1/messages", + model: "nope".into(), + } } #[rstest] - #[case::anthropic_body_passes_through( + #[case::anthropic_body( upstream(529, r#"{"type":"error","error":{"type":"overloaded_error","message":"busy","extra":1}}"#), Some("req_1"), json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy", "extra": 1}, "request_id": "req_1"}), @@ -194,30 +260,7 @@ mod tests { Some("caller"), json!({"type": "error", "error": {"type": "x", "message": "m"}, "request_id": "upstream"}), )] - #[case::bedrock_detail( - upstream(403, r#"{"detail":{"message":"denied"}}"#), - None, - json!({"type": "error", "error": {"type": "permission_error", "message": "denied"}}), - )] - #[case::aws_message( - upstream(429, r#"{"Message":"slow down"}"#), - None, - json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}}), - )] - #[case::plain_text_with_unmapped_status( - upstream(502, "bad gateway"), - None, - json!({"type": "error", "error": {"type": "api_error", "message": "bad gateway"}}), - )] - #[case::unknown_model( - Error::UnknownModel("nope".into()), - None, - json!({"type": "error", "error": { - "type": "invalid_request_error", - "message": "/v1/messages: Invalid model name passed in model=nope. Call `/v1/models` to view available models for your key.", - }}), - )] - fn body_follows_the_anthropic_exception_mapping( + fn an_anthropic_provider_body_passes_through( #[case] error: Error, #[case] request_id: Option<&str>, #[case] expected: Value, @@ -225,9 +268,74 @@ mod tests { assert_eq!(error.body(request_id), expected); } + #[rstest] + #[case::ascii("busy ".repeat(80))] + #[case::unicode("é".repeat(300))] + fn long_provider_errors_preserve_the_messages_envelope(#[case] message: String) { + let body = json!({ + "type": "error", + "error": {"type": "overloaded_error", "message": message}, + "request_id": "provider-request", + }); + let error = upstream(529, &body.to_string()); + let report = litellm_error::Report::new(&error).to_string(); + let diagnostic_body: String = body.to_string().chars().take(256).collect(); + assert_eq!(report.matches(&diagnostic_body).count(), 1, "{report}"); + assert!(report.ends_with("... (truncated)"), "{report}"); + assert!(!report.contains(&message), "{report}"); + assert!(!format!("{error:?}").contains(&message)); + assert_eq!(error.status(), StatusCode::from_u16(529).unwrap()); + assert_eq!(error.body(Some("gateway-request")), body); + } + + #[rstest] + #[case::bedrock_detail( + upstream(403, r#"{"detail":{"message":"denied"}}"#), + "permission_error", + "denied" + )] + #[case::aws_message( + upstream(429, r#"{"Message":"slow down"}"#), + "rate_limit_error", + "slow down" + )] + #[case::plain_text_with_unmapped_status( + upstream(502, "bad gateway"), + "api_error", + "bad gateway" + )] + #[case::unknown_model( + unknown_model(), + "invalid_request_error", + "/v1/messages: Invalid model name passed in model=nope. Call `/v1/models` to view available models for your key." + )] + fn other_failures_get_the_anthropic_envelope_with_their_message( + #[case] error: Error, + #[case] error_type: &str, + #[case] message: &str, + ) { + let body = error.body(Some("req_1")); + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], error_type); + assert_eq!(body["error"]["message"], message); + assert_eq!(body["request_id"], "req_1"); + } + #[rstest] #[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)] - #[case::rejected_request(Error::Route(InferenceError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] + #[case::declined( + Error::Route(InferenceError::from(litellm_http::Error::computed_header( + "authorization" + ))), + StatusCode::NOT_IMPLEMENTED + )] + #[case::rejected_request( + Error::Route(InferenceError::invalid_request(litellm_llms::ErrorDetail::Expected { + subject: "top_k", + expected: "an integer", + })), + StatusCode::BAD_REQUEST + )] #[case::missing_key( Error::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", @@ -236,10 +344,83 @@ mod tests { StatusCode::UNAUTHORIZED, )] #[case::lost_connection( - Error::Route(InferenceError::Transport(TransportError::Network("reset".into()))), + Error::Route( + litellm_http::Error::from(tokio_tungstenite::tungstenite::Error::ConnectionClosed) + .into() + ), StatusCode::INTERNAL_SERVER_ERROR, )] + #[case::invalid_body(invalid_json(), StatusCode::BAD_REQUEST)] + #[case::missing_model(BodyError::MissingModel.into(), StatusCode::BAD_REQUEST)] + #[case::unknown_model(unknown_model(), StatusCode::BAD_REQUEST)] + #[case::oversized_body(Error::BodyTooLarge, StatusCode::PAYLOAD_TOO_LARGE)] + #[case::unserved_path(Error::Unsupported("/v1/completions".into()), StatusCode::NOT_IMPLEMENTED)] + #[case::host_protocol( + litellm_host_http::Error::::Protocol.into(), + StatusCode::INTERNAL_SERVER_ERROR + )] + #[case::ocr_provider_answer( + Error::Ocr(OcrError::Provider { status: 429, body: "slow".into(), headers: Default::default() }), + StatusCode::TOO_MANY_REQUESTS + )] + #[case::ocr_missing_credential( + Error::Ocr(OcrError::MissingCredential { credential: "Azure AI credentials", guidance: "set a key" }), + StatusCode::UNAUTHORIZED + )] + #[case::ocr_request_fault(Error::Ocr(OcrError::MissingDocumentUrl), StatusCode::BAD_REQUEST)] + #[case::ocr_response_fault( + Error::Ocr(OcrError::EmptyContent), + StatusCode::INTERNAL_SERVER_ERROR + )] fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { assert_eq!(error.status(), status); } + + #[test] + fn a_body_fault_names_the_body_and_keeps_its_cause_once() { + let cause = serde_json::from_str::("{").unwrap_err().to_string(); + let body = invalid_json().openai_body(); + let message = body["error"]["message"].as_str().unwrap(); + assert!( + message.starts_with("invalid request body: invalid JSON: "), + "{message}" + ); + assert_eq!(message.matches(&cause).count(), 1, "{message}"); + assert_eq!(body["error"]["type"], "invalid_request_error"); + assert_eq!(body["error"]["code"], 400); + } + + fn auth_outage() -> Error { + Error::Auth(litellm_gateway_auth::Error::Unavailable( + litellm_error::Source::new(std::io::Error::other("private-backend-detail")), + )) + } + + #[rstest] + #[case::openai(auth_outage().openai_body()["error"]["message"].clone())] + #[case::anthropic(auth_outage().body(None)["error"]["message"].clone())] + #[case::responses_sse(Value::from(auth_outage().responses_sse_frame()))] + #[case::anthropic_sse(Value::from(auth_outage().sse_frame()))] + fn an_auth_failure_hides_its_backend_cause(#[case] rendered: Value) { + let top = auth_outage().to_string(); + let rendered = rendered.to_string(); + assert!(rendered.contains(&top), "{rendered}"); + assert!(!rendered.contains("private-backend-detail"), "{rendered}"); + } + + #[rstest] + #[case::provider_answer(upstream(429, "slow down"))] + #[case::caller_fault(invalid_json())] + #[case::host_bug(litellm_host_http::Error::::Protocol.into())] + fn the_responses_sse_frame_is_the_responses_error_event(#[case] error: Error) { + let frame = error.responses_sse_frame(); + let data = frame + .strip_prefix("event: error\ndata: ") + .and_then(|data| data.strip_suffix("\n\n")) + .unwrap_or_else(|| panic!("an event: error frame, got {frame:?}")); + let event = serde_json::from_str::(data).unwrap(); + assert_eq!(event["type"], "error"); + assert_eq!(event["code"], error.status().as_u16().to_string()); + assert_eq!(event["message"], error.openai_body()["error"]["message"]); + } } diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 4663151f1de..43966f2bee7 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -25,7 +25,7 @@ use litellm_inference_transcription::AudioTranscriptionRoute; use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; use litellm_secrets::source::SecretSource; -pub use error::Error; +pub use error::{BodyError, Error}; pub use litellm_router::{Deployment, Router as ModelRouter}; pub use request::{JsonObject, RequestId}; diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index d612a042e03..58ecd40204d 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -46,7 +46,7 @@ async fn handle( headers: &HeaderMap, body: Map, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, "/v1/messages", &body)?; request::authorize_model(identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(identity, body)?; let route = gateway.messages.clone(); diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs index 03d064b6bae..f200fd7b141 100644 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -35,7 +35,7 @@ async fn handle( .get("x-req-format") .and_then(|value| value.to_str().ok()) .map(str::to_owned); - let deployment = request::resolve_deployment(gateway, &body)?; + let deployment = request::resolve_deployment(gateway, "/ocr", &body)?; request::authorize_model(identity, deployment, &body).await?; let document = match upload { Some(upload) => OcrDocumentInput::Bytes { diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index b398da2013b..d41794242ff 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -6,7 +6,7 @@ use axum::{ }; use serde_json::{Map, Value}; -use crate::{Deployment, Error, Gateway}; +use crate::{BodyError, Deployment, Error, Gateway}; pub(crate) const MAX_FILE_BYTES: usize = 50 * 1024 * 1024; pub(crate) const MAX_BODY_BYTES: usize = MAX_FILE_BYTES + 1024 * 1024; @@ -21,6 +21,19 @@ pub struct JsonObject(pub Map); pub struct RequestId(pub Option); +pub(crate) struct Path(pub T); + +impl FromRequestParts for Path { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + axum::extract::Path::::from_request_parts(parts, state) + .await + .map(|axum::extract::Path(value)| Self(value)) + .map_err(Error::InvalidPath) + } +} + impl FromRequestParts for RequestId { type Rejection = std::convert::Infallible; @@ -43,10 +56,10 @@ impl FromRequest for JsonObject { if error.status() == axum::http::StatusCode::PAYLOAD_TOO_LARGE { Error::BodyTooLarge } else { - Error::InvalidBody(error.to_string()) + BodyError::Read(error).into() } })?; - object(&body).map(Self) + Ok(Self(object(&body)?)) } } @@ -77,16 +90,16 @@ impl FromRequest for InferenceBody { } let multipart = Multipart::from_request(request, state) .await - .map_err(|error| Error::InvalidBody(error.to_string()))?; + .map_err(BodyError::Multipart)?; parse_multipart(multipart).await } } -fn object(body: &[u8]) -> Result, Error> { +fn object(body: &[u8]) -> Result, BodyError> { match serde_json::from_slice(body) { Ok(Value::Object(body)) => Ok(body), - Ok(_) => Err(Error::InvalidBody("expected a JSON object".into())), - Err(error) => Err(Error::InvalidBody(error.to_string())), + Ok(_) => Err(BodyError::NotAnObject), + Err(error) => Err(BodyError::Json(error)), } } @@ -98,7 +111,7 @@ pub(crate) async fn authorize_model( let name = body .get("model") .and_then(Value::as_str) - .ok_or_else(|| Error::InvalidBody("model is required".into()))?; + .ok_or(BodyError::MissingModel)?; let access = litellm_gateway_auth::AccessRequest::Model { name: name.into(), deployment: deployment.model.clone(), @@ -112,16 +125,20 @@ pub(crate) async fn authorize_model( pub(crate) fn resolve_deployment<'a>( gateway: &'a Gateway, + route: &'static str, body: &Map, ) -> Result<&'a Deployment, Error> { let model = body .get("model") .and_then(Value::as_str) - .ok_or_else(|| Error::InvalidBody("model is required".into()))?; + .ok_or(BodyError::MissingModel)?; gateway .models .get(model) - .ok_or_else(|| Error::UnknownModel(model.to_owned())) + .ok_or_else(|| Error::UnknownModel { + route, + model: model.to_owned(), + }) } async fn parse_multipart(mut multipart: Multipart) -> Result { @@ -159,7 +176,7 @@ fn multipart_error(error: axum::extract::multipart::MultipartError) -> Error { if error.status() == axum::http::StatusCode::PAYLOAD_TOO_LARGE { return Error::BodyTooLarge; } - Error::InvalidBody(error.to_string()) + BodyError::MultipartField(error).into() } pub(crate) async fn unsupported(OriginalUri(uri): OriginalUri) -> impl IntoResponse { diff --git a/litellm-rust/crates/gateway-inference/src/responses.rs b/litellm-rust/crates/gateway-inference/src/responses.rs index 77bebf610e1..df272460d21 100644 --- a/litellm-rust/crates/gateway-inference/src/responses.rs +++ b/litellm-rust/crates/gateway-inference/src/responses.rs @@ -4,7 +4,6 @@ use axum::{Json, body::Bytes, extract::State, response::Response}; use litellm_gateway_auth::AuthenticatedRequest; use litellm_host_http::Sse; use litellm_inference_responses::{route::Responses, types::ResponsesCall}; -use serde_json::json; use crate::{Error, Gateway, JsonObject, request}; @@ -13,7 +12,7 @@ pub(crate) async fn create( identity: AuthenticatedRequest, JsonObject(body): JsonObject, ) -> Result { - let deployment = request::resolve_deployment(&gateway, &body)?; + let deployment = request::resolve_deployment(&gateway, "/responses", &body)?; request::authorize_model(&identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(&identity, body)?; let route = gateway.responses.clone(); @@ -40,11 +39,7 @@ pub(crate) async fn create( }; let machine = route.machine(call, cache_options.policy); let stream = Sse::::new(Json, |error| { - let error = Error::from(error); - Bytes::from(format!( - "event: error\ndata: {}\n\n", - json!({"type": "error", "code": error.status().as_u16().to_string(), "message": error.to_string(), "param": null}) - )) + Bytes::from(Error::from(error).responses_sse_frame()) }); let headers = crate::caching::CacheHeaders::default(); let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?; diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index bf11e67e0fc..7d6d1f160d1 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -6,7 +6,6 @@ use axum::{ }; use rstest::rstest; use serde_json::json; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; use wiremock::{ Mock, MockServer, ResponseTemplate, @@ -90,34 +89,11 @@ async fn invalid_messages_stays_an_anthropic_error() { assert_eq!(body["error"]["type"], "invalid_request_error"); } -/// Answers with the SSE head and one event, then drops the connection short of the -/// announced body length. -async fn truncating_upstream() -> String { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = vec![0; 4096]; - let _ = socket.read(&mut request).await; - socket - .write_all( - format!( - "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{FIRST_EVENT}", - FIRST_EVENT.len() * 2 - ) - .as_bytes(), - ) - .await - .unwrap(); - }); - base -} - const FIRST_EVENT: &str = "event: message_start\ndata: {}\n\n"; #[tokio::test] async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame() { - let base = truncating_upstream().await; + let base = support::truncating_upstream(FIRST_EVENT).await; let request = Request::post("/v1/messages") .header("content-type", "application/json") .body(Body::from( diff --git a/litellm-rust/crates/gateway-inference/tests/responses.rs b/litellm-rust/crates/gateway-inference/tests/responses.rs index ded7d73c5f3..ed7c8e6fe76 100644 --- a/litellm-rust/crates/gateway-inference/tests/responses.rs +++ b/litellm-rust/crates/gateway-inference/tests/responses.rs @@ -116,3 +116,32 @@ async fn responses_authorize_the_public_model_before_provider_execution() { assert_eq!(response.status(), 403); assert!(upstream.received_requests().await.unwrap().is_empty()); } + +const RESPONSE_CREATED: &str = "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"response-1\",\"model\":\"test-model\",\"output\":[]}}\n\n"; + +#[rstest] +#[tokio::test] +async fn a_stream_that_fails_after_opening_ends_with_the_responses_error_event() { + let base = support::truncating_upstream(RESPONSE_CREATED).await; + let response = support::post( + support::app("openai/test-model", &base), + "/v1/responses", + json!({"model": "public/model", "input": "hello", "stream": true}), + ) + .await; + assert_eq!(response.status(), 200); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + let text = std::str::from_utf8(&body).unwrap(); + let frame = text + .strip_prefix(RESPONSE_CREATED) + .and_then(|rest| rest.strip_prefix("event: error\ndata: ")) + .unwrap_or_else(|| panic!("the delivered event then one error frame, got {text:?}")); + let error: serde_json::Value = serde_json::from_str(frame.trim_end()).unwrap(); + assert_eq!(error["type"], "error"); + assert_eq!(error["code"], "500"); + assert!( + error["message"] + .as_str() + .is_some_and(|message| !message.is_empty()) + ); +} diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs index b0699ff10b5..24b8c5cc694 100644 --- a/litellm-rust/crates/gateway-inference/tests/routes.rs +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -125,7 +125,10 @@ async fn chat_errors_come_from_core( .await; assert_eq!(response.status(), status); let body = support::json(response).await; - assert_eq!(body["error"]["message"], error.to_string()); + assert_eq!( + body["error"]["message"], + litellm_error::Report::new(&error).to_string() + ); assert_eq!(body["error"]["code"], status); assert!(upstream.received_requests().await.unwrap().is_empty()); } @@ -193,17 +196,20 @@ async fn transcription_aliases_reach_core_validation(#[case] path: &str) { } #[rstest] -#[case::no_extension("test")] -#[case::unsupported_extension("test.invalid")] +#[case::no_extension("test", json!(null))] +#[case::unsupported_extension("test.invalid", json!("invalid"))] #[tokio::test] -async fn upload_audio_format_validation_matches_core(#[case] filename: &str) { +async fn upload_audio_format_validation_matches_core( + #[case] filename: &str, + #[case] format: Value, +) { let upstream = MockServer::start().await; let app = support::app("bedrock/test-model", &upstream.uri()); let path = "/v1/audio/transcriptions"; let expected = support::post( app.clone(), path, - json!({"model": "public/model", "audio": {"data": "YWJj", "format": null}}), + json!({"model": "public/model", "audio": {"data": "YWJj", "format": format}}), ) .await; let payload = format!( diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 0e37bb0aff9..737dc0c830c 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -15,6 +15,7 @@ use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsRes use litellm_inference::resources::CoreResources; use litellm_secrets::{SecretValue, source::SecretSource}; use serde_json::Value; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; pub struct NoSecrets; @@ -166,3 +167,26 @@ impl litellm_gateway_auth::IdentityResolver for TestPermissions { }) } } + +/// Answers with the SSE head and one event, then drops the connection short of the +/// announced body length. +pub async fn truncating_upstream(first_event: &'static str) -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{first_event}", + first_event.len() * 2 + ) + .as_bytes(), + ) + .await + .unwrap(); + }); + base +} diff --git a/litellm-rust/crates/gateway-management/Cargo.toml b/litellm-rust/crates/gateway-management/Cargo.toml index 53ff2ea59aa..502760f8fbe 100644 --- a/litellm-rust/crates/gateway-management/Cargo.toml +++ b/litellm-rust/crates/gateway-management/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] base64.workspace = true litellm-auth-types.workspace = true +litellm-error.workspace = true litellm-gateway-auth.workspace = true rand.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/gateway-management/src/error.rs b/litellm-rust/crates/gateway-management/src/error.rs index 5da7977c7cb..7194395d668 100644 --- a/litellm-rust/crates/gateway-management/src/error.rs +++ b/litellm-rust/crates/gateway-management/src/error.rs @@ -5,7 +5,7 @@ pub enum Error { #[error("key already exists")] AlreadyExists, #[error("key storage unavailable")] - Storage(#[source] Box), + Storage(#[source] litellm_error::Source), #[error("key generation unavailable")] Entropy(#[source] rand::Error), } diff --git a/litellm-rust/crates/gateway-management/tests/keys.rs b/litellm-rust/crates/gateway-management/tests/keys.rs index dbfedb7ac4d..bf0221cf1be 100644 --- a/litellm-rust/crates/gateway-management/tests/keys.rs +++ b/litellm-rust/crates/gateway-management/tests/keys.rs @@ -9,7 +9,7 @@ use std::{ use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ - KeyError, hash_token, + hash_token, keys::{KeyHash, KeyLookup, KeyStatus, verify_key}, }; use litellm_gateway_management::{ @@ -73,7 +73,7 @@ async fn generated_key_can_be_retrieved_verified_and_revoked( ); assert!(matches!( verify_key(lookup.as_ref(), &generated.token, || now).await, - Err(KeyError::Invalid) + Err(litellm_gateway_auth::Error::InvalidToken) )); assert_eq!(keys.revoke(hash).await.unwrap(), Revocation::Revoked); } @@ -137,9 +137,7 @@ async fn storage_failures_are_returned_for_all_operations(now: SystemTime) { assert!(matches!(keys.get(&hash).await, Err(Error::Storage(_)))); let error = keys.revoke(&hash).await.unwrap_err(); assert!(matches!(&error, Error::Storage(_))); - assert_eq!( - std::error::Error::source(&error).unwrap().to_string(), - "backend detail" - ); + let cause = litellm_error::find::(&error).expect("storage cause"); + assert_eq!(cause.to_string(), "backend detail"); assert!(!error.to_string().contains("backend detail")); } diff --git a/litellm-rust/crates/gateway-management/tests/keys/support.rs b/litellm-rust/crates/gateway-management/tests/keys/support.rs index d0efec4c780..71511c7eaba 100644 --- a/litellm-rust/crates/gateway-management/tests/keys/support.rs +++ b/litellm-rust/crates/gateway-management/tests/keys/support.rs @@ -1,10 +1,7 @@ use std::{collections::BTreeMap, sync::Mutex}; use litellm_gateway_auth::keys::KeyStatus; -use litellm_gateway_auth::{ - KeyError, - keys::{KeyHash, KeyLookup, LookupFuture}, -}; +use litellm_gateway_auth::keys::{KeyHash, KeyLookup, LookupFuture}; use litellm_gateway_management::{ Error, keys::{KeyRecord, KeyStore, NewKey, Revocation, StoreFuture}, @@ -56,7 +53,9 @@ impl KeyLookup for MemoryStore { self.get(hash) .await .map(|record| record.map(|record| record.status)) - .map_err(|error| KeyError::Lookup(Box::new(error))) + .map_err(|error| { + litellm_gateway_auth::Error::Unavailable(litellm_error::Source::new(error)) + }) }) } } @@ -64,9 +63,9 @@ impl KeyLookup for MemoryStore { pub struct Unavailable; fn unavailable() -> Result { - Err(Error::Storage(Box::new(std::io::Error::other( - "backend detail", - )))) + Err(Error::Storage(litellm_error::Source::new( + std::io::Error::other("backend detail"), + ))) } impl KeyStore for Unavailable { diff --git a/litellm-rust/crates/gateway-mcp/AGENTS.md b/litellm-rust/crates/gateway-mcp/AGENTS.md index aa608f1e720..c4b00e3acbf 100644 --- a/litellm-rust/crates/gateway-mcp/AGENTS.md +++ b/litellm-rust/crates/gateway-mcp/AGENTS.md @@ -6,6 +6,8 @@ The HTTP router serves `/mcp`, `/{server}/mcp`, `/mcp/{server}`, `/mcp/{server}/ `NativeGateway` aggregates upstream tools, prompts, resources, and resource templates. Tool and prompt calls resolve registered server prefixes, including prefixes containing hyphens. Resource names receive prefixes while their URIs remain unchanged, and reads require exactly one selected server. Tool allowlists apply to discovery and execution. Upstream pagination is drained with cycle detection. Aggregate tool listings retain healthy servers and expose sanitized outcomes under `_meta["litellm.ai/server_outcomes"]`, while scoped REST listings preserve upstream failures +Every HTTP failure, including extractor rejections and the legacy SSE and session transport checks, is an `Error` with its status from `status()` and the `{"detail"}` envelope. The detail is the `litellm_error::Report` of the error, except for upstream and authorization failures, whose causes can carry upstream URLs and backend detail. Startup failures are `ConnectError`, whose `Debug` redacts the SDK's initialization cause for the same reason + ## Integration Construct `Server` entries from connected SDK peers and pass a `Registry` to `NativeGateway::new`. Keep their `RunningService` handles alive in the host and cancel them during shutdown. SDK peers can use stdio, Streamable HTTP, or another SDK transport. Connection establishment, outbound credentials, HTTP client pools, and reconnection remain owned by the host diff --git a/litellm-rust/crates/gateway-mcp/Cargo.toml b/litellm-rust/crates/gateway-mcp/Cargo.toml index 07706dcb6a4..bed48bad58d 100644 --- a/litellm-rust/crates/gateway-mcp/Cargo.toml +++ b/litellm-rust/crates/gateway-mcp/Cargo.toml @@ -12,6 +12,7 @@ futures-util.workspace = true litellm-config.workspace = true litellm-secrets.workspace = true litellm-auth-types.workspace = true +litellm-error.workspace = true base64.workspace = true http.workspace = true uuid.workspace = true @@ -19,12 +20,13 @@ url.workspace = true moka.workspace = true sha2.workspace = true tower = { version = "0.5", features = ["util"] } -rmcp = { version = "=3.4.1", default-features = false, features = ["server", "client", "elicitation", "transport-streamable-http-server", "transport-streamable-http-client", "transport-child-process"] } +rmcp = { version = "=3.4.1", default-features = false, features = ["server", "client", "elicitation", "transport-streamable-http-server", "transport-streamable-http-client-reqwest", "transport-child-process"] } serde.workspace = true serde_json.workspace = true thiserror.workspace = true tokio.workspace = true tokio-util = "0.7" +veil.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/gateway-mcp/src/configured.rs b/litellm-rust/crates/gateway-mcp/src/configured.rs index bb4f5594f1f..cc7a52fc029 100644 --- a/litellm-rust/crates/gateway-mcp/src/configured.rs +++ b/litellm-rust/crates/gateway-mcp/src/configured.rs @@ -22,25 +22,25 @@ use rmcp::{ use sha2::{Digest, Sha256}; use tokio_util::sync::CancellationToken; -use crate::{ConnectError, Limits, NativeGateway, Registry, Server, ServerInfo}; +use crate::{ConfigFault, ConnectError, Limits, NativeGateway, Registry, Server, ServerInfo}; pub struct ConfiguredGateway { pub operations: Arc, services: Vec>, } -fn invalid(server: &str, message: impl Into) -> ConnectError { +fn invalid(server: &str, fault: ConfigFault) -> ConnectError { ConnectError::Configuration { server: server.into(), - message: message.into(), + fault, } } fn timeout(name: &str, config: &McpServer) -> Result { let duration = Duration::try_from_secs_f64(config.timeout.unwrap_or(30.0)) - .map_err(|_| invalid(name, "timeout must be finite and positive"))?; + .map_err(|source| invalid(name, ConfigFault::InvalidTimeout(source)))?; if duration.is_zero() { - return Err(invalid(name, "timeout must be positive")); + return Err(invalid(name, ConfigFault::ZeroTimeout)); } Ok(duration) } @@ -75,27 +75,24 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> { .as_ref() .is_some_and(|alias| alias.trim().is_empty()) { - return Err(invalid(name, "server identifiers must not be blank")); + return Err(invalid(name, ConfigFault::BlankIdentifier)); } if let Some(field) = config.unsupported.keys().next() { return Err(invalid( name, - format!("setting '{field}' is not supported by the Rust MCP host"), + ConfigFault::UnsupportedSetting(field.clone()), )); } if config .max_concurrent_requests .is_some_and(|limit| limit == 0 || limit > tokio::sync::Semaphore::MAX_PERMITS) { - return Err(invalid(name, "max_concurrent_requests is out of range")); + return Err(invalid(name, ConfigFault::ConcurrencyOutOfRange)); } if config.mcp_info.contains_key("protocol_version") || config.mcp_info.contains_key("mcp_server_cost_info") { - return Err(invalid( - name, - "mcp_info protocol and cost settings are not supported", - )); + return Err(invalid(name, ConfigFault::UnsupportedMcpInfo)); } if !matches!( config.auth_type, @@ -108,17 +105,11 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> { | McpAuth::Token ) ) { - return Err(invalid( - name, - "configured upstream auth mode is not supported by the Rust MCP host", - )); + return Err(invalid(name, ConfigFault::UnsupportedAuth)); } match config.transport { McpTransport::Sse => { - return Err(invalid( - name, - "legacy SSE upstreams are not supported by the pinned SDK; use transport: http", - )); + return Err(invalid(name, ConfigFault::SseTransport)); } McpTransport::Http => { if config.url.is_none() @@ -126,10 +117,7 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> { || !config.args.is_empty() || !config.env.is_empty() { - return Err(invalid( - name, - "HTTP transport requires url and does not accept command, args or env", - )); + return Err(invalid(name, ConfigFault::HttpShape)); } } McpTransport::Stdio => { @@ -139,53 +127,49 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> { .is_none_or(|command| command.trim().is_empty()) || config.url.is_some() { - return Err(invalid( - name, - "stdio transport requires command and does not accept url", - )); + return Err(invalid(name, ConfigFault::StdioShape)); } if config.authentication_token.is_some() || !config.static_headers.is_empty() || config.upstream_token_header.is_some() || !matches!(config.auth_type, None | Some(McpAuth::None)) { - return Err(invalid( - name, - "stdio credentials must be supplied through env", - )); + return Err(invalid(name, ConfigFault::StdioCredentials)); } } } if matches!(config.auth_type, None | Some(McpAuth::None)) { if config.authentication_token.is_some() || config.upstream_token_header.is_some() { - return Err(invalid( - name, - "authentication_token and upstream_token_header require an auth_type", - )); + return Err(invalid(name, ConfigFault::TokenWithoutAuth)); } } else if config.authentication_token.is_none() { - return Err(invalid( - name, - "configured auth_type requires authentication_token", - )); + return Err(invalid(name, ConfigFault::AuthWithoutToken)); } Ok(()) } -async fn resolve(value: &SecretValue, secrets: &dyn SecretSource) -> Result { +async fn resolve( + name: &str, + value: &SecretValue, + secrets: &dyn SecretSource, +) -> Result { match value.expose().strip_prefix("os.environ/") { Some(key) => secrets .get_secret_str(key) .await? .map(|secret| secret.expose().to_owned()) - .ok_or_else(|| invalid("configuration", "referenced secret is unavailable")), + .ok_or_else(|| invalid(name, ConfigFault::MissingSecret)), None => Ok(value.expose().to_owned()), } } -fn header(name: &str, value: &str) -> Result<(HeaderName, HeaderValue), ConnectError> { +fn header( + server: &str, + name: &str, + value: &str, +) -> Result<(HeaderName, HeaderValue), ConnectError> { let key = HeaderName::try_from(name) - .map_err(|_| invalid("configuration", "invalid upstream header name"))?; + .map_err(|source| invalid(server, ConfigFault::HeaderName(source)))?; if matches!( key.as_str(), "host" @@ -198,13 +182,10 @@ fn header(name: &str, value: &str) -> Result<(HeaderName, HeaderValue), ConnectE | "mcp-protocol-version" | "last-event-id" ) { - return Err(invalid( - "configuration", - "upstream header conflicts with HTTP framing", - )); + return Err(invalid(server, ConfigFault::FramingHeader)); } let mut value = HeaderValue::try_from(value) - .map_err(|_| invalid("configuration", "invalid upstream header value"))?; + .map_err(|source| invalid(server, ConfigFault::HeaderValue(source)))?; value.set_sensitive(true); Ok((key, value)) } @@ -231,21 +212,19 @@ fn basic(value: &str) -> String { } async fn headers( + name: &str, config: &McpServer, secrets: &dyn SecretSource, ) -> Result, ConnectError> { - let entries = try_join_all( - config - .static_headers - .iter() - .map(|(key, value)| async move { header(key, &resolve(value, secrets).await?) }), - ) + let entries = try_join_all(config.static_headers.iter().map(|(key, value)| async move { + header(name, key, &resolve(name, value, secrets).await?) + })) .await?; let auth = match (&config.authentication_token, config.auth_type) { (Some(token), Some(kind)) => { - let token = resolve(token, secrets).await?; + let token = resolve(name, token, secrets).await?; if token.trim().is_empty() { - return Err(invalid("configuration", "upstream credential is empty")); + return Err(invalid(name, ConfigFault::EmptyCredential)); } let (slot, value) = match kind { McpAuth::ApiKey => ("x-api-key", token), @@ -253,9 +232,10 @@ async fn headers( McpAuth::Token => ("authorization", scheme(&token, "token")), McpAuth::Basic => ("authorization", basic(&token)), McpAuth::Authorization => ("authorization", token), - _ => return Err(invalid("configuration", "unsupported credential mode")), + _ => return Err(invalid(name, ConfigFault::UnsupportedAuth)), }; Some(header( + name, config.upstream_token_header.as_deref().unwrap_or(slot), &value, )?) @@ -349,36 +329,34 @@ async fn connect_service( let configured = config .url .as_ref() - .ok_or_else(|| invalid(name, "HTTP transport requires url"))?; - let url = resolve(configured, secrets).await?; - let parsed = url::Url::parse(&url).map_err(|_| invalid(name, "invalid HTTP URL"))?; + .ok_or_else(|| invalid(name, ConfigFault::HttpShape))?; + let url = resolve(name, configured, secrets).await?; + let parsed = url::Url::parse(&url) + .map_err(|source| invalid(name, ConfigFault::InvalidUrl(source)))?; if !matches!(parsed.scheme(), "http" | "https") || parsed.host_str().is_none() || !parsed.username().is_empty() || parsed.password().is_some() || parsed.fragment().is_some() { - return Err(invalid( - name, - "HTTP URL must use http or https without user credentials or fragments", - )); + return Err(invalid(name, ConfigFault::UrlShape)); } let transport = StreamableHttpClientTransport::with_client( client, StreamableHttpClientTransportConfig::with_uri(url) - .custom_headers(headers(config, secrets).await?), + .custom_headers(headers(name, config, secrets).await?), ); ().serve_with_ct(transport, shutdown).await } McpTransport::Stdio => { let environment = try_join_all(config.env.iter().map(|(key, value)| async move { - Ok::<_, ConnectError>((key, resolve(value, secrets).await?)) + Ok::<_, ConnectError>((key, resolve(name, value, secrets).await?)) })) .await?; let program = config .command .as_deref() - .ok_or_else(|| invalid(name, "stdio requires command"))?; + .ok_or_else(|| invalid(name, ConfigFault::StdioShape))?; let mut command = tokio::process::Command::new(program); command .args(&config.args) @@ -387,7 +365,7 @@ async fn connect_service( ().serve_with_ct(TokioChildProcess::new(command)?, shutdown) .await } - McpTransport::Sse => return Err(invalid(name, "unsupported transport")), + McpTransport::Sse => return Err(invalid(name, ConfigFault::SseTransport)), }; initialized.map_err(|source| ConnectError::Initialize { server: name.into(), @@ -452,7 +430,7 @@ mod tests { "mcp_servers:\n test:\n auth_type: {mode}\n authentication_token: '{token}'\n" )) .unwrap(); - let built = headers(&config.mcp_servers["test"], &Secrets) + let built = headers("test", &config.mcp_servers["test"], &Secrets) .await .unwrap(); let key = HeaderName::try_from(slot).unwrap(); @@ -464,7 +442,7 @@ mod tests { #[tokio::test] async fn explicit_credential_slot_overrides_static_header() { let config = litellm_config::Config::from_yaml("mcp_servers: {test: {auth_type: api_key, authentication_token: os.environ/KEY, upstream_token_header: x-custom, static_headers: {X-Custom: old, x-extra: extra}}}").unwrap(); - let built = headers(&config.mcp_servers["test"], &Secrets) + let built = headers("test", &config.mcp_servers["test"], &Secrets) .await .unwrap(); assert_eq!( diff --git a/litellm-rust/crates/gateway-mcp/src/error.rs b/litellm-rust/crates/gateway-mcp/src/error.rs index 1aa72e34703..9f08abaf26c 100644 --- a/litellm-rust/crates/gateway-mcp/src/error.rs +++ b/litellm-rust/crates/gateway-mcp/src/error.rs @@ -1,58 +1,231 @@ +use std::time::TryFromFloatSecsError; + use axum::{ Json, + extract::rejection::{JsonRejection, PathRejection, QueryRejection}, http::StatusCode, response::{IntoResponse, Response}, }; +use http::header::{InvalidHeaderName, InvalidHeaderValue, ToStrError}; +use litellm_error::{Report, Source}; use rmcp::{ model::{ErrorCode, ErrorData}, - service::ServiceError, + service::{ClientInitializeError, ServiceError}, }; use serde_json::json; #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("{0}")] - InvalidRequest(String), + #[error("{0} is required in request body")] + MissingField(&'static str), + #[error("Gateway listings do not accept an upstream cursor")] + UpstreamCursor, + #[error("read_resource requires exactly one allowed MCP server")] + AmbiguousResourceServer, + #[error("Invalid x-mcp-servers header")] + InvalidServersHeader(#[source] ToStrError), + #[error("invalid JSON body")] + Json(#[source] JsonRejection), + #[error("invalid query string")] + Query(#[source] QueryRejection), + #[error("invalid path")] + Path(#[source] PathRejection), + #[error("request body is too large")] + BodyTooLarge(#[source] axum::Error), #[error("User not allowed to access this MCP server or operation.")] Forbidden, - #[error("{0}")] - Configuration(String), + #[error("request host is not allowed")] + HostNotAllowed, + #[error("request origin is not allowed")] + OriginNotAllowed, + #[error("MCP session belongs to another caller")] + ForeignSession, + #[error("MCP session not found")] + UnknownSession, + #[error("MCP session is busy")] + SessionBusy, #[error("MCP upstream request failed")] Upstream(#[from] ServiceError), + #[error("MCP upstream returned invalid pagination")] + InvalidPagination, #[error("MCP operation returned an unexpected result")] UnexpectedResult, + #[error("MCP authorization unavailable")] + AuthorizationUnavailable(#[source] Source), #[error("MCP request cancelled")] Cancelled, } impl Error { - pub fn into_mcp(self) -> ErrorData { + pub fn status(&self) -> StatusCode { match self { - Self::Upstream(ServiceError::McpError(error)) => error, - Self::InvalidRequest(message) => ErrorData::invalid_params(message, None), - Self::Forbidden => ErrorData::new(ErrorCode(-32003), self.to_string(), None), - _ => ErrorData::internal_error(self.to_string(), None), + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) => StatusCode::BAD_REQUEST, + Self::Json(rejection) => rejection.status(), + Self::Query(rejection) => rejection.status(), + Self::Path(rejection) => rejection.status(), + Self::BodyTooLarge(_) => StatusCode::PAYLOAD_TOO_LARGE, + Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession => StatusCode::FORBIDDEN, + Self::UnknownSession => StatusCode::NOT_FOUND, + Self::SessionBusy => StatusCode::TOO_MANY_REQUESTS, + Self::Upstream(_) | Self::InvalidPagination => StatusCode::BAD_GATEWAY, + Self::Cancelled => StatusCode::REQUEST_TIMEOUT, + Self::UnexpectedResult | Self::AuthorizationUnavailable(_) => { + StatusCode::INTERNAL_SERVER_ERROR + } + } + } + + fn is_request(&self) -> bool { + match self { + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) + | Self::Json(_) + | Self::Query(_) + | Self::Path(_) + | Self::BodyTooLarge(_) => true, + Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession + | Self::UnknownSession + | Self::SessionBusy + | Self::Upstream(_) + | Self::InvalidPagination + | Self::UnexpectedResult + | Self::AuthorizationUnavailable(_) + | Self::Cancelled => false, + } + } + + fn detail(&self) -> String { + match self { + Self::Upstream(_) | Self::AuthorizationUnavailable(_) => self.to_string(), + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) + | Self::Json(_) + | Self::Query(_) + | Self::Path(_) + | Self::BodyTooLarge(_) + | Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession + | Self::UnknownSession + | Self::SessionBusy + | Self::InvalidPagination + | Self::UnexpectedResult + | Self::Cancelled => Report::new(self).to_string(), + } + } + + fn upstream(&self) -> Option<&ServiceError> { + match self { + Self::Upstream(error) => Some(error), + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) + | Self::Json(_) + | Self::Query(_) + | Self::Path(_) + | Self::BodyTooLarge(_) + | Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession + | Self::UnknownSession + | Self::SessionBusy + | Self::InvalidPagination + | Self::UnexpectedResult + | Self::AuthorizationUnavailable(_) + | Self::Cancelled => None, + } + } + + pub(crate) fn is_upstream(&self) -> bool { + self.upstream().is_some() + } + + pub(crate) fn is_timeout(&self) -> bool { + self.upstream().is_some_and(|error| match error { + ServiceError::Timeout { .. } => true, + ServiceError::McpError(_) + | ServiceError::TransportSend(_) + | ServiceError::TransportClosed + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => false, + _ => false, + }) + } + + pub(crate) fn is_unreachable(&self) -> bool { + self.upstream().is_some_and(|error| match error { + ServiceError::TransportSend(_) | ServiceError::TransportClosed => true, + ServiceError::McpError(_) + | ServiceError::Timeout { .. } + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => false, + _ => false, + }) + } + + pub(crate) fn mcp_error(&self) -> Option<&ErrorData> { + self.upstream().and_then(|error| match error { + ServiceError::McpError(error) => Some(error), + ServiceError::TransportSend(_) + | ServiceError::TransportClosed + | ServiceError::Timeout { .. } + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => None, + _ => None, + }) + } + + pub fn into_mcp(self) -> ErrorData { + if let Some(error) = self.mcp_error() { + return error.clone(); + } + let message = self.detail(); + if self.status() == StatusCode::FORBIDDEN { + ErrorData::new(ErrorCode(-32003), message, None) + } else if self.is_request() { + ErrorData::invalid_params(message, None) + } else { + ErrorData::internal_error(message, None) } } } impl IntoResponse for Error { fn into_response(self) -> Response { - let status = match &self { - Self::InvalidRequest(_) => StatusCode::BAD_REQUEST, - Self::Forbidden => StatusCode::FORBIDDEN, - Self::Upstream(_) => StatusCode::BAD_GATEWAY, - Self::Cancelled => StatusCode::REQUEST_TIMEOUT, - Self::Configuration(_) | Self::UnexpectedResult => StatusCode::INTERNAL_SERVER_ERROR, - }; - (status, Json(json!({"detail": self.to_string()}))).into_response() + (self.status(), Json(json!({"detail": self.detail()}))).into_response() } } -#[derive(thiserror::Error)] +#[derive(thiserror::Error, veil::Redact)] pub enum ConnectError { - #[error("MCP server {server}: {message}")] - Configuration { server: String, message: String }, + #[error("MCP server {server} is misconfigured")] + Configuration { + server: String, + #[source] + fault: ConfigFault, + }, #[error("could not resolve MCP credentials")] Secret(#[from] litellm_secrets::Error), #[error("could not start MCP child process")] @@ -61,38 +234,174 @@ pub enum ConnectError { Initialize { server: String, #[source] - source: Box, + #[redact] + source: Box, }, #[error("MCP upstream {0} initialization timed out")] Timeout(String), #[error("MCP upstream startup cancelled")] Cancelled, - #[error(transparent)] - Registry(#[from] Error), + #[error("MCP server identifiers must not be empty")] + EmptyIdentifier, + #[error("MCP server identifiers and prefixes must be unique")] + DuplicateIdentifier, } -impl std::fmt::Debug for ConnectError { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(self, formatter) - } +#[derive(Debug, thiserror::Error)] +pub enum ConfigFault { + #[error("timeout must be finite and positive")] + InvalidTimeout(#[source] TryFromFloatSecsError), + #[error("timeout must be positive")] + ZeroTimeout, + #[error("server identifiers must not be blank")] + BlankIdentifier, + #[error("setting '{0}' is not supported by the Rust MCP host")] + UnsupportedSetting(String), + #[error("max_concurrent_requests is out of range")] + ConcurrencyOutOfRange, + #[error("mcp_info protocol and cost settings are not supported")] + UnsupportedMcpInfo, + #[error("configured upstream auth mode is not supported by the Rust MCP host")] + UnsupportedAuth, + #[error("legacy SSE upstreams are not supported by the pinned SDK; use transport: http")] + SseTransport, + #[error("HTTP transport requires url and does not accept command, args or env")] + HttpShape, + #[error("stdio transport requires command and does not accept url")] + StdioShape, + #[error("stdio credentials must be supplied through env")] + StdioCredentials, + #[error("authentication_token and upstream_token_header require an auth_type")] + TokenWithoutAuth, + #[error("configured auth_type requires authentication_token")] + AuthWithoutToken, + #[error("referenced secret is unavailable")] + MissingSecret, + #[error("upstream credential is empty")] + EmptyCredential, + #[error("invalid upstream header name")] + HeaderName(#[source] InvalidHeaderName), + #[error("invalid upstream header value")] + HeaderValue(#[source] InvalidHeaderValue), + #[error("upstream header conflicts with HTTP framing")] + FramingHeader, + #[error("invalid HTTP URL")] + InvalidUrl(#[source] url::ParseError), + #[error("HTTP URL must use http or https without user credentials or fragments")] + UrlShape, } #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; - #[rstest::rstest] + const SECRET_URL: &str = "https://example.test/mcp?token=private-secret"; + + #[rstest] fn startup_diagnostics_do_not_expose_upstream_details() { let error = ConnectError::Initialize { server: "docs".into(), - source: Box::new(rmcp::service::ClientInitializeError::ConnectionClosed( - "https://example.test/mcp?token=private-secret".into(), - )), + source: Box::new(ClientInitializeError::ConnectionClosed(SECRET_URL.into())), }; - assert_eq!( - format!("{error:?}"), - "could not initialize MCP upstream docs" - ); + let debug = format!("{error:?}"); + assert!(debug.contains("docs")); + assert!(!debug.contains("private-secret")); assert!(std::error::Error::source(&error).is_some()); } + + #[rstest] + fn configuration_faults_reach_the_startup_diagnostic() { + let error = ConnectError::Configuration { + server: "docs".into(), + fault: ConfigFault::ZeroTimeout, + }; + assert!(format!("{error:?}").contains("ZeroTimeout")); + let report = Report::new(&error).to_string(); + assert!(report.contains("docs"), "{report}"); + assert_eq!( + report.matches("timeout must be positive").count(), + 1, + "{report}" + ); + } + + #[rstest] + #[case::missing_field(Error::MissingField("server_id"), StatusCode::BAD_REQUEST)] + #[case::upstream_cursor(Error::UpstreamCursor, StatusCode::BAD_REQUEST)] + #[case::forbidden(Error::Forbidden, StatusCode::FORBIDDEN)] + #[case::foreign_session(Error::ForeignSession, StatusCode::FORBIDDEN)] + #[case::unknown_session(Error::UnknownSession, StatusCode::NOT_FOUND)] + #[case::busy_session(Error::SessionBusy, StatusCode::TOO_MANY_REQUESTS)] + #[case::invalid_pagination(Error::InvalidPagination, StatusCode::BAD_GATEWAY)] + #[case::upstream( + Error::Upstream(ServiceError::TransportClosed), + StatusCode::BAD_GATEWAY + )] + #[case::cancelled(Error::Cancelled, StatusCode::REQUEST_TIMEOUT)] + fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { + assert_eq!(error.status(), status); + } + + fn mcp_failure() -> Error { + Error::Upstream(ServiceError::McpError(ErrorData::invalid_params( + "bad", None, + ))) + } + + #[rstest] + #[case::timeout( + Error::Upstream(ServiceError::Timeout { timeout: std::time::Duration::from_secs(1) }), + true, + false, + false + )] + #[case::closed(Error::Upstream(ServiceError::TransportClosed), false, true, false)] + #[case::mcp_error(mcp_failure(), false, false, true)] + #[case::unexpected_response( + Error::Upstream(ServiceError::UnexpectedResponse), + false, + false, + false + )] + #[case::gateway_cancelled(Error::Cancelled, false, false, false)] + fn upstream_failures_answer_what_went_wrong( + #[case] error: Error, + #[case] timeout: bool, + #[case] unreachable: bool, + #[case] mcp_error: bool, + ) { + assert_eq!(error.is_timeout(), timeout); + assert_eq!(error.is_unreachable(), unreachable); + assert_eq!(error.mcp_error().is_some(), mcp_error); + } + + #[rstest] + fn an_upstream_mcp_error_passes_through_unchanged() { + let mapped = mcp_failure().into_mcp(); + assert_eq!(mapped.code, ErrorCode::INVALID_PARAMS); + assert_eq!(mapped.message, "bad"); + } + + #[rstest] + #[case::request(Error::MissingField("name"), ErrorCode::INVALID_PARAMS)] + #[case::forbidden(Error::Forbidden, ErrorCode(-32003))] + #[case::upstream_fault(Error::InvalidPagination, ErrorCode::INTERNAL_ERROR)] + fn mcp_error_code_follows_the_fault(#[case] error: Error, #[case] code: ErrorCode) { + assert_eq!(error.into_mcp().code, code); + } + + #[rstest] + #[case::upstream(Error::Upstream(ServiceError::Cancelled { + reason: Some(SECRET_URL.into()) + }))] + #[case::authorization(Error::AuthorizationUnavailable(Source::new(std::io::Error::other( + SECRET_URL + ))))] + fn private_causes_stay_out_of_client_messages(#[case] error: Error) { + assert!(Report::new(&error).to_string().contains("private-secret")); + assert!(!error.detail().contains("private-secret")); + assert!(!error.into_mcp().message.contains("private-secret")); + } } diff --git a/litellm-rust/crates/gateway-mcp/src/legacy.rs b/litellm-rust/crates/gateway-mcp/src/legacy.rs index 429fd933135..1673e1b7b19 100644 --- a/litellm-rust/crates/gateway-mcp/src/legacy.rs +++ b/litellm-rust/crates/gateway-mcp/src/legacy.rs @@ -1,8 +1,8 @@ use std::{sync::Arc, time::Duration}; use axum::{ - Json, Router, - extract::{Query, Request, State}, + Router, + extract::{FromRequest, Request, State}, http::{StatusCode, request::Parts}, response::{ IntoResponse, Response, Sse, @@ -18,7 +18,11 @@ use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use uuid::Uuid; -use crate::{Context, HttpConfig, McpServer, Operations, sessions::fingerprint}; +use crate::{ + Context, Error, HttpConfig, McpServer, Operations, + request::{Json, Query}, + sessions::fingerprint, +}; #[derive(Clone)] struct Session { @@ -63,13 +67,13 @@ pub(crate) fn router( } impl Legacy { - fn validate(&self, parts: &Parts) -> Result<(), StatusCode> { + fn validate(&self, parts: &Parts) -> Result<(), Error> { let host = parts .headers .get("host") .and_then(|value| value.to_str().ok()) .and_then(|value| value.parse::().ok()) - .ok_or(StatusCode::FORBIDDEN)?; + .ok_or(Error::HostNotAllowed)?; let allowed = self.hosts.iter().any(|allowed| { host.as_str().eq_ignore_ascii_case(allowed) || host @@ -78,17 +82,16 @@ impl Legacy { .eq_ignore_ascii_case(allowed) }); if !allowed { - return Err(StatusCode::FORBIDDEN); + return Err(Error::HostNotAllowed); } - if let Some(origin) = parts.headers.get("origin") { - let origin = origin.to_str().map_err(|_| StatusCode::FORBIDDEN)?; - if !self - .origins - .iter() - .any(|allowed| same_origin(origin, allowed)) - { - return Err(StatusCode::FORBIDDEN); - } + if let Some(origin) = parts.headers.get("origin") + && !origin.to_str().is_ok_and(|origin| { + self.origins + .iter() + .any(|allowed| same_origin(origin, allowed)) + }) + { + return Err(Error::OriginNotAllowed); } Ok(()) } @@ -111,27 +114,22 @@ fn same_origin(value: &str, allowed: &str) -> bool { } } -async fn connect(State(state): State, request: Request) -> Response { +async fn connect(State(state): State, request: Request) -> Result { let owner = fingerprint(&request); let (parts, _) = request.into_parts(); - if let Err(status) = state.validate(&parts) { - return status.into_response(); - } - if let Err(error) = state + state.validate(&parts)?; + state .operations .authorize(Context { parts, server: None, mcp: None, }) - .await - { - return error.into_response(); - } + .await?; let connection = start_session(state, owner).await; - Sse::new(connection.events()) + Ok(Sse::new(connection.events()) .keep_alive(KeepAlive::default()) - .into_response() + .into_response()) } struct Connection { @@ -216,29 +214,21 @@ async fn message( State(state): State, Query(query): Query, request: Request, -) -> Response { +) -> Result { let owner = fingerprint(&request); let (parts, body) = request.into_parts(); - if let Err(status) = state.validate(&parts) { - return status.into_response(); - } - let Some(session) = state.sessions.get(&query.session_id).await else { - return StatusCode::NOT_FOUND.into_response(); - }; - if session.owner != owner { - return StatusCode::FORBIDDEN.into_response(); - } - let request = Request::from_parts(parts.clone(), body); - let Json(mut message) = - match as axum::extract::FromRequest<()>>::from_request( - request, - &(), - ) + state.validate(&parts)?; + let session = state + .sessions + .get(&query.session_id) .await - { - Ok(message) => message, - Err(error) => return error.into_response(), - }; + .ok_or(Error::UnknownSession)?; + if session.owner != owner { + return Err(Error::ForeignSession); + } + let Json(mut message) = + Json::::from_request(Request::from_parts(parts.clone(), body), &()) + .await?; match &mut message { ClientJsonRpcMessage::Request(request) => { request.request.extensions_mut().insert(parts); @@ -249,8 +239,8 @@ async fn message( _ => (), } match session.sender.try_send(message) { - Ok(()) => (StatusCode::ACCEPTED, "Accepted").into_response(), - Err(mpsc::error::TrySendError::Full(_)) => StatusCode::TOO_MANY_REQUESTS.into_response(), - Err(mpsc::error::TrySendError::Closed(_)) => StatusCode::NOT_FOUND.into_response(), + Ok(()) => Ok((StatusCode::ACCEPTED, "Accepted").into_response()), + Err(mpsc::error::TrySendError::Full(_)) => Err(Error::SessionBusy), + Err(mpsc::error::TrySendError::Closed(_)) => Err(Error::UnknownSession), } } diff --git a/litellm-rust/crates/gateway-mcp/src/lib.rs b/litellm-rust/crates/gateway-mcp/src/lib.rs index dd72940e24b..fcb7fd98383 100644 --- a/litellm-rust/crates/gateway-mcp/src/lib.rs +++ b/litellm-rust/crates/gateway-mcp/src/lib.rs @@ -3,6 +3,7 @@ mod error; mod legacy; mod native; mod relay; +mod request; mod routes; mod server; mod sessions; @@ -13,7 +14,7 @@ use http::request::Parts; use rmcp::{RoleServer, model::*, service::RequestContext}; pub use configured::ConfiguredGateway; -pub use error::{ConnectError, Error}; +pub use error::{ConfigFault, ConnectError, Error}; pub use native::{Limits, NativeGateway, Registry, Server, ServerResolver}; pub use relay::RelayClient; pub use rmcp; diff --git a/litellm-rust/crates/gateway-mcp/src/native.rs b/litellm-rust/crates/gateway-mcp/src/native.rs index 3a60ce5aca7..7a83f1316e9 100644 --- a/litellm-rust/crates/gateway-mcp/src/native.rs +++ b/litellm-rust/crates/gateway-mcp/src/native.rs @@ -178,14 +178,15 @@ impl NativeGateway { .ok_or(Error::Forbidden)?; let mut upstream = request; upstream.name = tool.tool.name.clone(); - match self + let result = self .send(server, CallToolRequest::new(upstream).into(), &context) - .await - { - Err(Error::Upstream(rmcp::service::ServiceError::McpError(error))) => { - Ok(CallToolResult::error(vec![ContentBlock::text(error.message)]).into()) - } - result => result, + .await; + match result.as_ref().err().and_then(Error::mcp_error) { + Some(error) => Ok(CallToolResult::error(vec![ContentBlock::text( + error.message.clone(), + )]) + .into()), + None => result, } } Operation::ListPrompts(params) => { @@ -321,9 +322,7 @@ impl NativeGateway { return Err(Error::Forbidden); } let [server] = servers.as_ref() else { - return Err(Error::InvalidRequest( - "read_resource requires exactly one allowed MCP server".into(), - )); + return Err(Error::AmbiguousResourceServer); }; self.send(server, ReadResourceRequest::new(request).into(), &context) .await @@ -363,9 +362,7 @@ impl Operations for NativeGateway { fn reject_cursor(params: &Option) -> Result<(), Error> { match params.as_ref().and_then(|params| params.cursor.as_ref()) { - Some(_) => Err(Error::InvalidRequest( - "Gateway listings do not accept an upstream cursor".into(), - )), + Some(_) => Err(Error::UpstreamCursor), None => Ok(()), } } diff --git a/litellm-rust/crates/gateway-mcp/src/native/catalog.rs b/litellm-rust/crates/gateway-mcp/src/native/catalog.rs index 448729ec686..6dd16d1fe63 100644 --- a/litellm-rust/crates/gateway-mcp/src/native/catalog.rs +++ b/litellm-rust/crates/gateway-mcp/src/native/catalog.rs @@ -28,9 +28,7 @@ impl NativeGateway { let state = match next { None => None, Some(cursor) if seen.contains(&cursor) || seen.len() >= 1000 => { - return Err(Error::InvalidRequest( - "MCP upstream returned invalid pagination".into(), - )); + return Err(Error::InvalidPagination); } Some(cursor) => Some(( Some(cursor.clone()), @@ -95,14 +93,9 @@ impl NativeGateway { Ok(tools) => ServerOutcome::Ok { tool_count: tools.len(), }, - Err(Error::Upstream(rmcp::service::ServiceError::Timeout { .. })) => { - ServerOutcome::Timeout - } - Err(Error::Upstream( - rmcp::service::ServiceError::TransportClosed - | rmcp::service::ServiceError::TransportSend(_), - )) => ServerOutcome::Unreachable, - Err(Error::Upstream(_)) => ServerOutcome::UpstreamError, + Err(error) if error.is_timeout() => ServerOutcome::Timeout, + Err(error) if error.is_unreachable() => ServerOutcome::Unreachable, + Err(error) if error.is_upstream() => ServerOutcome::UpstreamError, Err(_) => ServerOutcome::Internal, }; (server.prefix(), outcome) diff --git a/litellm-rust/crates/gateway-mcp/src/native/registry.rs b/litellm-rust/crates/gateway-mcp/src/native/registry.rs index 0aa051d6f32..231737e9772 100644 --- a/litellm-rust/crates/gateway-mcp/src/native/registry.rs +++ b/litellm-rust/crates/gateway-mcp/src/native/registry.rs @@ -1,5 +1,5 @@ use super::NativeGateway; -use crate::{Context, Error, GatewayFuture, ServerInfo}; +use crate::{ConnectError, Context, Error, GatewayFuture, ServerInfo}; use rmcp::{Peer, RoleClient}; use std::sync::Arc; @@ -39,14 +39,12 @@ pub trait ServerResolver: Send + Sync + 'static { pub struct Registry(Arc<[Server]>); impl Registry { - pub fn new(servers: impl Into>) -> Result { + pub fn new(servers: impl Into>) -> Result { let servers = servers.into(); for (index, server) in servers.iter().enumerate() { let identifiers = [&server.info.server_id, &server.info.server_name]; if identifiers.iter().any(|value| value.is_empty()) || server.prefix().is_empty() { - return Err(Error::Configuration( - "MCP server identifiers must not be empty".into(), - )); + return Err(ConnectError::EmptyIdentifier); } if servers[..index].iter().any(|previous| { previous.prefix() == server.prefix() @@ -57,9 +55,7 @@ impl Registry { .as_deref() .is_some_and(|alias| previous.matches(alias)) }) { - return Err(Error::Configuration( - "MCP server identifiers and prefixes must be unique".into(), - )); + return Err(ConnectError::DuplicateIdentifier); } } Ok(Self(servers)) @@ -79,11 +75,7 @@ impl NativeGateway { .parts .headers .get("x-mcp-servers") - .map(|value| { - value - .to_str() - .map_err(|_| Error::InvalidRequest("Invalid x-mcp-servers header".into())) - }) + .map(|value| value.to_str().map_err(Error::InvalidServersHeader)) .transpose()?; let scope = context.server.as_deref().or(header); let Some(scope) = scope else { diff --git a/litellm-rust/crates/gateway-mcp/src/request.rs b/litellm-rust/crates/gateway-mcp/src/request.rs new file mode 100644 index 00000000000..eaa9a6f1840 --- /dev/null +++ b/litellm-rust/crates/gateway-mcp/src/request.rs @@ -0,0 +1,44 @@ +use axum::extract::{FromRequest, FromRequestParts, Request}; +use http::request::Parts; +use serde::de::DeserializeOwned; + +use crate::Error; + +pub(crate) struct Json(pub T); + +impl FromRequest for Json { + type Rejection = Error; + + async fn from_request(request: Request, state: &S) -> Result { + axum::Json::::from_request(request, state) + .await + .map(|axum::Json(value)| Self(value)) + .map_err(Error::Json) + } +} + +pub(crate) struct Query(pub T); + +impl FromRequestParts for Query { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + axum::extract::Query::::from_request_parts(parts, state) + .await + .map(|axum::extract::Query(value)| Self(value)) + .map_err(Error::Query) + } +} + +pub(crate) struct Path(pub T); + +impl FromRequestParts for Path { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + axum::extract::Path::::from_request_parts(parts, state) + .await + .map(|axum::extract::Path(value)| Self(value)) + .map_err(Error::Path) + } +} diff --git a/litellm-rust/crates/gateway-mcp/src/routes.rs b/litellm-rust/crates/gateway-mcp/src/routes.rs index b5d314e4882..56d67a5bc77 100644 --- a/litellm-rust/crates/gateway-mcp/src/routes.rs +++ b/litellm-rust/crates/gateway-mcp/src/routes.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use axum::{ Json, Router, - extract::{Path, Query, Request, State}, + extract::{Request, State}, http::request::Parts, middleware::{self, Next}, response::{IntoResponse, Response}, @@ -12,7 +12,11 @@ use rmcp::{model::*, transport::streamable_http_server::StreamableHttpServerConf use serde::{Deserialize, Serialize}; use tokio_util::sync::CancellationToken; -use crate::{Context, Error, McpServer, Operation, Operations, ToolCatalog, server::ServerScope}; +use crate::{ + Context, Error, McpServer, Operation, Operations, ToolCatalog, + request::{self, Path, Query}, + server::ServerScope, +}; pub struct HttpConfig { pub allowed_hosts: Vec, @@ -132,16 +136,16 @@ struct ToolCall { async fn rest_call( State(operations): State>, parts: Parts, - Json(call): Json, + request::Json(call): request::Json, ) -> Result { let server = call .server_id .filter(|id| !id.is_empty()) - .ok_or_else(|| Error::InvalidRequest("server_id is required in request body".into()))?; + .ok_or(Error::MissingField("server_id"))?; let name = call .name .filter(|name| !name.is_empty()) - .ok_or_else(|| Error::InvalidRequest("name is required in request body".into()))?; + .ok_or(Error::MissingField("name"))?; let params = CallToolRequestParams::new(name).with_arguments(call.arguments.unwrap_or_default()); let context = Context { diff --git a/litellm-rust/crates/gateway-mcp/src/sessions.rs b/litellm-rust/crates/gateway-mcp/src/sessions.rs index 54d0634bd3d..583a8b86341 100644 --- a/litellm-rust/crates/gateway-mcp/src/sessions.rs +++ b/litellm-rust/crates/gateway-mcp/src/sessions.rs @@ -4,7 +4,7 @@ use axum::{ Router, body::{Body, to_bytes}, extract::{Request, State}, - http::{Method, StatusCode}, + http::Method, response::IntoResponse, routing::any, }; @@ -19,7 +19,7 @@ use rmcp::{ use sha2::{Digest, Sha256}; use tower::ServiceExt; -use crate::{McpServer, server::ServerScope}; +use crate::{Error, McpServer, server::ServerScope}; #[derive(Clone)] pub struct SessionOwner(pub String); @@ -65,7 +65,7 @@ pub(crate) fn transport(server: McpServer, config: StreamableHttpServerConfig) - async fn dispatch( State(transport): State, request: Request, -) -> Result { +) -> Result { let owner = fingerprint(&request); let session = request .headers() @@ -75,8 +75,8 @@ async fn dispatch( if let Some(session) = &session { match transport.owners.get(session).await { Some(expected) if expected == owner => (), - Some(_) => return Err(StatusCode::FORBIDDEN), - None => return Err(StatusCode::NOT_FOUND), + Some(_) => return Err(Error::ForeignSession), + None => return Err(Error::UnknownSession), } } let method = request.method().clone(); @@ -84,7 +84,7 @@ async fn dispatch( let (parts, body) = request.into_parts(); let bytes = to_bytes(body, 4 * 1024 * 1024) .await - .map_err(|_| StatusCode::PAYLOAD_TOO_LARGE)?; + .map_err(Error::BodyTooLarge)?; let initialize = matches!(serde_json::from_slice::(&bytes), Ok(ClientJsonRpcMessage::Request(request)) if matches!(request.request, ClientRequest::InitializeRequest(_))); (Request::from_parts(parts, Body::from(bytes)), initialize) diff --git a/litellm-rust/crates/gateway-mcp/tests/gateway/http.rs b/litellm-rust/crates/gateway-mcp/tests/gateway/http.rs index 5feab0194d3..c5060ebe8f3 100644 --- a/litellm-rust/crates/gateway-mcp/tests/gateway/http.rs +++ b/litellm-rust/crates/gateway-mcp/tests/gateway/http.rs @@ -171,43 +171,44 @@ async fn resolver_receives_each_requests_identity( Duration::from_secs(2), )); let app = router(operations, HttpConfig::default()); - for token in ["allowed", "denied", "allowed"] { - let response = app - .clone() - .oneshot( - Request::builder() - .method(method) - .uri(path) - .header("host", "localhost") - .header("authorization", format!("Bearer {token}")) - .header("content-type", "application/json") - .header("accept", "application/json, text/event-stream") - .header("mcp-protocol-version", "2025-11-25") - .body(Body::from( - json!({"jsonrpc":"2.0","id":1,"method":"tools/list"}).to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - let status = response.status(); - let body = String::from_utf8( - to_bytes(response.into_body(), 65536) - .await - .unwrap() - .to_vec(), + let (status, body) = list_as(&app, method, path, "allowed").await; + assert_eq!(status, StatusCode::OK); + assert!(body.contains("echo-value"), "{body}"); + let (status, body) = list_as(&app, method, path, "denied").await; + assert_eq!(status == StatusCode::FORBIDDEN, method == "GET"); + assert!(body.contains("not allowed"), "{body}"); + let (status, body) = list_as(&app, method, path, "allowed").await; + assert_eq!(status, StatusCode::OK); + assert!(body.contains("echo-value"), "{body}"); +} + +async fn list_as( + app: &axum::Router, + method: &str, + path: &str, + token: &str, +) -> (StatusCode, String) { + let response = app + .clone() + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("host", "localhost") + .header("authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .header("accept", "application/json, text/event-stream") + .header("mcp-protocol-version", "2025-11-25") + .body(Body::from( + json!({"jsonrpc":"2.0","id":1,"method":"tools/list"}).to_string(), + )) + .unwrap(), ) + .await .unwrap(); - if token == "denied" { - if method == "GET" { - assert_eq!(status, StatusCode::FORBIDDEN); - } - assert!(body.contains("not allowed"), "{body}"); - } else { - assert_eq!(status, StatusCode::OK); - assert!(body.contains("echo-value"), "{body}"); - } - } + let status = response.status(); + let body = to_bytes(response.into_body(), 65536).await.unwrap(); + (status, String::from_utf8(body.to_vec()).unwrap()) } fn rpc_request(path: &str, token: &str, session: Option<&str>, body: Value) -> Request { @@ -352,10 +353,14 @@ async fn header_scope_cannot_broaden_path_scope( } #[rstest] -#[case::missing_server(json!({"name":"echo-value"}))] -#[case::missing_name(json!({"server_id":"alpha"}))] +#[case::missing_server(json!({"name":"echo-value"}), "server_id")] +#[case::missing_name(json!({"server_id":"alpha"}), "name")] #[tokio::test] -async fn rest_requires_server_and_tool(#[future] harness: Harness, #[case] body: Value) { +async fn rest_requires_server_and_tool( + #[future] harness: Harness, + #[case] body: Value, + #[case] field: &str, +) { let harness = harness.await; let response = harness .app @@ -368,6 +373,40 @@ async fn rest_requires_server_and_tool(#[future] harness: Harness, #[case] body: .await .unwrap(); assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert!(body["detail"].as_str().unwrap().contains(field), "{body}"); +} + +#[rstest] +#[case::malformed_call("/mcp-rest/tools/call", "POST", "{")] +#[case::unknown_list_parameter("/mcp-rest/tools/list?unknown=1", "GET", "")] +#[case::legacy_message_without_session("/mcp/sse/messages", "POST", "{}")] +#[tokio::test] +async fn rejected_requests_use_the_detail_envelope( + #[future] harness: Harness, + #[case] path: &str, + #[case] method: &str, + #[case] body: &'static str, +) { + let harness = harness.await; + let response = harness + .app + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("host", "localhost") + .header("content-type", "application/json") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + assert!(response.status().is_client_error(), "{}", response.status()); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert!(body["detail"].is_string(), "{body}"); } #[rstest] diff --git a/litellm-rust/crates/gateway-mcp/tests/gateway/native.rs b/litellm-rust/crates/gateway-mcp/tests/gateway/native.rs index aca19b38d6e..54fdabc0da7 100644 --- a/litellm-rust/crates/gateway-mcp/tests/gateway/native.rs +++ b/litellm-rust/crates/gateway-mcp/tests/gateway/native.rs @@ -148,10 +148,12 @@ async fn routes_prompts_and_resources_without_rewriting_uris(#[future] harness: "fixture://{file}" ); let read = Operation::ReadResource(ReadResourceRequestParams::new("fixture://document")); - assert!(matches!( - harness.gateway.execute(read.clone(), context(None)).await, - Err(Error::InvalidRequest(_)) - )); + let ambiguous = harness + .gateway + .execute(read.clone(), context(None)) + .await + .expect_err("read without a server scope"); + assert_eq!(ambiguous.status(), axum::http::StatusCode::BAD_REQUEST); let result = harness .gateway .execute(read, context(Some("alpha"))) diff --git a/litellm-rust/crates/gateway-ui/Cargo.toml b/litellm-rust/crates/gateway-ui/Cargo.toml index 94546ef9118..1c65665be72 100644 --- a/litellm-rust/crates/gateway-ui/Cargo.toml +++ b/litellm-rust/crates/gateway-ui/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true axum = { workspace = true, features = ["json", "original-uri", "query"] } axum-login.workspace = true base64.workspace = true @@ -22,6 +23,7 @@ tower-http = { version = "0.6.11", features = ["fs", "set-header"] } tower-sessions.workspace = true [dev-dependencies] +futures-util.workspace = true serde_json.workspace = true tower = { version = "0.5", features = ["util"] } litellm-auth-types.workspace = true diff --git a/litellm-rust/crates/gateway-ui/src/dashboard.rs b/litellm-rust/crates/gateway-ui/src/dashboard.rs index 2642506fde6..04ec8490411 100644 --- a/litellm-rust/crates/gateway-ui/src/dashboard.rs +++ b/litellm-rust/crates/gateway-ui/src/dashboard.rs @@ -2,7 +2,8 @@ use std::path::Path; use axum::{ Router, - extract::{Query, Request}, + extract::{FromRequestParts, Query, Request}, + http::request::Parts, routing::get, }; use serde::{Deserialize, Serialize}; @@ -34,6 +35,17 @@ struct LogoQuery { variant: Option, } +impl FromRequestParts for LogoQuery { + type Rejection = crate::Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + Query::::from_request_parts(parts, state) + .await + .map(|Query(query)| query) + .map_err(crate::Error::Query) + } +} + fn logo_file(query: &LogoQuery) -> &'static str { match (query.variant, query.theme) { (Some(Variant::Monogram), Some(Theme::Dark)) => "assets/logos/litellm_monogram_dark.svg", @@ -59,7 +71,7 @@ pub fn dashboard_assets(directory: impl AsRef) -> Router { ) .route( "/get_image", - get(move |Query(query): Query, request: Request| { + get(move |query: LogoQuery, request: Request| { ServeFile::new(logos.join(logo_file(&query))).oneshot(request) }), ) diff --git a/litellm-rust/crates/gateway-ui/src/error.rs b/litellm-rust/crates/gateway-ui/src/error.rs index 16e6e205d99..cce01ba76ce 100644 --- a/litellm-rust/crates/gateway-ui/src/error.rs +++ b/litellm-rust/crates/gateway-ui/src/error.rs @@ -1,9 +1,9 @@ use axum::{ Json, + extract::rejection::{JsonRejection, QueryRejection}, http::{StatusCode, header}, response::{IntoResponse, Response}, }; -use litellm_gateway_auth::UiAuthError; #[derive(serde::Serialize)] struct ErrorBody { @@ -21,39 +21,108 @@ pub enum Error { InvalidCredentials, #[error("Too many login attempts, try again in one minute")] RateLimited, - #[error("UI authentication unavailable")] - Auth(#[from] axum_login::Error), + #[error("invalid login request")] + Body(#[source] JsonRejection), + #[error("invalid query string")] + Query(#[source] QueryRejection), + #[error("UI login unavailable")] + Login(#[from] axum_login::Error), #[error("UI session unavailable")] Session(#[from] tower_sessions::session::Error), #[error("UI token unavailable")] Token(#[from] jsonwebtoken::errors::Error), - #[error("Invalid UI session")] - Unauthorized(#[from] UiAuthError), + #[error(transparent)] + Auth(#[from] litellm_gateway_auth::Error), +} + +impl Error { + pub fn status(&self) -> StatusCode { + match self { + Self::InvalidCredentials => StatusCode::UNAUTHORIZED, + Self::RateLimited => StatusCode::TOO_MANY_REQUESTS, + Self::Body(rejection) => rejection.status(), + Self::Query(rejection) => rejection.status(), + Self::Auth(error) => error.status(), + Self::Login(_) | Self::Session(_) | Self::Token(_) => StatusCode::INTERNAL_SERVER_ERROR, + } + } +} + +impl Error { + fn message(&self) -> String { + match self { + Self::Body(_) | Self::Query(_) => litellm_error::Report::new(self).to_string(), + Self::InvalidCredentials + | Self::RateLimited + | Self::Login(_) + | Self::Session(_) + | Self::Token(_) + | Self::Auth(_) => self.to_string(), + } + } } impl IntoResponse for Error { fn into_response(self) -> Response { - let status = match &self { - Self::InvalidCredentials | Self::Unauthorized(UiAuthError::Unauthorized) => { - StatusCode::UNAUTHORIZED + let status = self.status(); + let body = Json(ErrorBody { + error: ErrorMessage { + message: self.message(), + }, + }); + match status { + StatusCode::TOO_MANY_REQUESTS => { + (status, [(header::RETRY_AFTER, "60")], body).into_response() } - Self::RateLimited => StatusCode::TOO_MANY_REQUESTS, - _ => StatusCode::INTERNAL_SERVER_ERROR, - }; - let response = ( - status, - Json(ErrorBody { - error: ErrorMessage { - message: self.to_string(), - }, - }), - ) - .into_response(); - - if matches!(self, Self::RateLimited) { - return ([(header::RETRY_AFTER, "60")], response).into_response(); + _ => (status, body).into_response(), } - - response + } +} + +#[cfg(test)] +mod tests { + use litellm_gateway_auth::Error as AuthError; + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::rejected_session(Error::Auth(AuthError::InvalidSession), StatusCode::UNAUTHORIZED)] + #[case::session_layer_missing( + Error::Auth(AuthError::MissingSessionLayer), + StatusCode::INTERNAL_SERVER_ERROR + )] + #[case::operation_outside_session_scope( + Error::Auth(AuthError::Forbidden), + StatusCode::FORBIDDEN + )] + #[case::rate_limited(Error::RateLimited, StatusCode::TOO_MANY_REQUESTS)] + fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { + assert_eq!(error.status(), status); + } + + #[rstest] + #[case::rate_limited(Error::RateLimited, true)] + #[case::invalid_credentials(Error::InvalidCredentials, false)] + fn only_rate_limits_ask_to_retry_later(#[case] error: Error, #[case] retry_after: bool) { + assert_eq!( + error + .into_response() + .headers() + .contains_key(header::RETRY_AFTER), + retry_after + ); + } + + #[tokio::test] + async fn a_body_never_shows_what_the_backend_said() { + let backend = std::io::Error::other("vault token rejected"); + let response = Error::Auth(AuthError::Unavailable(litellm_error::Source::new(backend))) + .into_response(); + let body = axum::body::to_bytes(response.into_body(), 4096) + .await + .unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert!(!body.contains("vault token rejected"), "{body}"); } } diff --git a/litellm-rust/crates/gateway-ui/src/session.rs b/litellm-rust/crates/gateway-ui/src/session.rs index b80f40bb376..b5d847ad353 100644 --- a/litellm-rust/crates/gateway-ui/src/session.rs +++ b/litellm-rust/crates/gateway-ui/src/session.rs @@ -4,14 +4,18 @@ use serde::Serialize; use std::sync::Arc; use time::{Duration, OffsetDateTime}; -use axum::{Json, extract::State as ExtractState, response::IntoResponse}; +use axum::{ + Json, + extract::{FromRequest, Request, State as ExtractState}, + response::IntoResponse, +}; use governor::DefaultDirectRateLimiter; use tower_cookies::{Cookie, Cookies}; use tower_sessions::{Expiry, Session, cookie::SameSite}; use crate::Error; use litellm_gateway_auth::{ - AccessRequest, UI_CSRF_KEY, UiAction, UiAuthError, UiAuthSession, UiCredentials, UiSession, + AccessRequest, UI_CSRF_KEY, UiAction, UiAuthSession, UiCredentials, UiSession, }; pub struct State { @@ -60,12 +64,25 @@ struct Discovery { is_control_plane: bool, } +pub struct Credentials(UiCredentials); + +impl FromRequest for Credentials { + type Rejection = Error; + + async fn from_request(request: Request, state: &S) -> Result { + Json::::from_request(request, state) + .await + .map(|Json(credentials)| Self(credentials)) + .map_err(Error::Body) + } +} + pub async fn login( ExtractState(state): ExtractState>, mut auth: UiAuthSession, session: Session, cookies: Cookies, - Json(credentials): Json, + Credentials(credentials): Credentials, ) -> Result { state.login_limit.check().map_err(|_| Error::RateLimited)?; let user = auth @@ -107,13 +124,14 @@ pub async fn login( })) } -pub async fn info(session: Result) -> Result { +pub async fn info( + session: Result, +) -> Result { let session = session?; session .identity .authorize(AccessRequest::Ui(UiAction::SessionInfo)) - .await - .map_err(|_| UiAuthError::Unauthorized)?; + .await?; Ok(Json(SessionInfo { user_id: session.user.username, user_role: "proxy_admin", @@ -121,7 +139,7 @@ pub async fn info(session: Result) -> Result, + validated: Result, mut auth: UiAuthSession, cookies: Cookies, ) -> Result { @@ -129,8 +147,7 @@ pub async fn logout( validated .identity .authorize(AccessRequest::Ui(UiAction::Logout)) - .await - .map_err(|_| UiAuthError::Unauthorized)?; + .await?; auth.logout().await?; cookies.add( Cookie::build(("token", "")) diff --git a/litellm-rust/crates/gateway-ui/tests/assets.rs b/litellm-rust/crates/gateway-ui/tests/assets.rs index ba741fbf2ec..e2819033f95 100644 --- a/litellm-rust/crates/gateway-ui/tests/assets.rs +++ b/litellm-rust/crates/gateway-ui/tests/assets.rs @@ -200,6 +200,24 @@ async fn logo_discovery_points_to_served_image(dashboard: App) { ); } +#[rstest] +#[tokio::test] +async fn unknown_logo_theme_is_rejected_in_the_ui_envelope(dashboard: App) { + let response = dashboard + .router + .oneshot( + Request::get("/get_image?theme=purple") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert!(body["error"]["message"].as_str().unwrap().contains("theme")); +} + #[rstest] #[case::logo("/get_image")] #[case::logo_dark("/get_image?theme=dark")] diff --git a/litellm-rust/crates/gateway-ui/tests/routes.rs b/litellm-rust/crates/gateway-ui/tests/routes.rs index 4a3238f79c2..95246f33a35 100644 --- a/litellm-rust/crates/gateway-ui/tests/routes.rs +++ b/litellm-rust/crates/gateway-ui/tests/routes.rs @@ -330,24 +330,45 @@ async fn reauthentication_rotates_and_revokes_the_previous_session(app: App) { ); } +async fn failed_login(app: &App, username: &str) -> Response { + app.router + .clone() + .oneshot( + Request::post("/v2/login") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({"username": username, "password": "wrong-password"}).to_string(), + )) + .unwrap(), + ) + .await + .unwrap() +} + #[rstest] #[tokio::test] async fn failed_logins_are_generic_and_rate_limited(app: App) { - for attempt in 0..11 { - let response = app.router.clone().oneshot(Request::post("/v2/login") - .header(header::CONTENT_TYPE, "application/json") - .body(Body::from(json!({"username": format!("unknown-{attempt}"), "password": "wrong-password"}).to_string())).unwrap()).await.unwrap(); - if attempt < 10 { - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert_eq!( - body(response).await["error"]["message"], - "Invalid username or password" - ); - } else { - assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); - assert!(response.headers().contains_key(header::RETRY_AFTER)); - } - } + let first = failed_login(&app, "unknown-0").await; + assert_eq!(first.status(), StatusCode::UNAUTHORIZED); + assert_eq!( + body(first).await["error"]["message"], + "Invalid username or password" + ); + let usernames: Vec = (1..10) + .map(|attempt| format!("unknown-{attempt}")) + .collect(); + let burst = futures_util::future::join_all( + usernames + .iter() + .map(|username| failed_login(&app, username)), + ) + .await; + let statuses: Vec = burst.iter().map(Response::status).collect(); + assert_eq!(statuses, vec![StatusCode::UNAUTHORIZED; 9]); + let limited = failed_login(&app, "unknown-10").await; + assert_eq!(limited.status(), StatusCode::TOO_MANY_REQUESTS); + assert!(limited.headers().contains_key(header::RETRY_AFTER)); + assert!(body(limited).await["error"]["message"].is_string()); } #[rstest] @@ -365,6 +386,7 @@ async fn login_rejects_browser_form_posts(app: App) { .unwrap(); assert_eq!(response.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE); assert!(!response.headers().contains_key(header::SET_COOKIE)); + assert!(body(response).await["error"]["message"].is_string()); } #[rstest] @@ -468,6 +490,7 @@ async fn invalid_login_payloads_do_not_create_sessions( .unwrap(); assert_eq!(response.status(), status); assert!(!response.headers().contains_key(header::SET_COOKIE)); + assert!(body(response).await["error"]["message"].is_string()); } #[rstest] diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index b2a2f09bfc7..4d831aa2314 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -19,9 +19,10 @@ litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true litellm-auth-types.workspace = true litellm-config.workspace = true +litellm-error.workspace = true litellm-circuit-breaker.workspace = true litellm-redis-client.workspace = true -litellm-http = { workspace = true, features = ["mcp"] } +litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets.workspace = true litellm-tracing.workspace = true diff --git a/litellm-rust/crates/gateway/src/auth.rs b/litellm-rust/crates/gateway/src/auth.rs index 72892be783c..d5d8a229d0d 100644 --- a/litellm-rust/crates/gateway/src/auth.rs +++ b/litellm-rust/crates/gateway/src/auth.rs @@ -64,10 +64,8 @@ impl OperationAuthorizer for McpPolicy { } fn mcp_error(error: Error) -> litellm_gateway_mcp::Error { - match error { - Error::InvalidToken | Error::Expired | Error::Forbidden => { - litellm_gateway_mcp::Error::Forbidden - } - _ => litellm_gateway_mcp::Error::Configuration("MCP authorization unavailable".into()), + if error.is_denied() { + return litellm_gateway_mcp::Error::Forbidden; } + litellm_gateway_mcp::Error::AuthorizationUnavailable(litellm_error::Source::new(error)) } diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index 170a8284c18..e2f646c5448 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -6,6 +6,8 @@ pub enum Error { Environment, #[error("MCP host does not yet support configured setting {0}")] McpSetting(String), + #[error("MCP host does not yet support guardrails or policies")] + McpGuardrails, #[error(transparent)] Mcp(#[from] litellm_gateway_mcp::ConnectError), #[error(transparent)] @@ -17,6 +19,6 @@ pub enum Error { }, #[error(transparent)] RedisConnect(#[from] litellm_redis_client::ConnectError), - #[error("MCP requires a configured master key")] + #[error("MCP requires working gateway authentication")] Auth(#[from] litellm_gateway_auth::Error), } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 2ed2fc80a51..b8ff8e46033 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -109,16 +109,16 @@ pub async fn build_mcp( || !config.policies.is_empty() || !config.policy_attachments.is_empty() { - return Err(Error::McpSetting("guardrails or policies".into())); + return Err(Error::McpGuardrails); } Auth::from_config(config, secrets.clone()) .validate() .await?; - let client = pool.mcp_client(http)?; + let client = pool.client(http, ClientVariant::NoRedirect)?; Ok(Some( litellm_gateway_mcp::ConfiguredGateway::connect( &config.mcp_servers, - client, + (*client).clone(), secrets.as_ref(), shutdown, ) diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index 593fbfeb129..3d12efd730e 100644 --- a/litellm-rust/crates/gateway/tests/mcp.rs +++ b/litellm-rust/crates/gateway/tests/mcp.rs @@ -78,7 +78,9 @@ async fn upstream() -> Fixture { let transport = StreamableHttpService::new( || Ok(Upstream), Arc::new(LocalSessionManager::default()), - StreamableHttpServerConfig::default().with_cancellation_token(shutdown.clone()), + StreamableHttpServerConfig::default() + .with_allowed_hosts(["127.0.0.1", "upstream.invalid"]) + .with_cancellation_token(shutdown.clone()), ); let app = Router::new().route_service("/mcp", transport); let stop = shutdown.clone(); @@ -95,11 +97,29 @@ async fn upstream() -> Fixture { } } -#[rstest] -#[tokio::test] -async fn serves_configured_upstream_with_separate_admission_and_outbound_credentials( - #[future] upstream: Fixture, -) { +struct Served { + app: Router, + mcp: litellm_gateway_mcp::ConfiguredGateway, + shutdown: CancellationToken, + upstream: Fixture, +} + +impl Served { + async fn close(self) { + self.shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(5), self.mcp.close()) + .await + .unwrap(); + self.upstream.shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(5), self.upstream.task) + .await + .unwrap() + .unwrap(); + } +} + +#[fixture] +async fn served(#[future] upstream: Fixture) -> Served { let upstream = upstream.await; let config = Config::from_yaml(&format!( r#" @@ -146,30 +166,53 @@ mcp_servers: }, ); let app = litellm_gateway::router(inference, &config, None, Some(mcp_router)); - - for (path, method) in [ - ("/mcp", "POST"), - ("/mcp", "GET"), - ("/mcp", "DELETE"), - ("/mcp/sse", "GET"), - ("/mcp/sse/messages", "POST"), - ("/mcp-rest/tools/list", "GET"), - ("/mcp/enabled", "GET"), - ] { - let response = app - .clone() - .oneshot( - Request::builder() - .method(method) - .uri(path) - .header("host", "localhost") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), 401, "{method} {path}"); + Served { + app, + mcp, + shutdown, + upstream, } +} + +#[rstest] +#[case::streamable_post("/mcp", "POST")] +#[case::streamable_get("/mcp", "GET")] +#[case::streamable_delete("/mcp", "DELETE")] +#[case::legacy_sse("/mcp/sse", "GET")] +#[case::legacy_messages("/mcp/sse/messages", "POST")] +#[case::rest_list("/mcp-rest/tools/list", "GET")] +#[case::enabled("/mcp/enabled", "GET")] +#[tokio::test] +async fn every_mcp_route_requires_admission_credentials( + #[future] served: Served, + #[case] path: &str, + #[case] method: &str, +) { + let served = served.await; + let response = served + .app + .clone() + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("host", "localhost") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), 401); + served.close().await; +} + +#[rstest] +#[tokio::test] +async fn serves_configured_upstream_with_separate_admission_and_outbound_credentials( + #[future] served: Served, +) { + let served = served.await; + let app = served.app.clone(); let list = app .clone() .oneshot( @@ -211,6 +254,49 @@ mcp_servers: .unwrap(); assert_eq!(blocked.status(), 403); + served.close().await; +} + +#[rstest] +#[tokio::test] +async fn mcp_tools_are_discovered_through_the_configured_proxy(#[future] upstream: Fixture) { + let upstream = upstream.await; + let config = Config::from_yaml( + "general_settings: {master_key: gateway-key}\nmcp_servers: {docs: {url: 'http://upstream.invalid/mcp', timeout: 2}}", + ) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let proxy = upstream.url.trim_end_matches("/mcp"); + let http = litellm_http::HttpClientConfig { + proxies: litellm_http::EnvironmentProxies::from_environment(&|name: &str| { + (name == "HTTP_PROXY").then(|| proxy.to_owned()) + }), + ..inference.http.clone() + }; + let shutdown = CancellationToken::new(); + let mcp = litellm_gateway::build_mcp( + &config, + inference.secrets.clone(), + shutdown.clone(), + &inference.resources.pool, + &http, + ) + .await + .unwrap() + .unwrap(); + let app = litellm_gateway_mcp::router(mcp.operations.clone(), HttpConfig::default()); + let response = app + .oneshot( + Request::get("/mcp-rest/tools/list?server_id=docs") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), 200); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert_eq!(body["tools"][0]["name"], "echo"); shutdown.cancel(); tokio::time::timeout(Duration::from_secs(5), mcp.close()) .await @@ -222,6 +308,47 @@ mcp_servers: .unwrap(); } +#[rstest] +#[rstest] +#[tokio::test] +async fn mcp_uses_the_configured_ca_bundle() { + let directory = tempfile::tempdir().unwrap(); + let missing = directory.path().join("missing-ca.pem"); + let config = Config::from_yaml( + "general_settings: {master_key: gateway-key}\nmcp_servers: {docs: {url: 'http://127.0.0.1:1/mcp'}}", + ) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let http = litellm_http::HttpClientConfig { + verify: litellm_http::Verify::CaBundle(missing.clone()), + ..inference.http.clone() + }; + let error = litellm_gateway::build_mcp( + &config, + inference.secrets.clone(), + CancellationToken::new(), + &inference.resources.pool, + &http, + ) + .await + .err() + .unwrap(); + let litellm_gateway::Error::Http(error) = error else { + panic!("expected an HTTP configuration error"); + }; + assert_eq!(error.tls_source(), Some(litellm_http::TlsSource::CaBundle)); + assert_eq!( + error.to_string(), + format!("could not read {}", missing.display()) + ); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .kind(), + std::io::ErrorKind::NotFound, + ); +} + #[rstest] #[case::unsupported_auth("auth_type: oauth2", "auth mode")] #[case::unsupported_policy("allowed_params: {echo: [safe]}", "allowed_params")] @@ -252,9 +379,10 @@ async fn rejects_unsupported_or_invalid_settings_before_serving( &inference.http, ) .await; - let error = result.err().expect("startup must fail").to_string(); - assert!(error.contains(expected), "{error}"); - assert!(!error.contains("secret-value")); + let error = result.err().expect("startup must fail"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(expected), "{report}"); + assert!(!report.contains("secret-value")); } #[rstest] @@ -313,9 +441,10 @@ fn applies_configured_http_environment_before_building_clients() { .unwrap(); let error = litellm_gateway::build_inference(&config).err().unwrap(); assert!(matches!( - error, - litellm_gateway::Error::Http(litellm_http::Error::Read { path, .. }) - if path == missing + &error, + litellm_gateway::Error::Http(http) + if http.tls_source() == Some(litellm_http::TlsSource::CaBundle) + && http.to_string().contains(&missing.display().to_string()) )); } @@ -340,5 +469,8 @@ async fn refuses_unimplemented_global_mcp_policy(#[case] settings: &str) { .await .err() .unwrap(); - assert!(matches!(error, litellm_gateway::Error::McpSetting(_))); + assert!(matches!( + error, + litellm_gateway::Error::McpSetting(_) | litellm_gateway::Error::McpGuardrails + )); } diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 53a97e7098c..c876c1cbcbb 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -94,7 +94,10 @@ async fn authenticates_before_serving_mounted_inference_routes( let response = request.send().await.unwrap(); assert_eq!(response.status().as_u16(), status); if status == 400 { - let expected = Error::UnknownModel("unconfigured-model".into()); + let expected = Error::UnknownModel { + route: "/v1/messages", + model: "unconfigured-model".into(), + }; assert_eq!( response.json::().await.unwrap(), expected.body(Some("gateway-test")) diff --git a/litellm-rust/crates/host-native/src/driver.rs b/litellm-rust/crates/host-native/src/driver.rs index 7254721da87..af3248848a2 100644 --- a/litellm-rust/crates/host-native/src/driver.rs +++ b/litellm-rust/crates/host-native/src/driver.rs @@ -66,11 +66,10 @@ where MachineStep::Suspended(request) => request, }; let answered = match request { - HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => self - .interceptors - .result_ready(facts) - .await - .map(|()| reply.send(())), + HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => { + reply.send(self.interceptors.result_ready(facts).await); + Ok(()) + } HostRequest::HostCall(call) => self.services.handle_host_call(call).await, HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, @@ -82,10 +81,8 @@ where .await .map(|wire| reply.send(wire)), HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { - self.interceptors - .after_provider_response(raw) - .await - .map(|()| reply.send(())) + reply.send(self.interceptors.after_provider_response(raw).await); + Ok(()) } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { self.demand = Some(reply); diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index 8d220b19def..0b0107d0548 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-coroutine.workspace = true +litellm-error.workspace = true litellm-host.workspace = true bytes.workspace = true diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index 91fbc8c0e52..e36b7d70c28 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -5,7 +5,7 @@ use pyo3::types::PyDict; /// Converts requests, responses, stream values and errors at the Python boundary. pub trait PythonBinding: PythonOwned { - type Protocol: Protocol; + type Protocol: Protocol; /// The public exception a native failure maps to, kept as a value until the driver /// raises it. diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index f8e35aa09ba..dbf843a5295 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -112,7 +112,7 @@ enum Setup { /// The route's public exception for a native failure. When classification itself fails, /// that failure is raised with the native error's text as its `__context__`. fn classified(py: Python<'_>, binding: &H, error: ErrorOf) -> PyErr { - let native = error.to_string(); + let native = litellm_error::Report::new(&error).to_string(); let classifier_error = match binding.map_error(py, error) { Ok(failure) => return failure.into(), Err(classifier_error) => classifier_error, @@ -299,7 +299,7 @@ where facts, })) .await?; - reply.send(()); + reply.send(Ok(())); Ok(()) } HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { @@ -307,7 +307,7 @@ where ExecutionEvent::ProviderResponseReceived { raw: &raw }, )) .await?; - reply.send(()); + reply.send(Ok(())); Ok(()) } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { diff --git a/litellm-rust/crates/host-python/src/error.rs b/litellm-rust/crates/host-python/src/error.rs index b06a4a2fd41..8d47d97df25 100644 --- a/litellm-rust/crates/host-python/src/error.rs +++ b/litellm-rust/crates/host-python/src/error.rs @@ -1,3 +1,4 @@ +use litellm_error::Report; use pyo3::{ PyErr, exceptions::{PyRuntimeError, PyTypeError}, @@ -24,16 +25,19 @@ pub enum Error { SyncInsideTokio, #[error("the context body produced no value")] ContextBodyProducedNothing, + #[error("the context body ran more than once")] + ContextBodyRanTwice, + #[error("the blocking pool dropped a Python call")] + BlockingCallCancelled(#[source] tokio::task::JoinError), #[error("file read must return bytes or str, got {got}")] FileContentType { got: String }, } impl From for PyErr { fn from(error: Error) -> Self { + let message = Report::new(&error).to_string(); match error { - Error::SendBeforeStart | Error::FileContentType { .. } => { - PyTypeError::new_err(error.to_string()) - } + Error::SendBeforeStart | Error::FileContentType { .. } => PyTypeError::new_err(message), Error::MissingState | Error::ExecutionRunning | Error::ExecutionClosed @@ -41,7 +45,9 @@ impl From for PyErr { | Error::ExecutionClosedWhileRunning | Error::NativeExecutionClosed | Error::SyncInsideTokio - | Error::ContextBodyProducedNothing => PyRuntimeError::new_err(error.to_string()), + | Error::ContextBodyProducedNothing + | Error::ContextBodyRanTwice + | Error::BlockingCallCancelled(_) => PyRuntimeError::new_err(message), } } } @@ -51,7 +57,7 @@ impl From for PyErr { /// raised, which reaches the caller as it was raised. #[derive(Debug, thiserror::Error)] pub enum InvokeError { - #[error("{0}")] + #[error(transparent)] Native(E), #[error(transparent)] Python(#[from] PyErr), diff --git a/litellm-rust/crates/host-python/src/fork_gate.rs b/litellm-rust/crates/host-python/src/fork_gate.rs index c4842dd9223..82698931f8f 100644 --- a/litellm-rust/crates/host-python/src/fork_gate.rs +++ b/litellm-rust/crates/host-python/src/fork_gate.rs @@ -12,13 +12,23 @@ pub(crate) struct ForkGate { fork_only_pid: AtomicU32, } -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum Refused { + #[error( + "this process is reserved for forking workers and cannot run native routes; \ + move the call into a worker, after the fork" + )] ReservedForForking, + #[error( + "this process was forked after the native runtime started, and runtime threads \ + do not survive fork(); start workers with spawn or forkserver, or fork before \ + the first native call" + )] ForkedAfterStart, } -#[derive(Debug, PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq, thiserror::Error)] +#[error("the native runtime already started in this process")] pub struct RuntimeAlreadyStarted; impl ForkGate { diff --git a/litellm-rust/crates/host-python/src/gil.rs b/litellm-rust/crates/host-python/src/gil.rs index aacecb3af81..5f54fae81ed 100644 --- a/litellm-rust/crates/host-python/src/gil.rs +++ b/litellm-rust/crates/host-python/src/gil.rs @@ -1,8 +1,9 @@ use std::sync::{ - Arc, Mutex, + Arc, atomic::{AtomicU64, Ordering}, }; +use parking_lot::Mutex; use pyo3::prelude::*; use crate::Error; @@ -36,12 +37,8 @@ impl PythonContext { let value = Arc::clone(&value); pyo3::types::PyCFunction::new_closure(py, None, None, move |_, _| { Python::attach(|py| -> PyResult<()> { - let body = body - .lock() - .expect("context body slot poisoned") - .take() - .expect("the context body ran more than once"); - *value.lock().expect("context value slot poisoned") = Some(body(py)); + let body = body.lock().take().ok_or(Error::ContextBodyRanTwice)?; + *value.lock() = Some(body(py)); Ok(()) }) })? @@ -49,7 +46,6 @@ impl PythonContext { copy.call_method1("run", (callback,))?; value .lock() - .expect("context value slot poisoned") .take() .ok_or_else(|| Error::ContextBodyProducedNothing.into()) } @@ -84,7 +80,7 @@ where match tokio::task::spawn_blocking(move || Python::attach(|py| context.enter(py, f))).await { Ok(value) => value, Err(error) if error.is_panic() => std::panic::resume_unwind(error.into_panic()), - Err(error) => panic!("the blocking pool dropped a Python call: {error}"), + Err(error) => Err(Error::BlockingCallCancelled(error).into()), } } diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index 2361128bf43..cb4de2e4763 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -56,25 +56,27 @@ impl Execution { fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult> { let (mut body, started) = { let mut execution = slf.borrow_mut(); - let started = match (&execution.state, &input) { - (ExecutionState::Created(_), Resume::Send(value)) if value.is_none(py) => false, - (ExecutionState::Created(_), Resume::Send(_)) => { + match ( + std::mem::replace(&mut execution.state, ExecutionState::Running), + &input, + ) { + (ExecutionState::Created(body), Resume::Send(value)) if value.is_none(py) => { + (body, false) + } + (ExecutionState::Suspended(body), _) => (body, true), + (ExecutionState::Running, _) => return Err(Error::ExecutionRunning.into()), + (state @ ExecutionState::Created(_), Resume::Send(_)) => { + execution.state = state; return Err(Error::SendBeforeStart.into()); } - (ExecutionState::Suspended(_), _) => true, - (ExecutionState::Running, _) => { - return Err(Error::ExecutionRunning.into()); - } - (ExecutionState::Closed, _) => { - return Err(Error::ExecutionClosed.into()); - } - (ExecutionState::Created(_), Resume::Throw(_)) => { + (state @ ExecutionState::Created(_), Resume::Throw(_)) => { + execution.state = state; return Err(Error::ExecutionOutOfOrder.into()); } - }; - match std::mem::replace(&mut execution.state, ExecutionState::Running) { - ExecutionState::Created(body) | ExecutionState::Suspended(body) => (body, started), - _ => unreachable!(), + (ExecutionState::Closed, _) => { + execution.state = ExecutionState::Closed; + return Err(Error::ExecutionClosed.into()); + } } }; let result = match input { diff --git a/litellm-rust/crates/host-python/src/marshal.rs b/litellm-rust/crates/host-python/src/marshal.rs index ae9066379ae..465521cf9d2 100644 --- a/litellm-rust/crates/host-python/src/marshal.rs +++ b/litellm-rust/crates/host-python/src/marshal.rs @@ -14,7 +14,8 @@ pub fn from_py_argument(value: &Bound<'_, PyAny>) -> PyResult where T: DeserializeOwned, { - pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string())) + pythonize::depythonize(value) + .map_err(|error| PyValueError::new_err(litellm_error::Report::new(&error).to_string())) } pub fn from_py(value: &Bound<'_, PyAny>) -> PyResult diff --git a/litellm-rust/crates/host-python/src/runtime.rs b/litellm-rust/crates/host-python/src/runtime.rs index c128242ef14..e8872c9340a 100644 --- a/litellm-rust/crates/host-python/src/runtime.rs +++ b/litellm-rust/crates/host-python/src/runtime.rs @@ -7,6 +7,7 @@ use std::time::Duration; use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted}; use crate::{Error, Pythonized, panic_to_pyerr, release_gil}; use futures_util::FutureExt; +use litellm_error::Report; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use serde::Serialize; @@ -42,19 +43,23 @@ pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> { /// Claims process-bound native state before runtime startup or tokenizer execution. pub fn enter_native() -> PyResult<()> { - FORK_GATE - .enter(std::process::id()) - .map_err(|refused| match refused { - Refused::ReservedForForking => ProcessReservedForForking::new_err( - "this process is reserved for forking workers and cannot run native routes; \ - move the call into a worker, after the fork", - ), - Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err( - "this process was forked after the native runtime started, and runtime threads \ - do not survive fork(); start workers with spawn or forkserver, or fork before \ - the first native call", - ), - }) + Ok(FORK_GATE.enter(std::process::id())?) +} + +impl From for PyErr { + fn from(refused: Refused) -> Self { + let message = Report::new(&refused).to_string(); + match refused { + Refused::ReservedForForking => ProcessReservedForForking::new_err(message), + Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err(message), + } + } +} + +impl From for PyErr { + fn from(error: RuntimeAlreadyStarted) -> Self { + PyRuntimeError::new_err(Report::new(&error).to_string()) + } } #[expect(clippy::disallowed_methods, reason = "this is the gated door")] diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index ef09135698e..003fce74309 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -92,6 +92,8 @@ impl std::fmt::Display for Error { } } +impl std::error::Error for Error {} + impl From for Error { fn from(fault: MachineFault) -> Self { Self(format!("{fault:?}")) diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index 51015e99906..d3354fa8622 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -10,6 +10,7 @@ futures-util.workspace = true litellm-auth.workspace = true litellm-coroutine.workspace = true serde_json.workspace = true +thiserror.workspace = true tokio = { workspace = true, features = ["sync"] } [dev-dependencies] diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index f0fd7434c6d..3cff781792f 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -212,7 +212,7 @@ mod tests { panic!("on_event yields Emit"); }; assert_eq!(raw.body, "raw"); - reply.send(()); + reply.send(Ok(())); let Ok(MachineStep::Complete((sent, ()))) = machine.resume().await else { panic!("the call completes with the answers"); diff --git a/litellm-rust/crates/host/src/machine/context.rs b/litellm-rust/crates/host/src/machine/context.rs index b2c9e95c563..1e53605ed7b 100644 --- a/litellm-rust/crates/host/src/machine/context.rs +++ b/litellm-rust/crates/host/src/machine/context.rs @@ -109,7 +109,7 @@ where .request_reply(|reply| { HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) }) - .await + .await? } async fn after_provider_response(&self, raw: RawResponse) -> Result<(), P::Error> { @@ -117,7 +117,7 @@ where .request_reply(|reply| { HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) }) - .await + .await? } } diff --git a/litellm-rust/crates/host/src/machine/coroutine.rs b/litellm-rust/crates/host/src/machine/coroutine.rs index 825536e1d9b..db69f1dceb2 100644 --- a/litellm-rust/crates/host/src/machine/coroutine.rs +++ b/litellm-rust/crates/host/src/machine/coroutine.rs @@ -10,12 +10,14 @@ use super::{ use crate::protocol::{HostRequest, Protocol}; /// The machine's own failures, distinct from anything the provider call reports. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub enum MachineFault { /// The host dropped an op's reply unanswered, or went away while the call waited. + #[error("the host abandoned the call")] Abandoned, /// The host resumed the call out of turn. - Protocol(ResumeError), + #[error("the host resumed the call out of turn")] + Protocol(#[source] ResumeError), } pub type ExecuteFuture::Response> = diff --git a/litellm-rust/crates/host/src/protocol.rs b/litellm-rust/crates/host/src/protocol.rs index 9e5aff8aae8..c24557b0348 100644 --- a/litellm-rust/crates/host/src/protocol.rs +++ b/litellm-rust/crates/host/src/protocol.rs @@ -15,14 +15,14 @@ pub trait Protocol: Send + Sync + 'static { pub enum HostRequest { HostCall(P::HostCall), - Intercept(InterceptRequest), + Intercept(InterceptRequest), Stream(StreamDelivery

), } -pub enum InterceptRequest { +pub enum InterceptRequest { ResultReady { facts: crate::interceptors::ExecutionFacts, - reply: Reply<()>, + reply: Reply>, }, BeforeProviderRequest { wire: Box, @@ -31,7 +31,7 @@ pub enum InterceptRequest { }, AfterProviderResponse { raw: RawResponse, - reply: Reply<()>, + reply: Reply>, }, } diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 7d64467f2f7..6b3e3cc7852 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -7,13 +7,11 @@ repository.workspace = true [features] test-support = [] -mcp = ["dep:rmcp", "dep:reqwest-mcp"] [dependencies] -rmcp = { version = "=3.4.1", default-features = false, features = ["client", "transport-streamable-http-client-reqwest"], optional = true } -reqwest-mcp = { package = "reqwest", version = "0.13.2", default-features = false, features = ["rustls-no-provider", "json", "stream", "http2"], optional = true } http.workspace = true litellm-core-utils.workspace = true +litellm-error.workspace = true hyper-util.workspace = true reqwest.workspace = true rustls.workspace = true @@ -25,6 +23,7 @@ strum.workspace = true thiserror.workspace = true tokio.workspace = true tracing.workspace = true +url.workspace = true veil.workspace = true webpki-roots.workspace = true @@ -32,5 +31,6 @@ webpki-roots.workspace = true futures-util.workspace = true rcgen = "0.14.10" tempfile.workspace = true +tokio-rustls = { version = "0.26", default-features = false, features = ["ring"] } rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/http/src/client.rs b/litellm-rust/crates/http/src/client.rs index 1f7017d083b..bb7883f665b 100644 --- a/litellm-rust/crates/http/src/client.rs +++ b/litellm-rust/crates/http/src/client.rs @@ -9,17 +9,7 @@ impl Client { } #[cfg(any(test, feature = "test-support"))] - pub fn for_test(client: reqwest::Client) -> Self { - Self(client) - } - - #[cfg(any(test, feature = "test-support"))] - pub fn plain_for_test() -> Self { - Self(reqwest::Client::new()) - } - - #[cfg(any(test, feature = "test-support"))] - pub fn no_redirect_for_test() -> Self { + pub fn for_test() -> Self { Self( reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) @@ -27,6 +17,11 @@ impl Client { .expect("a client without TLS or proxy settings builds"), ) } + + #[cfg(any(test, feature = "test-support"))] + pub fn wrap_for_test(client: reqwest::Client) -> Self { + Self(client) + } } impl Deref for Client { diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 2f36784bc70..5e5f2d66fa2 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -16,6 +16,7 @@ pub enum Verify { Disabled, CaBundle(PathBuf), BuiltInRoots, + NativeRoots, } #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -98,7 +99,7 @@ impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder { fn try_from(config: &HttpClientConfig) -> Result { let base = reqwest::Client::builder() - .use_preconfigured_tls(rustls::ClientConfig::try_from(config)?) + .tls_backend_preconfigured(rustls::ClientConfig::try_from(config)?) .connect_timeout(config.connect_timeout) .pool_idle_timeout(config.pool_idle_timeout); let with_keepalive = match config.tcp_keepalive { @@ -135,7 +136,7 @@ mod tests { use rstest::rstest; use super::*; - use crate::TlsSource; + use crate::{TlsSource, error::Repr}; fn settings(ssl_verify: Option, ssl_cert_file: Option<&str>) -> HttpSettings { HttpSettings { @@ -305,11 +306,10 @@ mod tests { }; assert!(matches!( reqwest::ClientBuilder::try_from(&config), - Err(Error::Read { - path: reported, - tls_source: TlsSource::CaBundle, - .. - }) if reported == path + Err(error) if matches!( + error.repr(), + Repr::Read { path: reported, tls_source: TlsSource::CaBundle, .. } if *reported == path + ) )); } @@ -326,11 +326,10 @@ mod tests { std::fs::remove_file(&path).unwrap(); assert!(matches!( result, - Err(Error::InvalidPem { - path: reported, - tls_source: TlsSource::CaBundle, - .. - }) if reported == path + Err(error) if matches!( + error.repr(), + Repr::InvalidPem { path: reported, tls_source: TlsSource::CaBundle, .. } if *reported == path + ) )); } } diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index eafb4d2976b..29c8b988c0b 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -1,4 +1,10 @@ -use std::path::PathBuf; +use std::{borrow::Cow, fmt, io, path::PathBuf, sync::Arc}; + +use tokio::time::error::Elapsed; +use tokio_tungstenite::tungstenite; + +const UPSTREAM_BODY_MAX_CHARS: usize = 256; +const TRUNCATION_MARKER: &str = "... (truncated)"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum TlsSource { @@ -6,32 +12,666 @@ pub enum TlsSource { ClientIdentity, } -#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] -pub enum Error { - #[error("could not read {}: {message}", path.display())] +#[derive(Clone, Debug, thiserror::Error)] +#[error(transparent)] +pub struct Error(Arc); + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Repr { + #[error("could not read {}", path.display())] Read { path: PathBuf, - message: String, tls_source: TlsSource, + source: io::Error, }, - #[error("{} is not a PEM file: {message}", path.display())] + #[error("could not load {}", path.display())] InvalidPem { path: PathBuf, - message: String, tls_source: TlsSource, + source: PemFault, }, - #[error("could not build the HTTP client: {0}")] - Client(String), - #[error("request body could not be serialized: {0}")] - RequestBody(String), - #[error("request forwards a header the signer computes: {0}")] + #[error("could not build the HTTP client")] + Client(#[source] reqwest::Error), + #[error("could not build the TLS configuration")] + TlsConfig(#[source] rustls::Error), + + #[error("request body could not be serialized")] + RequestBody(#[source] serde_json::Error), + #[error("invalid request")] + Request(#[source] http::Error), + #[error("invalid request")] + RequestUrl(#[source] reqwest::Error), + #[error("invalid request URL")] + ParseUrl(#[source] url::ParseError), + #[error("request forwards the header {0:?}, which the signer computes")] ComputedHeader(String), - #[error("request signing failed: {0}")] - Signature(String), + #[error("request signing failed: a streaming body cannot be signed")] + UnsignableBody, + #[error("request forwards a header that cannot be signed")] + UnsignableHeader(#[source] litellm_error::Source), + #[error("request signing failed")] + Signature(#[source] litellm_error::Source), + + #[error("could not reach the provider")] + Connect(#[source] reqwest::Error), + #[error("upstream network error")] + Network(#[source] reqwest::Error), + #[error("upstream request timed out")] + Timeout(#[source] reqwest::Error), + #[error("upstream request failed with status {status}: {body}")] + Status { status: u16, body: UpstreamBody }, + + #[error("invalid WebSocket URL")] + WebSocketUrl(#[source] url::ParseError), + #[error("unsupported WebSocket URL scheme {0:?}")] + WebSocketScheme(String), + #[error("could not rewrite the WebSocket URL scheme to {0}")] + WebSocketSchemeRewrite(&'static str), + #[error("WebSocket handshake returned {0}")] + Handshake(HandshakeFault), + #[error("WebSocket connection timed out")] + HandshakeTimeout(#[source] Elapsed), + #[error("upstream WebSocket error")] + WebSocket(#[source] tungstenite::Error), + #[error("the WebSocket connection is closed")] + WebSocketClosed, + + #[error("media URL rejected by network policy")] + BlockedUrl, + #[error("could not resolve the media host")] + Resolve(#[source] io::Error), + #[error("media download is disabled")] + DownloadDisabled, + #[error("media download exceeds the maximum size")] + DownloadTooLarge, + #[error("too many redirects while fetching media")] + TooManyRedirects, + #[error("media redirect is missing a Location header")] + MissingRedirectLocation, + #[error("invalid media redirect")] + InvalidRedirect(#[source] RedirectFault), + #[error("media download failed with status {0}")] + DownloadStatus(u16), + #[error("media download failed")] + Download(#[source] reqwest::Error), + #[error("media download timed out")] + DownloadTimeout(#[source] Elapsed), +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum HandshakeFault { + #[error("a wrong Sec-WebSocket-Accept")] + Accept, + #[error("an invalid Upgrade header")] + Upgrade, + #[error("an invalid Connection header")] + Connection, + #[error("an invalid subprotocol")] + Subprotocol, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum PemFault { + #[error("no certificates found")] + NoCertificates, + #[error("could not parse PEM")] + Parse(#[from] rustls::pki_types::pem::Error), + #[error("rustls rejected the certificate or key")] + Rejected(#[from] rustls::Error), +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum RedirectFault { + #[error("Location header is not valid text")] + Location(#[from] http::header::ToStrError), + #[error("Location is not a valid URL")] + Url(#[from] url::ParseError), +} + +impl Repr { + pub(crate) fn download(error: reqwest::Error) -> Self { + Self::Download(error.without_url()) + } +} + +impl Error { + pub fn upstream(status: u16, body: String) -> Self { + Repr::Status { + status, + body: UpstreamBody(body), + } + .into() + } + + pub fn computed_header(name: impl Into) -> Self { + Repr::ComputedHeader(name.into()).into() + } + + pub fn unsignable_body() -> Self { + Repr::UnsignableBody.into() + } + + pub fn unsignable_header(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Repr::UnsignableHeader(litellm_error::Source::new(cause)).into() + } + + pub fn websocket_closed() -> Self { + Repr::WebSocketClosed.into() + } + + pub fn signature(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Repr::Signature(litellm_error::Source::new(cause)).into() + } + + /// The request may have reached the provider, so retrying elsewhere could + /// repeat a billed call. + pub fn dispatched(&self) -> bool { + match &*self.0 { + Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) => true, + Repr::WebSocketClosed + | Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_request(&self) -> bool { + match &*self.0 { + Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableHeader(_) + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::BlockedUrl + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::UnsignableBody + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::Resolve(_) + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_response(&self) -> bool { + match &*self.0 { + Repr::MissingRedirectLocation | Repr::InvalidRedirect(_) => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_computed_header(&self) -> bool { + match &*self.0 { + Repr::ComputedHeader(_) => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_connect(&self) -> bool { + match &*self.0 { + Repr::Connect(_) => true, + Repr::ComputedHeader(_) + | Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => false, + } + } + + pub fn is_timeout(&self) -> bool { + match &*self.0 { + Repr::Timeout(_) | Repr::HandshakeTimeout(_) | Repr::DownloadTimeout(_) => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) => false, + } + } + + pub fn upstream_response(&self) -> Option<(u16, &str)> { + match &*self.0 { + Repr::Status { status, body } => Some((*status, &body.0)), + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => None, + } + } + + pub fn status(&self) -> Option { + match &*self.0 { + Repr::Status { status, .. } | Repr::DownloadStatus(status) => Some(*status), + Repr::Timeout(_) | Repr::DownloadTimeout(_) => Some(408), + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::Download(_) => None, + } + } + + pub fn tls_source(&self) -> Option { + match &*self.0 { + Repr::Read { tls_source, .. } | Repr::InvalidPem { tls_source, .. } => { + Some(*tls_source) + } + Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ParseUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Timeout(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::HandshakeTimeout(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) + | Repr::DownloadTimeout(_) => None, + } + } + + #[cfg(test)] + pub(crate) fn repr(&self) -> &Repr { + &self.0 + } +} + +pub(crate) struct UpstreamBody(String); + +impl fmt::Display for UpstreamBody { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&bounded(&self.0), f) + } +} + +impl fmt::Debug for UpstreamBody { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(&bounded(&self.0), f) + } +} + +fn bounded(body: &str) -> Cow<'_, str> { + match body.char_indices().nth(UPSTREAM_BODY_MAX_CHARS) { + None => Cow::Borrowed(body), + Some((cut, _)) => Cow::Owned(format!("{}{TRUNCATION_MARKER}", &body[..cut])), + } +} + +impl PartialEq for Error { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +impl Eq for Error {} + +impl From for Error { + fn from(repr: Repr) -> Self { + Self(Arc::new(repr)) + } } impl From for Error { fn from(error: reqwest::Error) -> Self { - Self::Client(error.without_url().to_string()) + let error = error.without_url(); + if error.is_timeout() { + Repr::Timeout(error).into() + } else if error.is_connect() || error.is_builder() { + Repr::Connect(error).into() + } else { + Repr::Network(error).into() + } + } +} + +impl From for Error { + fn from(error: tungstenite::Error) -> Self { + Repr::WebSocket(error).into() + } +} + +#[cfg(test)] +mod tests { + use litellm_error::Report; + use rstest::rstest; + + use super::{TRUNCATION_MARKER, UPSTREAM_BODY_MAX_CHARS}; + + #[tokio::test] + async fn transport_errors_remove_urls_and_keep_dispatch_context() { + let error = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get("http://localhost:invalid/private?api_key=secret") + .send() + .await + .expect_err("invalid port"); + let error = crate::Error::from(error); + let report = Report::new(&error).to_string(); + assert!(!error.dispatched()); + assert!(!report.contains("secret"), "{report}"); + assert!(!report.contains("private"), "{report}"); + } + + fn root_cause<'a>( + error: &'a (dyn std::error::Error + 'static), + ) -> &'a (dyn std::error::Error + 'static) { + error.source().map_or(error, root_cause) + } + + #[tokio::test] + async fn connect_failure_reports_its_root_cause_once() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let address = listener.local_addr().expect("address"); + drop(listener); + let error = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}/private?api_key=secret")) + .send() + .await + .expect_err("nothing listens on the port"); + let error = crate::Error::from(error); + let root_cause = root_cause(&error).to_string(); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches(&root_cause).count(), 1, "{report}"); + assert!(!error.to_string().contains(&root_cause)); + assert!(error.is_connect()); + assert!(!error.is_timeout()); + assert!(!report.contains("secret"), "{report}"); + assert!(!report.contains("private"), "{report}"); + } + + #[tokio::test] + async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() { + use std::time::Duration; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let address = listener.local_addr().expect("address"); + let request = reqwest::Client::builder() + .no_proxy() + .build() + .expect("client") + .get(format!("http://{address}")) + .timeout(Duration::from_millis(200)) + .send(); + let (response, accepted) = tokio::join!( + request, + tokio::time::timeout(Duration::from_secs(2), listener.accept()) + ); + let _connection = accepted + .expect("accept deadline") + .expect("accepted connection"); + let error = response.expect_err("server does not respond"); + assert!(error.is_timeout()); + let error = crate::Error::from(error); + assert!(error.dispatched()); + assert!(error.is_timeout()); + assert!(!error.is_connect()); + assert_eq!(error.status(), Some(408)); + } + + #[rstest] + #[case::empty(String::new(), String::new())] + #[case::short("slow down".to_owned(), "slow down".to_owned())] + #[case::ascii_over_the_limit( + format!("{}private-tail", "a".repeat(UPSTREAM_BODY_MAX_CHARS)), + format!("{}{TRUNCATION_MARKER}", "a".repeat(UPSTREAM_BODY_MAX_CHARS)) + )] + #[case::at_the_limit("\u{00e9}".repeat(UPSTREAM_BODY_MAX_CHARS), "\u{00e9}".repeat(UPSTREAM_BODY_MAX_CHARS))] + #[case::over_the_limit( + format!("{}private-tail", "\u{00e9}".repeat(UPSTREAM_BODY_MAX_CHARS)), + format!("{}{TRUNCATION_MARKER}", "\u{00e9}".repeat(UPSTREAM_BODY_MAX_CHARS)) + )] + fn upstream_preserves_the_body_and_bounds_only_diagnostics( + #[case] body: String, + #[case] expected: String, + ) { + let error = crate::Error::upstream(429, body.clone()); + let expected_message = format!("upstream request failed with status 429: {expected}"); + assert_eq!(error.to_string(), expected_message); + assert_eq!(Report::new(&error).to_string(), expected_message); + let debug = format!("{error:?}"); + assert!(debug.contains(&format!("{expected:?}")), "{debug}"); + assert!(!debug.contains("private-tail"), "{debug}"); + assert!(!format!("{error:#?}").contains("private-tail")); + assert_eq!(error.upstream_response(), Some((429, body.as_str()))); + assert_eq!(error.status(), Some(429)); } } diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index d53da607d39..02129c62083 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -7,8 +7,6 @@ mod client; mod config; mod error; -#[cfg(feature = "mcp")] -mod mcp; pub mod media; pub mod outbound; mod pool; @@ -16,7 +14,6 @@ mod proxy; pub mod request; mod settings; mod tls; -pub mod transport; pub mod websocket; pub use client::Client; diff --git a/litellm-rust/crates/http/src/mcp.rs b/litellm-rust/crates/http/src/mcp.rs deleted file mode 100644 index 2b101fb9c98..00000000000 --- a/litellm-rust/crates/http/src/mcp.rs +++ /dev/null @@ -1,76 +0,0 @@ -use std::{ - collections::HashMap, - net::{IpAddr, Ipv4Addr}, - sync::{Mutex, PoisonError}, - time::{Duration, Instant}, -}; - -use crate::{Error, HttpClientConfig}; - -struct Entry { - client: reqwest_mcp::Client, - built_at: Instant, -} - -#[derive(Default)] -pub(super) struct Pool(Mutex>); - -impl Pool { - pub fn client( - &self, - config: &HttpClientConfig, - ttl: Duration, - ) -> Result { - let mut clients = self.0.lock().unwrap_or_else(PoisonError::into_inner); - if let Some(entry) = clients.get(config) - && entry.built_at.elapsed() < ttl - { - return Ok(entry.client.clone()); - } - let client = build(config)?; - clients.insert( - config.clone(), - Entry { - client: client.clone(), - built_at: Instant::now(), - }, - ); - Ok(client) - } -} - -fn build(config: &HttpClientConfig) -> Result { - let base = reqwest_mcp::Client::builder() - .tls_backend_preconfigured(rustls::ClientConfig::try_from(config)?) - .connect_timeout(config.connect_timeout) - .pool_idle_timeout(config.pool_idle_timeout) - .redirect(reqwest_mcp::redirect::Policy::none()); - let keepalive = match config.tcp_keepalive { - Some(value) => base - .tcp_keepalive(value.idle) - .tcp_keepalive_interval(value.interval) - .tcp_keepalive_retries(value.retries), - None => base, - }; - let address = if config.force_ipv4 { - keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED)) - } else { - keepalive - }; - let protocol = if config.http2 { - address - } else { - address.http1_only() - }; - let agent = match &config.user_agent { - Some(value) => protocol.user_agent(value), - None => protocol, - }; - config - .proxies - .mcp_proxies() - .into_iter() - .fold(agent.no_proxy(), reqwest_mcp::ClientBuilder::proxy) - .build() - .map_err(|error| Error::Client(error.without_url().to_string())) -} diff --git a/litellm-rust/crates/http/src/media.rs b/litellm-rust/crates/http/src/media.rs index 22a881831cf..4acaca5e502 100644 --- a/litellm-rust/crates/http/src/media.rs +++ b/litellm-rust/crates/http/src/media.rs @@ -12,29 +12,7 @@ use reqwest::{ dns::{Addrs, Name, Resolve, Resolving}, }; -use crate::{Client, ClientVariant, HttpClientConfig, HttpClientPool}; - -#[derive(Debug, thiserror::Error)] -pub enum Error { - #[error("media URL rejected by network policy")] - BlockedUrl, - #[error("media download is disabled")] - DownloadDisabled, - #[error("media download exceeds the maximum size")] - DownloadTooLarge, - #[error("too many redirects while fetching media")] - TooManyRedirects, - #[error("media redirect is missing a Location header")] - MissingRedirectLocation, - #[error("invalid media redirect")] - InvalidRedirect, - #[error("media download failed with status {0}")] - Http(u16), - #[error("media download timed out")] - Timeout, - #[error("{0}")] - Transport(#[from] crate::transport::Error), -} +use crate::{Client, ClientVariant, Error, HttpClientConfig, HttpClientPool, error::Repr}; #[derive(Clone, Debug, PartialEq, Eq)] pub struct UrlPolicy { @@ -125,7 +103,7 @@ impl MediaFetcher { pool: &HttpClientPool, config: &HttpClientConfig, url_policy: UrlPolicy, - ) -> Result { + ) -> Result { let uses_proxy: ProxyMatch = Arc::new(config.proxies.matcher()); Self::with_resolution( pool, @@ -142,7 +120,7 @@ impl MediaFetcher { url_policy: UrlPolicy, address_resolver: Arc, uses_proxy: ProxyMatch, - ) -> Result { + ) -> Result { Ok(Self { pinned: pool.client(config, ClientVariant::Media)?, unpinned: pool.client(config, ClientVariant::UnpinnedMedia)?, @@ -167,11 +145,11 @@ impl MediaFetcher { pub async fn fetch(&self, url: Url, policy: DownloadPolicy) -> Result { if policy.max_bytes == 0 { - return Err(Error::DownloadDisabled); + return Err(Repr::DownloadDisabled.into()); } tokio::time::timeout(policy.timeout, self.fetch_before_deadline(url, policy)) .await - .map_err(|_| Error::Timeout)? + .map_err(Repr::DownloadTimeout)? } async fn fetch_before_deadline( @@ -187,22 +165,25 @@ impl MediaFetcher { .get(url.clone()) .send() .await - .map_err(crate::transport::Error::from)?; + .map_err(Repr::download)?; if response.status().is_redirection() { if redirects_followed == policy.max_redirects { - return Err(Error::TooManyRedirects); + return Err(Repr::TooManyRedirects.into()); } let location = response .headers() .get(reqwest::header::LOCATION) - .and_then(|value| value.to_str().ok()) - .ok_or(Error::MissingRedirectLocation)?; - url = url.join(location).map_err(|_| Error::InvalidRedirect)?; + .ok_or(Repr::MissingRedirectLocation)? + .to_str() + .map_err(|error| Repr::InvalidRedirect(error.into()))?; + url = url + .join(location) + .map_err(|error| Repr::InvalidRedirect(error.into()))?; redirects_followed += 1; continue; } if !response.status().is_success() { - return Err(Error::Http(response.status().as_u16())); + return Err(Repr::DownloadStatus(response.status().as_u16()).into()); } enforce_download_size(response.content_length().unwrap_or(0), policy.max_bytes)?; let content_type = response @@ -215,11 +196,7 @@ impl MediaFetcher { .unwrap_or("application/octet-stream") .to_string(); let mut bytes = Vec::new(); - while let Some(chunk) = response - .chunk() - .await - .map_err(crate::transport::Error::from)? - { + while let Some(chunk) = response.chunk().await.map_err(Repr::download)? { enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?; bytes.extend_from_slice(&chunk); } @@ -238,13 +215,13 @@ impl MediaFetcher { || !url.username().is_empty() || url.password().is_some() { - return Err(Error::BlockedUrl); + return Err(Repr::BlockedUrl.into()); } - let host = url.host_str().ok_or(Error::BlockedUrl)?; + let host = url.host_str().ok_or(Repr::BlockedUrl)?; if self.allow_private_network { return Ok(&self.pinned); } - let port = url.port_or_known_default().ok_or(Error::BlockedUrl)?; + let port = url.port_or_known_default().ok_or(Repr::BlockedUrl)?; if self.url_policy.allows(host, port) { return Ok(&self.unpinned); } @@ -262,27 +239,29 @@ impl MediaFetcher { .trim_end_matches(']') .parse::() { - return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl); + return (!is_blocked_ip(ip)) + .then_some(()) + .ok_or(Repr::BlockedUrl.into()); } let addresses = self .address_resolver .resolve(host, port) .await - .map_err(|error| crate::transport::Error::Network(error.to_string()))?; + .map_err(Repr::Resolve)?; validate_addresses(&addresses) } } fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), Error> { if length > max_bytes { - return Err(Error::DownloadTooLarge); + return Err(Repr::DownloadTooLarge.into()); } Ok(()) } fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), Error> { if addresses.is_empty() || addresses.iter().any(|address| is_blocked_ip(address.ip())) { - return Err(Error::BlockedUrl); + return Err(Repr::BlockedUrl.into()); } Ok(()) } @@ -356,10 +335,8 @@ impl Resolve for PublicDnsResolver { .await .map_err(|error| Box::new(error) as Box)? .collect::>(); - validate_addresses(&addresses).map_err(|_| { - Box::new(io::Error::other("destination rejected by network policy")) - as Box - })?; + validate_addresses(&addresses) + .map_err(|error| Box::new(error) as Box)?; Ok(Box::new(addresses.into_iter()) as Addrs) }) } @@ -522,7 +499,7 @@ mod tests { b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", ) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let media = MediaFetcher::for_test(client) .fetch(url, policy(3, 0)) .await @@ -538,13 +515,13 @@ mod tests { b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc", ) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let error = MediaFetcher::for_test(client) .fetch(url, policy(2, 0)) .await .expect_err("oversize body is rejected"); server.await.expect("server completes"); - assert!(matches!(error, Error::DownloadTooLarge)); + assert!(matches!(error.repr(), Repr::DownloadTooLarge)); } #[tokio::test] @@ -553,13 +530,13 @@ mod tests { b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n", ) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let error = MediaFetcher::for_test(client) .fetch(url, policy(3, 0)) .await .expect_err("stream crossing limit is rejected"); server.await.expect("server completes"); - assert!(matches!(error, Error::DownloadTooLarge)); + assert!(matches!(error.repr(), Repr::DownloadTooLarge)); } #[tokio::test] @@ -595,7 +572,7 @@ mod tests { .expect_err("private redirect is rejected"); let requests = server.await.expect("server completes"); assert_eq!(requests.len(), 1); - assert!(matches!(error, Error::BlockedUrl)); + assert!(matches!(error.repr(), Repr::BlockedUrl)); } #[tokio::test] @@ -622,7 +599,27 @@ mod tests { .await .expect_err("fetch times out"); server.await.expect("server completes"); - assert!(matches!(error, Error::Timeout)); + assert!(matches!(error.repr(), Repr::DownloadTimeout(_))); + assert!(!error.dispatched()); + } + + #[rstest::rstest] + #[case::closed_before_answering(b"")] + #[case::reset_mid_body(b"HTTP/1.1 200 OK\r\nContent-Length: 10\r\n\r\nab")] + #[tokio::test] + async fn a_failed_download_never_reached_the_provider(#[case] response: &'static [u8]) { + let (url, server) = serve(response).await; + let url = url.join("?sig=secret").expect("valid test URL"); + let error = MediaFetcher::for_test(Client::for_test()) + .fetch(url, policy(10, 0)) + .await + .expect_err("download fails"); + server.await.expect("server completes"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(!error.dispatched()); + assert!(!error.is_request()); + assert!(!report.contains("secret"), "{report}"); + assert!(!report.contains("document"), "{report}"); } #[tokio::test] @@ -652,8 +649,13 @@ mod tests { let url = Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses"); assert!(matches!( - fetcher.fetch(url, policy(1, 0)).await, - Err(Error::BlockedUrl) + fetcher + .fetch(url, policy(1, 0)) + .await + .err() + .as_ref() + .map(Error::repr), + Some(Repr::BlockedUrl) )); } @@ -683,7 +685,10 @@ mod tests { let result = fetcher(UNROUTABLE, HashSet::from(["localhost"]), other_port, false) .fetch(url, policy(2, 0)) .await; - assert!(matches!(result, Err(Error::BlockedUrl))); + assert!(matches!( + result.as_ref().err().map(Error::repr), + Some(Repr::BlockedUrl) + )); } #[test] @@ -742,6 +747,9 @@ mod tests { ) .fetch(url, policy(2, 0)) .await; - assert!(matches!(blocked, Err(Error::BlockedUrl))); + assert!(matches!( + blocked.as_ref().err().map(Error::repr), + Some(Repr::BlockedUrl) + )); } } diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs index c2cfb00d79b..20b0b1be972 100644 --- a/litellm-rust/crates/http/src/outbound.rs +++ b/litellm-rust/crates/http/src/outbound.rs @@ -1,131 +1,67 @@ -//! The request a route hands to the transport. The body is serialized once, -//! when the request is built, and a [`RequestSigner`] sees those exact bytes. +//! Building the request a route hands to the transport. The body is serialized +//! once, into the [`reqwest::Request`] that is sent, and a [`RequestSigner`] +//! reads those exact bytes back out of it. //! //! Host hooks may rewrite the wire request (redaction, guardrails) and a //! signature such as AWS SigV4 covers the body, so a route builds this after -//! its hooks ran and cannot change or re-serialize it afterwards. +//! its hooks ran and signs last. use std::time::Duration; +use reqwest::{ + Request, Url, + header::{CONTENT_TYPE, HeaderMap}, +}; use serde::Serialize; -use crate::{ - Error, - request::{HeaderPolicy, has_header, with_headers}, -}; - -#[derive(Clone, Copy, Debug)] -pub struct UnsignedRequest<'a> { - pub url: &'a str, - pub headers: &'a [(String, String)], - pub body: &'a [u8], -} +use crate::{Error, error::Repr}; /// Returns the headers to add to the request; it never sees a mutable request. pub trait RequestSigner: Send + Sync { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error>; + fn sign(&self, request: &Request) -> Result; } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct OutboundRequest { - url: String, - headers: Vec<(String, String)>, - body: Vec, +pub fn json( + url: &str, + headers: &[(String, String)], + body: &impl Serialize, timeout: Option, +) -> Result { + let body = serde_json::to_vec(body).map_err(Repr::RequestBody)?; + let url = Url::parse(url).map_err(Repr::ParseUrl)?; + let with_headers = headers.iter().fold( + http::Request::post(url.as_str()), + |builder, (name, value)| builder.header(name, value), + ); + let has_content_type = with_headers + .headers_ref() + .is_some_and(|headers| headers.contains_key(CONTENT_TYPE)); + let builder = if has_content_type { + with_headers + } else { + with_headers.header(CONTENT_TYPE, "application/json") + }; + let mut request = from_http(builder.body(body).map_err(Repr::Request)?)?; + *request.timeout_mut() = timeout; + Ok(request) } -impl OutboundRequest { - pub fn json( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - ) -> Result { - Self::build(url, headers, body, timeout, None) - } +pub fn from_http>(request: http::Request) -> Result { + Request::try_from(request).map_err(|error| Repr::RequestUrl(error.without_url()).into()) +} - pub fn signed_json( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - signer: &dyn RequestSigner, - ) -> Result { - Self::build(url, headers, body, timeout, Some(signer)) - } - - fn build( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - signer: Option<&dyn RequestSigner>, - ) -> Result { - let body = - serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?; - let content_type = (!has_header(&headers, "content-type")) - .then(|| ("content-type".to_string(), "application/json".to_string())); - let unsigned: Vec<(String, String)> = headers.into_iter().chain(content_type).collect(); - let signature = signer - .map(|signer| { - signer.sign(UnsignedRequest { - url: &url, - headers: &unsigned, - body: &body, - }) - }) - .transpose()? - .unwrap_or_default(); - Ok(Self { - url, - headers: unsigned.into_iter().chain(signature).collect(), - body, - timeout, - }) - } - - pub fn url(&self) -> &str { - &self.url - } - - pub fn headers(&self) -> &[(String, String)] { - &self.headers - } - - pub fn header(&self, name: &str) -> Option<&str> { - self.headers - .iter() - .find(|(key, _)| key.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.as_str()) - } - - pub fn body(&self) -> &[u8] { - &self.body - } - - pub fn timeout(&self) -> Option { - self.timeout - } - - pub async fn send(self, client: &crate::Client) -> Result { - let builder = with_headers( - client.post(&self.url).body(self.body), - &self.headers, - HeaderPolicy::All, - ); - match self.timeout { - Some(timeout) => builder.timeout(timeout), - None => builder, - } - .send() - .await - } +pub fn signed(mut request: Request, signer: &dyn RequestSigner) -> Result { + let signature = signer.sign(&request)?; + request.headers_mut().extend(signature); + Ok(request) } #[cfg(test)] mod tests { use std::sync::Mutex; + use reqwest::header::HeaderValue; + use rstest::rstest; use serde_json::json; use super::*; @@ -134,61 +70,73 @@ mod tests { struct Recording(Mutex>); impl RequestSigner for Recording { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { - *self.0.lock().unwrap() = request.body.to_vec(); - Ok(vec![("authorization".into(), "signed".into())]) + fn sign(&self, request: &Request) -> Result { + *self.0.lock().unwrap() = request.body().unwrap().as_bytes().unwrap().to_vec(); + Ok(HeaderMap::from_iter([( + reqwest::header::AUTHORIZATION, + HeaderValue::from_static("signed"), + )])) } } #[test] fn the_signer_sees_exactly_the_bytes_that_are_sent() { let signer = Recording::default(); - let request = OutboundRequest::signed_json( - "https://provider.test/".into(), - vec![("x-caller".into(), "kept".into())], - &json!({"b": 1, "a": [true, null]}), - None, + let request = signed( + json( + "https://provider.test/", + &[("x-caller".into(), "kept".into())], + &json!({"b": 1, "a": [true, null]}), + Some(Duration::from_secs(7)), + ) + .unwrap(), &signer, ) .unwrap(); - assert_eq!(request.body(), signer.0.lock().unwrap().as_slice()); - assert_eq!(request.header("authorization"), Some("signed")); - assert_eq!(request.header("x-caller"), Some("kept")); + assert_eq!( + request.body().unwrap().as_bytes().unwrap(), + signer.0.lock().unwrap().as_slice() + ); + assert_eq!(request.headers()["authorization"], "signed"); + assert_eq!(request.headers()["x-caller"], "kept"); + assert_eq!(request.timeout(), Some(&Duration::from_secs(7))); } #[test] fn the_content_type_is_part_of_what_the_signer_sees() { struct RequiresContentType; impl RequestSigner for RequiresContentType { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { - has_header(request.headers, "content-type") - .then(Vec::new) - .ok_or_else(|| Error::Signature("content-type was not signed".into())) + fn sign(&self, request: &Request) -> Result { + request + .headers() + .contains_key(CONTENT_TYPE) + .then(HeaderMap::new) + .ok_or_else(|| Error::signature(std::fmt::Error)) } } - let defaulted = OutboundRequest::signed_json( - "u".into(), - Vec::new(), - &json!({}), - None, + let defaulted = signed( + json("https://u.test/", &[], &json!({}), None).unwrap(), &RequiresContentType, ) .unwrap(); - assert_eq!(defaulted.header("content-type"), Some("application/json")); + assert_eq!(defaulted.headers()[CONTENT_TYPE], "application/json"); - let provider = OutboundRequest::signed_json( - "u".into(), - vec![("Content-Type".into(), "application/x-amz-json-1.1".into())], - &json!({}), - None, + let provider = signed( + json( + "https://u.test/", + &[("Content-Type".into(), "application/x-amz-json-1.1".into())], + &json!({}), + None, + ) + .unwrap(), &RequiresContentType, ) .unwrap(); assert_eq!( - provider.header("content-type"), - Some("application/x-amz-json-1.1") + provider.headers()[CONTENT_TYPE], + "application/x-amz-json-1.1" ); assert_eq!(provider.headers().len(), 1); } @@ -197,14 +145,104 @@ mod tests { fn a_signer_failure_produces_no_request() { struct Refuses; impl RequestSigner for Refuses { - fn sign(&self, _request: UnsignedRequest<'_>) -> Result, Error> { - Err(Error::ComputedHeader("authorization".into())) + fn sign(&self, _request: &Request) -> Result { + Err(Error::computed_header("authorization")) } } + let error = signed( + json("https://u.test/", &[], &json!({}), None).unwrap(), + &Refuses, + ) + .unwrap_err(); + assert!(error.is_computed_header()); + } + + #[rstest] + #[case::unicode_domain("https://例え.test/v1/chat/completions")] + #[case::unicode_path("https://provider.test/模型/v1/chat/completions")] + #[case::unicode_query("https://provider.test/v1/chat/completions?region=日本")] + #[case::ascii("https://provider.test/v1/chat/completions?region=us-east-1")] + #[case::escaped("https://provider.test/%E6%A8%A1%E5%9E%8B?q=%E6%97%A5%E6%9C%AC")] + fn provider_urls_keep_reqwest_normalization(#[case] url: &str) { + let body = json!({"text": "日本"}); + let timeout = Duration::from_secs(7); + let expected = reqwest::Client::new() + .post(url) + .header(CONTENT_TYPE, "application/json") + .header("x-caller", "kept") + .body(serde_json::to_vec(&body).unwrap()) + .timeout(timeout) + .build() + .unwrap(); + let signer = Recording::default(); + let request = signed( + json( + url, + &[("x-caller".into(), "kept".into())], + &body, + Some(timeout), + ) + .unwrap(), + &signer, + ) + .unwrap(); + + assert_eq!(request.url(), expected.url()); + assert_eq!(request.method(), expected.method()); + assert_eq!(request.timeout(), expected.timeout()); assert_eq!( - OutboundRequest::signed_json("u".into(), Vec::new(), &json!({}), None, &Refuses), - Err(Error::ComputedHeader("authorization".into())) + request.headers()[CONTENT_TYPE], + expected.headers()[CONTENT_TYPE] + ); + assert_eq!( + request.headers()["x-caller"], + expected.headers()["x-caller"] + ); + assert_eq!(request.headers()["authorization"], "signed"); + assert_eq!( + request.body().unwrap().as_bytes(), + expected.body().unwrap().as_bytes() + ); + assert_eq!( + request.body().unwrap().as_bytes().unwrap(), + signer.0.lock().unwrap().as_slice() ); } + + #[rstest] + #[case::relative_url("not a url", "kept")] + #[case::invalid_host("https://[", "kept")] + #[case::invalid_port("https://u.test:invalid", "kept")] + #[case::invalid_header("https://u.test/", "\u{7f}")] + fn a_request_that_cannot_be_sent_fails_when_it_is_built( + #[case] url: &str, + #[case] header: &str, + ) { + assert!( + reqwest::Client::new() + .post(url) + .header("x-caller", header) + .build() + .is_err() + ); + assert!(matches!( + json(url, &[("x-caller".into(), header.into())], &json!({}), None), + Err(error) if error.is_request() + )); + } + + #[test] + fn a_relative_http_request_is_the_callers_fault_and_never_reached_the_provider() { + let request = http::Request::get("/private?api_key=secret") + .body(Vec::::new()) + .unwrap(); + + let error = from_http(request).unwrap_err(); + + assert!(error.is_request()); + assert!(!error.dispatched()); + let report = litellm_error::Report::new(&error).to_string(); + assert!(!report.contains("secret"), "{report}"); + } } diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 78932f9e29f..3bdc63453ef 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -6,7 +6,12 @@ use std::{ use reqwest::dns::Resolve; -use crate::{client::Client, config::HttpClientConfig, error::Error, proxy::EnvironmentProxies}; +use crate::{ + client::Client, + config::{HttpClientConfig, Verify}, + error::{Error, Repr}, + proxy::EnvironmentProxies, +}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { @@ -14,6 +19,7 @@ pub enum ClientVariant { NoRedirect, Media, UnpinnedMedia, + WebSocket, } const CLIENT_TTL: Duration = Duration::from_secs(3600); @@ -29,8 +35,6 @@ pub struct HttpClientPool { media_resolver: Arc, ttl: Duration, clients: Mutex, - #[cfg(feature = "mcp")] - mcp: crate::mcp::Pool, } impl HttpClientPool { @@ -43,19 +47,9 @@ impl HttpClientPool { media_resolver, ttl, clients: Mutex::default(), - #[cfg(feature = "mcp")] - mcp: crate::mcp::Pool::default(), } } - #[cfg(feature = "mcp")] - pub fn mcp_client( - &self, - config: &HttpClientConfig, - ) -> Result { - self.mcp.client(config, self.ttl) - } - pub fn client( &self, config: &HttpClientConfig, @@ -71,6 +65,14 @@ impl HttpClientPool { client_certificate: None, ..config.clone() }, + ClientVariant::WebSocket => HttpClientConfig { + http2: false, + verify: match &config.verify { + Verify::BuiltInRoots => Verify::NativeRoots, + verify => verify.clone(), + }, + ..config.clone() + }, ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), }; let key = (effective, variant); @@ -81,7 +83,8 @@ impl HttpClientPool { } let client = self .apply(variant, reqwest::ClientBuilder::try_from(&key.0)?) - .build()?; + .build() + .map_err(Repr::Client)?; self.lock().insert( key, PooledClient { @@ -103,12 +106,12 @@ impl HttpClientPool { ) -> reqwest::ClientBuilder { match variant { ClientVariant::Provider => builder, - ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => { + ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia | ClientVariant::WebSocket => { builder.redirect(reqwest::redirect::Policy::none()) } ClientVariant::Media => builder .redirect(reqwest::redirect::Policy::none()) - .dns_resolver2(Arc::clone(&self.media_resolver)), + .dns_resolver(Arc::clone(&self.media_resolver)), } } } @@ -390,56 +393,4 @@ mod tests { .is_err() ); } - #[cfg(feature = "mcp")] - #[rstest::rstest] - #[tokio::test] - async fn mcp_clients_reuse_connections_and_apply_proxy_headers_without_redirects() { - let (address, connections, requests) = serve("HTTP/1.1 302 Found").await; - let pool = pool(); - let config = HttpClientConfig { - proxies: proxied_through(&format!("http://{address}")), - ..config("mcp-pool-test") - }; - for _ in 0..2 { - let response = pool - .mcp - .client(&config, pool.ttl) - .unwrap() - .get("http://upstream.invalid/mcp") - .timeout(Duration::from_secs(2)) - .send() - .await - .unwrap(); - assert_eq!(response.status(), 302); - response.bytes().await.unwrap(); - } - assert_eq!(connections.load(Ordering::SeqCst), 1); - let requests = requests.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!( - requests - .iter() - .all(|request| request.starts_with("GET http://upstream.invalid/mcp HTTP/1.1")) - ); - assert!( - requests - .iter() - .all(|request| request.contains("user-agent: mcp-pool-test")) - ); - } - - #[cfg(feature = "mcp")] - #[rstest::rstest] - #[tokio::test] - async fn mcp_client_uses_host_tls_configuration() { - let directory = tempfile::tempdir().unwrap(); - let config = HttpClientConfig { - verify: Verify::CaBundle(directory.path().join("missing.pem")), - ..config("mcp-test") - }; - assert!(matches!( - pool().mcp_client(&config), - Err(Error::Read { .. }) - )); - } } diff --git a/litellm-rust/crates/http/src/proxy.rs b/litellm-rust/crates/http/src/proxy.rs index e7e54ed6ae2..eb960d8200d 100644 --- a/litellm-rust/crates/http/src/proxy.rs +++ b/litellm-rust/crates/http/src/proxy.rs @@ -63,19 +63,6 @@ impl EnvironmentProxies { .map(|proxy| proxy.no_proxy(no_proxy.clone())) .collect() } - #[cfg(feature = "mcp")] - pub(crate) fn mcp_proxies(&self) -> Vec { - let no_proxy = reqwest_mcp::NoProxy::from_string(&self.no); - [ - reqwest_mcp::Proxy::http(self.http.as_str()), - reqwest_mcp::Proxy::https(self.https.as_str()), - reqwest_mcp::Proxy::all(self.all.as_str()), - ] - .into_iter() - .filter_map(Result::ok) - .map(|proxy| proxy.no_proxy(no_proxy.clone())) - .collect() - } } #[cfg(test)] diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index cfebe8c6b01..11e19f7fa24 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -9,9 +9,7 @@ pub struct HeaderError { use litellm_core_utils::core_helpers::json_type_name; use serde_json::{Map, Value}; -/// Max characters of an upstream error body echoed across the call boundary -/// before truncation, so provider bodies are bounded and data-minimized. -const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256; +use crate::Error; pub enum HeaderPolicy<'a> { All, @@ -40,25 +38,27 @@ pub fn with_headers( }) } -pub async fn http_request( - request: reqwest::RequestBuilder, -) -> Result { - request.send().await +pub async fn http_request(request: reqwest::RequestBuilder) -> Result { + Ok(request.send().await?) } pub async fn execute_http_request( client: &reqwest::Client, request: reqwest::Request, -) -> Result { - client.execute(request).await +) -> Result { + Ok(client.execute(request).await?) } -pub fn truncate_error_body(body: &str) -> String { - if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS { - return body.to_string(); - } - let truncated: String = body.chars().take(UPSTREAM_ERROR_BODY_MAX_CHARS).collect(); - format!("{truncated}... (truncated)") +pub fn header_pairs(headers: &reqwest::header::HeaderMap) -> Vec<(String, String)> { + headers + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.to_string(), value.to_string())) + }) + .collect() } pub fn string_headers( @@ -203,22 +203,6 @@ mod tests { assert_ne!(request.headers()["content-length"], "0"); } - #[test] - fn truncate_leaves_short_bodies_untouched() { - assert_eq!(truncate_error_body("short"), "short"); - } - - #[test] - fn truncate_bounds_long_bodies_by_characters() { - let body = "\u{00e9}".repeat(UPSTREAM_ERROR_BODY_MAX_CHARS + 10); - let truncated = truncate_error_body(&body); - assert!(truncated.ends_with("... (truncated)")); - assert_eq!( - truncated.chars().count(), - UPSTREAM_ERROR_BODY_MAX_CHARS + "... (truncated)".chars().count() - ); - } - #[test] fn string_headers_rejects_non_string_values() { let headers = Map::from_iter([("x-trace".to_string(), json!(7))]); diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index 0b1c44bf6fe..a83b0018b2b 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -1,4 +1,4 @@ -use std::{fmt, path::Path, str::FromStr, sync::Arc}; +use std::{path::Path, str::FromStr, sync::Arc}; use rustls::{ CipherSuite, ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme, @@ -9,7 +9,7 @@ use rustls::{ use crate::{ config::{ClientIdentity, HttpClientConfig, Verify}, - error::{Error, TlsSource}, + error::{Error, PemFault, Repr, TlsSource}, }; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -187,7 +187,7 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { }); let builder = ClientConfig::builder_with_provider(Arc::clone(&provider)) .with_safe_default_protocol_versions() - .map_err(|error| Error::Client(error.to_string()))?; + .map_err(Repr::TlsConfig)?; let verified = match &config.verify { Verify::Disabled => builder .dangerous() @@ -195,6 +195,12 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(), }), + Verify::NativeRoots => { + let native = rustls_native_certs::load_native_certs(); + let mut roots = RootCertStore::empty(); + roots.add_parsable_certificates(native.certs); + builder.with_root_certificates(roots) + } Verify::CaBundle(path) => { builder.with_root_certificates(bundle_roots(path, TlsSource::CaBundle)?) } @@ -224,7 +230,7 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { fn bundle_roots(path: &Path, source: TlsSource) -> Result { let certificates = certificates(path, source)?; if certificates.is_empty() { - return Err(invalid_pem(path, source, "no certificates found")); + return Err(invalid_pem(path, source, PemFault::NoCertificates)); } let mut store = RootCertStore::empty(); for certificate in certificates { @@ -242,7 +248,7 @@ fn client_identity( let source = TlsSource::ClientIdentity; let chain = certificates(certificate, source)?; if chain.is_empty() { - return Err(invalid_pem(certificate, source, "no certificates found")); + return Err(invalid_pem(certificate, source, PemFault::NoCertificates)); } let private_key = PrivateKeyDer::from_pem_slice(&read(key, source)?) .map_err(|error| invalid_pem(key, source, error))?; @@ -256,19 +262,23 @@ fn certificates(path: &Path, source: TlsSource) -> Result Result, Error> { - std::fs::read(path).map_err(|error| Error::Read { - path: path.to_path_buf(), - message: error.to_string(), - tls_source: source, + std::fs::read(path).map_err(|error| { + Repr::Read { + path: path.to_path_buf(), + tls_source: source, + source: error, + } + .into() }) } -fn invalid_pem(path: &Path, source: TlsSource, message: impl fmt::Display) -> Error { - Error::InvalidPem { +fn invalid_pem(path: &Path, tls_source: TlsSource, fault: impl Into) -> Error { + Repr::InvalidPem { path: path.to_path_buf(), - message: message.to_string(), - tls_source: source, + tls_source, + source: fault.into(), } + .into() } #[derive(Debug)] @@ -448,11 +458,10 @@ mod tests { std::fs::remove_file(&path).unwrap(); assert!(matches!( result, - Err(Error::InvalidPem { - path: reported, - tls_source: TlsSource::ClientIdentity, - .. - }) if reported == path + Err(error) if matches!( + error.repr(), + Repr::InvalidPem { path: reported, tls_source: TlsSource::ClientIdentity, .. } if *reported == path + ) )); } diff --git a/litellm-rust/crates/http/src/transport.rs b/litellm-rust/crates/http/src/transport.rs deleted file mode 100644 index 8814925bbf2..00000000000 --- a/litellm-rust/crates/http/src/transport.rs +++ /dev/null @@ -1,109 +0,0 @@ -#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] -pub enum Error { - #[error("upstream request failed with status {status}: {body}")] - Http { status: u16, body: String }, - #[error("upstream network error: {0}")] - Network(String), - #[error("could not reach the provider: {0}")] - Connect(String), -} - -impl Error { - pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self { - let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder()); - let message = describe(error); - if before_dispatch { - Self::Connect(message) - } else { - Self::Network(message) - } - } -} - -impl From for Error { - fn from(error: reqwest::Error) -> Self { - Self::Network(describe(error)) - } -} - -fn describe(error: reqwest::Error) -> String { - let error = error.without_url(); - std::iter::successors(std::error::Error::source(&error), |cause| cause.source()) - .fold(error.to_string(), |message, cause| { - format!("{message}: {cause}") - }) -} - -#[cfg(test)] -mod tests { - #[tokio::test] - async fn transport_errors_remove_urls_and_keep_dispatch_context() { - let error = reqwest::Client::builder() - .no_proxy() - .build() - .expect("client") - .get("http://localhost:invalid/private?api_key=secret") - .send() - .await - .expect_err("invalid port"); - let error = crate::transport::Error::from_reqwest_before_dispatch(error); - assert!(matches!(error, crate::transport::Error::Connect(_))); - assert!(!error.to_string().contains("secret")); - assert!(!error.to_string().contains("private")); - } - - fn root_cause(error: &dyn std::error::Error) -> Option { - match error.source() { - Some(cause) => root_cause(cause).or_else(|| Some(cause.to_string())), - None => None, - } - } - - #[tokio::test] - async fn network_error_message_names_the_underlying_cause() { - let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); - let address = listener.local_addr().expect("address"); - drop(listener); - let error = reqwest::Client::builder() - .no_proxy() - .build() - .expect("client") - .get(format!("http://{address}/private?api_key=secret")) - .send() - .await - .expect_err("nothing listens on the port"); - let root_cause = root_cause(&error).expect("reqwest reports a cause"); - let message = crate::transport::Error::from(error).to_string(); - assert!(message.contains(&root_cause), "{message}"); - assert!(!message.contains("secret")); - } - - #[tokio::test] - async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() { - use std::time::Duration; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind"); - let address = listener.local_addr().expect("address"); - let request = reqwest::Client::builder() - .no_proxy() - .build() - .expect("client") - .get(format!("http://{address}")) - .timeout(Duration::from_millis(200)) - .send(); - let (response, accepted) = tokio::join!( - request, - tokio::time::timeout(Duration::from_secs(2), listener.accept()) - ); - let _connection = accepted - .expect("accept deadline") - .expect("accepted connection"); - let error = response.expect_err("server does not respond"); - assert!(error.is_timeout()); - assert!(matches!( - crate::transport::Error::from_reqwest_before_dispatch(error), - crate::transport::Error::Network(_) - )); - } -} diff --git a/litellm-rust/crates/http/src/websocket.rs b/litellm-rust/crates/http/src/websocket.rs index e904392152a..b2a6571e2b8 100644 --- a/litellm-rust/crates/http/src/websocket.rs +++ b/litellm-rust/crates/http/src/websocket.rs @@ -1,62 +1,137 @@ -use std::{ - io, - sync::{Arc, OnceLock}, -}; +use std::time::Duration; -use rustls::{ClientConfig, RootCertStore}; -use tokio::net::TcpStream; +use http::{ + HeaderMap, HeaderValue, StatusCode, + header::{ + CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_PROTOCOL, + SEC_WEBSOCKET_VERSION, UPGRADE, + }, +}; +use reqwest::Url; use tokio_tungstenite::{ - Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, - tungstenite::{client::IntoClientRequest, error::TlsError, handshake::client::Response}, + WebSocketStream, + tungstenite::{ + handshake::{client::generate_key, derive_accept_key}, + protocol::Role, + }, }; -pub type UpstreamWebSocket = WebSocketStream>; +use crate::{ + client::Client, + error::{Error, HandshakeFault, Repr}, +}; -static TLS_CONFIG: OnceLock> = OnceLock::new(); - -fn build_tls_config() -> Result> { - let native = rustls_native_certs::load_native_certs(); - let mut store = RootCertStore::empty(); - let (added, _ignored) = store.add_parsable_certificates(native.certs); - if added == 0 { - return Err(Box::new(tokio_tungstenite::tungstenite::Error::Io( - io::Error::other(format!( - "no usable native root certificates: {:?}", - native.errors - )), - ))); - } - ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) - .with_safe_default_protocol_versions() - .map(|builder| builder.with_root_certificates(store).with_no_client_auth()) - .map_err(|error| { - Box::new(tokio_tungstenite::tungstenite::Error::Tls( - TlsError::Rustls(error), - )) - }) -} - -fn tls_config() -> Result, Box> { - if let Some(config) = TLS_CONFIG.get() { - return Ok(Arc::clone(config)); - } - let built = Arc::new(build_tls_config()?); - Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) -} +pub type UpstreamWebSocket = WebSocketStream; #[tracing::instrument(name = "litellm.websocket.handshake", level = "debug", skip_all)] -pub async fn connect_upstream( - request: R, -) -> Result<(UpstreamWebSocket, Response), Box> -where - R: IntoClientRequest + Unpin, -{ - let request = request.into_client_request().map_err(Box::new)?; - let connector = match request.uri().scheme_str() { - Some("wss") => Some(Connector::Rustls(tls_config()?)), - _ => None, - }; - connect_async_tls_with_config(request, None, false, connector) - .await - .map_err(Box::new) +pub async fn connect( + client: &Client, + url: &str, + headers: HeaderMap, + timeout: Option, +) -> Result { + match timeout { + Some(timeout) => tokio::time::timeout(timeout, handshake(client, url, headers)) + .await + .map_err(Repr::HandshakeTimeout)?, + None => handshake(client, url, headers).await, + } +} + +async fn handshake( + client: &Client, + url: &str, + headers: HeaderMap, +) -> Result { + let offered_protocols = headers.get(SEC_WEBSOCKET_PROTOCOL).cloned(); + let key = generate_key(); + let response = client + .get(http_url(url)?) + .headers(handshake_headers(headers, &key)?) + .send() + .await?; + let status = response.status(); + if status != StatusCode::SWITCHING_PROTOCOLS { + return Err(Error::upstream( + status.as_u16(), + response.text().await.unwrap_or_default(), + )); + } + let expected = derive_accept_key(key.as_bytes()); + if response + .headers() + .get(SEC_WEBSOCKET_ACCEPT) + .map(HeaderValue::as_bytes) + != Some(expected.as_bytes()) + { + return Err(Repr::Handshake(HandshakeFault::Accept).into()); + } + validate_response_headers(response.headers(), offered_protocols.as_ref())?; + let upgraded = response.upgrade().await?; + Ok(WebSocketStream::from_raw_socket(upgraded, Role::Client, None).await) +} + +fn validate_response_headers( + headers: &HeaderMap, + offered_protocols: Option<&HeaderValue>, +) -> Result<(), Error> { + if !headers + .get(UPGRADE) + .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket")) + { + return Err(Repr::Handshake(HandshakeFault::Upgrade).into()); + } + if !headers.get_all(CONNECTION).iter().any(|value| { + value.to_str().is_ok_and(|value| { + value + .split(',') + .any(|token| token.trim().eq_ignore_ascii_case("upgrade")) + }) + }) { + return Err(Repr::Handshake(HandshakeFault::Connection).into()); + } + let protocol_matches = match (offered_protocols, headers.get(SEC_WEBSOCKET_PROTOCOL)) { + (_, None) => true, + (Some(offered), Some(selected)) => offered + .to_str() + .ok() + .zip(selected.to_str().ok()) + .is_some_and(|(offered, selected)| { + offered + .split(',') + .any(|protocol| protocol.trim() == selected) + }), + _ => false, + }; + if !protocol_matches { + return Err(Repr::Handshake(HandshakeFault::Subprotocol).into()); + } + Ok(()) +} + +fn http_url(url: &str) -> Result { + let mut url = Url::parse(url).map_err(Repr::WebSocketUrl)?; + let scheme = match url.scheme() { + "ws" | "http" => "http", + "wss" | "https" => "https", + other => return Err(Repr::WebSocketScheme(other.to_owned()).into()), + }; + url.set_scheme(scheme) + .map_err(|()| Repr::WebSocketSchemeRewrite(scheme))?; + Ok(url) +} + +fn handshake_headers(headers: HeaderMap, key: &str) -> Result { + let key = HeaderValue::from_str(key).map_err(|error| Repr::Request(error.into()))?; + Ok([ + (CONNECTION, HeaderValue::from_static("Upgrade")), + (UPGRADE, HeaderValue::from_static("websocket")), + (SEC_WEBSOCKET_VERSION, HeaderValue::from_static("13")), + (SEC_WEBSOCKET_KEY, key), + ] + .into_iter() + .fold(headers, |mut headers, (name, value)| { + headers.insert(name, value); + headers + })) } diff --git a/litellm-rust/crates/http/tests/websocket.rs b/litellm-rust/crates/http/tests/websocket.rs index 2fbaa3994e8..a5758d45e56 100644 --- a/litellm-rust/crates/http/tests/websocket.rs +++ b/litellm-rust/crates/http/tests/websocket.rs @@ -1,59 +1,560 @@ +use std::{ + net::SocketAddr, + path::PathBuf, + sync::{Arc, Mutex}, +}; + use futures_util::{SinkExt, StreamExt}; -use litellm_http::websocket::connect_upstream; -use rstest::rstest; -use tokio::net::TcpListener; +use http::{HeaderMap, HeaderValue}; +use litellm_http::{ + ClientIdentity, ClientVariant, EnvironmentProxies, Error, HttpClientConfig, HttpClientPool, + HttpSettings, Resolution, Verify, media::PublicDnsResolver, websocket::connect, +}; +use rcgen::{BasicConstraints, CertificateParams, CertifiedIssuer, IsCa, KeyPair}; +use rstest::{Context, rstest}; +use rustls::{ + RootCertStore, ServerConfig, + crypto::ring, + pki_types::{CertificateDer, PrivateKeyDer}, + server::WebPkiClientVerifier, +}; +use tokio::{ + io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + net::{TcpListener, TcpStream}, +}; +use tokio_rustls::TlsAcceptor; use tokio_tungstenite::{ accept_hdr_async, tungstenite::{ Message, - client::IntoClientRequest, handshake::server::{Request, Response}, }, }; +struct Pki { + directory: tempfile::TempDir, + ca: CertifiedIssuer<'static, KeyPair>, + server_chain: Vec>, + server_key: Vec, +} + +impl Pki { + fn new() -> Self { + let mut ca_params = CertificateParams::new(Vec::::new()).unwrap(); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca = CertifiedIssuer::self_signed(ca_params, KeyPair::generate().unwrap()).unwrap(); + let server_key = KeyPair::generate().unwrap(); + let server = CertificateParams::new(vec!["localhost".into()]) + .unwrap() + .signed_by(&server_key, &ca) + .unwrap(); + let client_key = KeyPair::generate().unwrap(); + let client = CertificateParams::new(vec!["client".into()]) + .unwrap() + .signed_by(&client_key, &ca) + .unwrap(); + let directory = tempfile::tempdir().unwrap(); + std::fs::write(directory.path().join("ca.pem"), ca.pem()).unwrap(); + std::fs::write( + directory.path().join("client.pem"), + format!("{}{}", client.pem(), client_key.serialize_pem()), + ) + .unwrap(); + Self { + directory, + server_chain: vec![server.der().clone()], + server_key: server_key.serialize_der(), + ca, + } + } + + fn path(&self, name: &str) -> PathBuf { + self.directory.path().join(name) + } + + fn acceptor(&self, require_client_certificate: bool) -> TlsAcceptor { + let provider = Arc::new(ring::default_provider()); + let builder = ServerConfig::builder_with_provider(Arc::clone(&provider)) + .with_safe_default_protocol_versions() + .unwrap(); + let builder = if require_client_certificate { + let mut roots = RootCertStore::empty(); + roots.add(self.ca.der().clone()).unwrap(); + builder.with_client_cert_verifier( + WebPkiClientVerifier::builder_with_provider(Arc::new(roots), provider) + .build() + .unwrap(), + ) + } else { + builder.with_no_client_auth() + }; + let mut config = builder + .with_single_cert( + self.server_chain.clone(), + PrivateKeyDer::try_from(self.server_key.clone()).unwrap(), + ) + .unwrap(); + config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + TlsAcceptor::from(Arc::new(config)) + } +} + +fn config(verify: Verify) -> HttpClientConfig { + HttpClientConfig { + verify, + http2: true, + ..Resolution::from(&HttpSettings::default()).config + } +} + +fn proxied_through(proxy: SocketAddr) -> EnvironmentProxies { + EnvironmentProxies::from_environment(&move |name: &str| { + (name == "HTTPS_PROXY").then(|| format!("http://{proxy}")) + }) +} + +async fn open( + config: &HttpClientConfig, + url: &str, + headers: HeaderMap, +) -> Result { + let client = HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client(config, ClientVariant::WebSocket) + .unwrap(); + connect(&client, url, headers, None).await +} + +async fn echo_once( + stream: S, + seen: Arc>>, +) { + #[expect( + clippy::result_large_err, + reason = "tungstenite requires an unboxed handshake error response" + )] + let Ok(mut socket) = accept_hdr_async(stream, |request: &Request, response: Response| { + *seen.lock().unwrap() = Some(request.clone()); + Ok(response) + }) + .await + else { + return; + }; + if let Some(Ok(message)) = socket.next().await { + let _ = socket.send(message).await; + } +} + +async fn serve_tls( + pki: &Pki, + require_client_certificate: bool, +) -> (SocketAddr, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let acceptor = pki.acceptor(require_client_certificate); + let seen = Arc::new(Mutex::new(None)); + let recorded = Arc::clone(&seen); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + if let Ok(stream) = acceptor.accept(stream).await { + echo_once(stream, Arc::clone(&recorded)).await; + } + } + }); + (address, seen) +} + +async fn read_head(stream: &mut TcpStream) -> String { + let mut head = Vec::new(); + while !head.ends_with(b"\r\n\r\n") { + let mut byte = [0u8; 1]; + if stream.read(&mut byte).await.unwrap() == 0 { + break; + } + head.push(byte[0]); + } + String::from_utf8(head).unwrap() +} + +async fn serve_tunnel_proxy(target: SocketAddr) -> (SocketAddr, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let seen = Arc::clone(&requests); + tokio::spawn(async move { + while let Ok((mut client, _)) = listener.accept().await { + let head = read_head(&mut client).await; + seen.lock() + .unwrap() + .push(head.lines().next().unwrap_or_default().to_owned()); + let mut upstream = TcpStream::connect(target).await.unwrap(); + client + .write_all(b"HTTP/1.1 200 Connection established\r\n\r\n") + .await + .unwrap(); + let _ = tokio::io::copy_bidirectional(&mut client, &mut upstream).await; + } + }); + (address, requests) +} + +async fn serve_raw(response: &'static str) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + read_head(&mut stream).await; + stream.write_all(response.as_bytes()).await.unwrap(); + let _ = stream.read(&mut [0u8; 1]).await; + }); + address +} + +async fn round_trip(socket: &mut litellm_http::websocket::UpstreamWebSocket) { + socket.send(Message::Text("ping".into())).await.unwrap(); + assert_eq!( + socket.next().await.unwrap().unwrap(), + Message::Text("ping".into()) + ); +} + #[rstest] #[case::without_query("/responses")] #[case::with_query("/responses?model=test-model")] #[tokio::test] -async fn connects_with_caller_headers_and_exchanges_frames(#[case] path: &'static str) { +async fn plain_websocket_sends_caller_headers_and_path(#[case] path: &'static str) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { + let seen = Arc::new(Mutex::new(None)); + let recorded = Arc::clone(&seen); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + echo_once(stream, recorded).await; + }); + let headers = HeaderMap::from_iter([( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer test-key"), + )]); + + let mut socket = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}{path}"), + headers, + ) + .await + .unwrap(); + + round_trip(&mut socket).await; + let request = seen.lock().unwrap().take().unwrap(); + assert_eq!(request.uri().path_and_query().unwrap().as_str(), path); + assert_eq!(request.headers()["authorization"], "Bearer test-key"); +} + +#[rstest] +#[case::ca_bundle_trusts_the_private_ca(|pki: &Pki| Verify::CaBundle(pki.path("ca.pem")), true)] +#[case::ssl_verify_off_skips_verification(|_: &Pki| Verify::Disabled, true)] +#[case::built_in_roots_reject_the_private_ca(|_: &Pki| Verify::BuiltInRoots, false)] +#[tokio::test] +async fn wss_honors_ssl_verify(#[case] verify: fn(&Pki) -> Verify, #[case] trusted: bool) { + let pki = Pki::new(); + let (address, _) = serve_tls(&pki, false).await; + + let result = open( + &config(verify(&pki)), + &format!("wss://localhost:{}/responses", address.port()), + HeaderMap::new(), + ) + .await; + + match result { + Ok(mut socket) if trusted => round_trip(&mut socket).await, + Err(error) if !trusted && !error.dispatched() => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } +} + +#[rstest] +#[case::plain_without_roots(false, false, true)] +#[case::tls_with_roots(true, true, true)] +#[case::tls_without_roots(true, false, false)] +#[tokio::test] +async fn websocket_uses_system_roots( + #[context] context: Context, + #[case] tls: bool, + #[case] trusted_roots: bool, + #[case] connects: bool, +) { + const CHILD_URL: &str = "LITELLM_TEST_SYSTEM_ROOTS_WSS_URL"; + if let Ok(url) = std::env::var(CHILD_URL) { + let result = open(&config(Verify::BuiltInRoots), &url, HeaderMap::new()).await; + match result { + Ok(mut socket) if connects => round_trip(&mut socket).await, + Err(error) if !connects && error.is_connect() => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } + return; + } + + let pki = Pki::new(); + let url = if tls { + let (address, _) = serve_tls(&pki, false).await; + format!("wss://localhost:{}/responses", address.port()) + } else { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + echo_once(stream, Arc::new(Mutex::new(None))).await; + }); + format!("ws://{address}/responses") + }; + let ca = if trusted_roots { + pki.path("ca.pem") + } else { + let empty = pki.path("empty.pem"); + std::fs::write(&empty, "").unwrap(); + empty + }; + let test = format!( + "{}::case_{}_{}", + context.name, + context.case.unwrap() + 1, + context.description.unwrap() + ); + let output = tokio::task::spawn_blocking(move || { + std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", &test, "--nocapture"]) + .env(CHILD_URL, url) + .env("SSL_CERT_FILE", ca) + .env("SSL_CERT_DIR", "") + .output() + .unwrap() + }) + .await + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed"), + "system-root WebSocket child failed:\n{stdout}\n{}", + String::from_utf8_lossy(&output.stderr) + ); +} + +#[rstest] +#[case::with_certificate(true)] +#[case::without_certificate(false)] +#[tokio::test] +async fn wss_presents_the_configured_client_certificate(#[case] configured: bool) { + let pki = Pki::new(); + let (address, _) = serve_tls(&pki, true).await; + let config = HttpClientConfig { + client_certificate: configured.then(|| ClientIdentity::Pem(pki.path("client.pem"))), + ..config(Verify::CaBundle(pki.path("ca.pem"))) + }; + + let result = open( + &config, + &format!("wss://localhost:{}/responses", address.port()), + HeaderMap::new(), + ) + .await; + + match result { + Ok(mut socket) if configured => round_trip(&mut socket).await, + Err(_) if !configured => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } +} + +#[rstest] +#[tokio::test] +async fn wss_tunnels_through_the_https_proxy() { + let pki = Pki::new(); + let (target, _) = serve_tls(&pki, false).await; + let (proxy, requests) = serve_tunnel_proxy(target).await; + let config = HttpClientConfig { + proxies: proxied_through(proxy), + ..config(Verify::CaBundle(pki.path("ca.pem"))) + }; + + let mut socket = open( + &config, + &format!("wss://localhost:{}/responses", target.port()), + HeaderMap::new(), + ) + .await + .unwrap(); + + round_trip(&mut socket).await; + assert_eq!( + *requests.lock().unwrap(), + [format!("CONNECT localhost:{} HTTP/1.1", target.port())] + ); +} + +#[rstest] +#[tokio::test] +async fn refused_upgrade_reports_the_status_and_body() { + let address = + serve_raw("HTTP/1.1 403 Forbidden\r\nContent-Length: 6\r\nConnection: close\r\n\r\ndenied") + .await; + + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + HeaderMap::new(), + ) + .await; + + let error = result.map(drop).unwrap_err(); + assert_eq!(error.upstream_response(), Some((403, "denied"))); +} + +#[rstest] +#[tokio::test] +async fn upgrade_with_a_wrong_accept_key_is_rejected() { + let address = serve_raw( + "HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Accept: wrong\r\n\r\n", + ) + .await; + + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + HeaderMap::new(), + ) + .await; + + assert!(matches!(result, Err(error) if error.dispatched() && !error.is_request())); +} + +#[rstest] +#[case::non_websocket_scheme("ftp://localhost/responses")] +#[case::unparseable("not a url")] +#[tokio::test] +async fn invalid_url_is_rejected_before_dispatch(#[case] url: &str) { + let result = open(&config(Verify::BuiltInRoots), url, HeaderMap::new()).await; + + assert!(matches!(result, Err(error) if !error.dispatched() && error.is_request())); +} + +#[rstest] +#[case::missing_upgrade(None, Some("Upgrade"), None, None, false)] +#[case::wrong_upgrade(Some("http"), Some("Upgrade"), None, None, false)] +#[case::missing_connection(Some("websocket"), None, None, None, false)] +#[case::wrong_connection(Some("websocket"), Some("keep-alive"), None, None, false)] +#[case::unsolicited_protocol(Some("websocket"), Some("Upgrade"), None, Some("other"), false)] +#[case::unoffered_protocol( + Some("websocket"), + Some("Upgrade"), + Some("events"), + Some("other"), + false +)] +#[case::declined_optional_protocol(Some("websocket"), Some("Upgrade"), Some("events"), None, true)] +#[case::matching_protocol( + Some("websocket"), + Some("Upgrade"), + Some("events"), + Some("events"), + true +)] +#[case::multiple_offers( + Some("websocket"), + Some("Upgrade"), + Some("events, realtime"), + Some("realtime"), + true +)] +#[case::case_insensitive_upgrade(Some("WebSocket"), Some("uPgRaDe"), None, None, true)] +#[case::connection_tokens(Some("websocket"), Some("keep-alive, Upgrade"), None, None, true)] +#[tokio::test] +async fn upgrade_validates_headers_and_negotiated_protocol( + #[case] upgrade: Option<&'static str>, + #[case] connection: Option<&'static str>, + #[case] offered: Option<&'static str>, + #[case] selected: Option<&'static str>, + #[case] accepted: bool, +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); #[expect( clippy::result_large_err, reason = "tungstenite requires an unboxed handshake error response" )] - let mut socket = accept_hdr_async(stream, move |request: &Request, response: Response| { - assert_eq!(request.uri().path_and_query().unwrap().as_str(), path); - assert_eq!(request.headers()["authorization"], "Bearer test-key"); + let result = accept_hdr_async(stream, move |_: &Request, mut response: Response| { + response.headers_mut().remove("upgrade"); + response.headers_mut().remove("connection"); + if let Some(value) = upgrade { + response + .headers_mut() + .insert("upgrade", HeaderValue::from_static(value)); + } + if let Some(value) = connection { + response + .headers_mut() + .insert("connection", HeaderValue::from_static(value)); + } + if let Some(value) = selected { + response + .headers_mut() + .insert("sec-websocket-protocol", HeaderValue::from_static(value)); + } Ok(response) }) - .await - .unwrap(); - assert_eq!( - socket.next().await.unwrap().unwrap(), - Message::Text("request".into()) - ); - socket.send(Message::Text("response".into())).await.unwrap(); - socket.close(None).await.unwrap(); + .await; + let Ok(mut socket) = result else { + return; + }; + if let Some(Ok(message)) = socket.next().await { + let _ = socket.send(message).await; + } }); - let mut request = format!("ws://{address}{path}") - .into_client_request() - .unwrap(); - request - .headers_mut() - .insert("authorization", "Bearer test-key".parse().unwrap()); - let (mut socket, response) = connect_upstream(request).await.unwrap(); - assert_eq!(response.status(), 101); - socket.send(Message::Text("request".into())).await.unwrap(); - assert_eq!( - socket.next().await.unwrap().unwrap(), - Message::Text("response".into()) - ); - assert!(matches!( - socket.next().await.unwrap().unwrap(), - Message::Close(_) - )); - server.await.unwrap(); + let headers = offered.map_or_else(HeaderMap::new, |value| { + HeaderMap::from_iter([( + http::header::SEC_WEBSOCKET_PROTOCOL, + HeaderValue::from_static(value), + )]) + }); + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + headers, + ) + .await; + + match result { + Ok(mut socket) if accepted => round_trip(&mut socket).await, + Err(error) if !accepted && error.dispatched() => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } +} + +#[rstest] +#[tokio::test] +async fn a_handshake_that_never_answers_times_out_as_dispatched() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + read_head(&mut stream).await; + let _ = stream.read(&mut [0u8; 1]).await; + }); + let client = HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client(&config(Verify::BuiltInRoots), ClientVariant::WebSocket) + .unwrap(); + + let error = connect( + &client, + &format!("ws://{address}/responses"), + HeaderMap::new(), + Some(std::time::Duration::from_millis(100)), + ) + .await + .map(drop) + .unwrap_err(); + + assert!(error.dispatched()); } diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index 130786e788d..8a779d44976 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -6,7 +6,6 @@ license.workspace = true repository.workspace = true [dependencies] -futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-cache-response.workspace = true litellm-core-utils.workspace = true @@ -16,11 +15,14 @@ litellm-inference.workspace = true litellm-llms.workspace = true litellm-llms-types.workspace = true litellm-secrets.workspace = true + serde_json = { workspace = true, features = ["preserve_order"] } +futures-util.workspace = true tracing.workspace = true [dev-dependencies] litellm-cache-memory.workspace = true +litellm-error.workspace = true litellm-host-native.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-inference-testing.workspace = true diff --git a/litellm-rust/crates/inference-chat/src/common_utils.rs b/litellm-rust/crates/inference-chat/src/common_utils.rs index 4063f33432d..1c177d58552 100644 --- a/litellm-rust/crates/inference-chat/src/common_utils.rs +++ b/litellm-rust/crates/inference-chat/src/common_utils.rs @@ -50,18 +50,20 @@ pub(super) fn resolve_request( config: &'static dyn BaseConfig, ) -> Result { let messages = serde_json::from_value::>(call.messages).map_err(|err| { - Error::InvalidRequest(litellm_llms::ErrorDetail::invalid( + Error::invalid_request(litellm_llms::ErrorDetail::invalid( "chat completions messages", err, )) })?; if messages.is_empty() { - return Err(Error::InvalidRequest( - "chat completions requires at least one message".into(), + return Err(Error::invalid_request( + litellm_llms::ErrorDetail::NoMessages { + route: "chat completions", + }, )); } if let Some(reason) = config.unsupported_reason(&messages, &call.optional_params) { - return Err(Error::Unsupported(reason.0)); + return Err(Error::Declined(reason)); } Ok(ResolvedChatCompletionsRequest { messages, diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index 595e851c6fd..1c724a47956 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -3,11 +3,7 @@ use std::time::Duration; use litellm_auth::AuthServices; use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}; -use litellm_http::{ - Client, - outbound::OutboundRequest, - request::{truncate_error_body, with_default_headers}, -}; +use litellm_http::{Client, request::with_default_headers}; use litellm_llms::base_llm::{ auth::{Authenticated, ValidatedEnvironment, resolve_auth}, chat::transformation::BaseConfig, @@ -91,39 +87,26 @@ pub(super) async fn execute( interceptors, observers, || async move { - let outbound = outbound_request( + let outbound = litellm_inference::outbound::outbound_request( Authenticated { headers: wire.headers, signer: authenticated.signer, }, - wire.url, + &wire.url, &wire.body, - call.timeout, + Some( + call.timeout + .unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)), + ), )?; - let response = litellm_inference::outbound::send(outbound, http) - .await - .map_err(|err| { - // Failing to establish the connection means the request never went out, - // so the host can still serve it. Everything else here, a timeout - // above all, may have reached the provider and been answered. - if err.is_connect() || err.is_builder() { - Error::Transport(litellm_http::transport::Error::Connect(err.to_string())) - } else { - Error::Transport(litellm_http::transport::Error::Network(err.to_string())) - } - })?; + let response = litellm_inference::outbound::send(outbound, http).await?; let status = response.status(); - let text = response.text().await.map_err(|err| { - Error::Transport(litellm_http::transport::Error::Network(err.to_string())) - })?; + let text = response.text().await.map_err(litellm_http::Error::from)?; if !status.is_success() { - return Err(Error::Transport(litellm_http::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - })); + return Err(litellm_http::Error::upstream(status.as_u16(), text).into()); } let raw = RawResponse { body: text.clone() }; if let Some(observers) = observers { @@ -137,7 +120,7 @@ pub(super) async fn execute( .map_err(Error::post_call)?; let body: Value = serde_json::from_str(&text).map_err(|err| { - Error::InvalidResponse(litellm_llms::ErrorDetail::invalid( + Error::invalid_response(litellm_llms::ErrorDetail::invalid( "chat completions response JSON", err, )) @@ -145,51 +128,12 @@ pub(super) async fn execute( config .transform_response(provider.model, ProviderChatResponseData { body }) .map_err(Error::from) - .map_err(as_response_error) + .map_err(Error::after_send) }, ) .await } -/// Re-tag an error raised while normalizing a response the provider already -/// returned. -/// -/// A config reports the same variants on either side of the call: a missing -/// field or an unsupported block can mean "this request cannot be translated" -/// before the request is sent and "this response cannot be normalized" here. Only the -/// second kind has already been billed, and a host that keeps a reference -/// implementation must not retry those, so collapse them to one variant that -/// can only mean the provider was already called. -pub(super) fn as_response_error(err: Error) -> Error { - match err { - already @ (Error::InvalidResponse(_) - | Error::Transport(litellm_http::transport::Error::Http { .. })) => already, - other => Error::InvalidResponse(other.to_string().into()), - } -} - -pub(super) fn outbound_request( - authenticated: Authenticated, - url: String, - body: &Value, - timeout: Option, -) -> Result { - litellm_inference::outbound::outbound_request( - authenticated, - url, - body, - Some(timeout.unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS))), - ) - .map_err(|error| match error { - // Python drops the caller's copy and prefers a forwarded Authorization - // over the signature, so leave the request to it. - litellm_http::Error::ComputedHeader(_) => { - Error::Unsupported("request forwards a header AWS SigV4 computes") - } - other => Error::Http(other), - }) -} - #[cfg(test)] mod tests { use std::sync::Mutex; @@ -283,7 +227,7 @@ mod tests { let (provider, config, request) = prepared(&upstream.uri()); execute( - &Client::plain_for_test(), + &Client::for_test(), &AuthServices::default(), config, provider, @@ -328,7 +272,7 @@ mod tests { let (provider, config, request) = prepared(&upstream.uri()); let error = execute( - &Client::plain_for_test(), + &Client::for_test(), &AuthServices::default(), config, provider, @@ -342,31 +286,12 @@ mod tests { .await .expect_err("the upstream failure fails the call"); - assert!(matches!( - error, - Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) - )); + assert!( + error + .upstream_response() + .is_some_and(|(status, _)| status == 500), + "{error:?}" + ); assert!(interceptors.raw.into_inner().unwrap().is_empty()); } - - #[rstest::rstest] - fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { - for original in [ - Error::MissingField("usage"), - Error::Unsupported("non-text response content block"), - Error::InvalidRequest("whatever".to_string().into()), - Error::Auth(litellm_auth::Error::InvalidHeader), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), Error::InvalidResponse(_)), - "{label} must not stay retryable once the provider has answered" - ); - } - let upstream = Error::Transport(litellm_http::transport::Error::Http { - status: 500, - body: "boom".to_string(), - }); - assert_eq!(as_response_error(upstream.clone()), upstream); - } } diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 70da8bc3255..4b3e32f04bc 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -239,7 +239,7 @@ async fn chat_cache_identity_follows_resolved_configuration_and_request_callback .store(usize::from(call >= 2), Ordering::SeqCst); let cache = ScopedCache::new(cache.clone(), CacheScope::Shared); ChatCompletionsRoute::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), Arc::new(Default::default()), secrets.clone(), ) diff --git a/litellm-rust/crates/inference-chat/tests/chat_completions.rs b/litellm-rust/crates/inference-chat/tests/chat_completions.rs index 461d4146a63..6ed513ba32a 100644 --- a/litellm-rust/crates/inference-chat/tests/chat_completions.rs +++ b/litellm-rust/crates/inference-chat/tests/chat_completions.rs @@ -5,7 +5,6 @@ use litellm_host::{ }; use std::time::Duration; -use litellm_http::transport::Error as TransportError; use litellm_inference_chat::{Error, types::ChatCompletionsRequest}; use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse; use rstest::{fixture, rstest}; @@ -180,7 +179,7 @@ async fn a_response_it_cannot_normalize_is_reported_as_already_sent( .await .expect_err("response cannot be normalized"); - assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!(error.is_response(), "{error:?}"); } #[rstest] @@ -201,13 +200,7 @@ async fn an_upstream_error_status_keeps_its_code_and_body( .await .expect_err("upstream rejects"); - assert_eq!( - error, - Error::Transport(TransportError::Http { - status, - body: "slow down".into() - }) - ); + assert_eq!(error.upstream_response(), Some((status, "slow down"))); } #[rstest] @@ -222,10 +215,25 @@ async fn a_connection_that_is_never_established_returns_a_connect_error( .await .expect_err("nothing is listening"); - assert!( - matches!(error, Error::Transport(TransportError::Connect(_))), - "{error:?}" - ); + assert!(!error.dispatched(), "{error:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_transport_error_does_not_echo_the_provider_url( + request: ChatCompletionsRequest<'static>, +) { + let base = format!("{UNREACHABLE_BASE}/tenant-token-1234"); + + let error = complete(ChatCompletionsRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("nothing is listening"); + + let report = litellm_error::Report::new(&error).to_string(); + assert!(!report.contains("tenant-token-1234"), "{report}"); } #[rstest] @@ -243,10 +251,8 @@ async fn a_timeout_after_sending_returns_a_network_error(request: ChatCompletion .await .expect_err("the call times out"); - assert!( - matches!(error, Error::Transport(TransportError::Network(_))), - "{error:?}" - ); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); } #[rstest] @@ -342,7 +348,7 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry( Ok(wire) } async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { - Err(Error::InvalidRequest("callback rejected".into())) + Err(litellm_llms::Error::MissingField("callback rejected").into()) } } let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; @@ -358,10 +364,10 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry( ) .await .unwrap_err(); - let Error::PostCallHook(source) = error else { - panic!("expected retained callback error") - }; - assert_eq!(*source, Error::InvalidRequest("callback rejected".into())); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("callback rejected"), "{report}"); assert_eq!(received(&upstream).await.len(), 1); } diff --git a/litellm-rust/crates/inference-chat/tests/request.rs b/litellm-rust/crates/inference-chat/tests/request.rs index 5988a7baee5..121f7347c24 100644 --- a/litellm-rust/crates/inference-chat/tests/request.rs +++ b/litellm-rust/crates/inference-chat/tests/request.rs @@ -211,7 +211,7 @@ async fn rejects_empty_messages_before_resolving_credentials() { }) .await .expect_err("empty messages are rejected"); - assert!(matches!(error, Error::InvalidRequest(_))); + assert!(error.is_request(), "{error:?}"); } #[rstest] @@ -233,26 +233,27 @@ async fn rejects_an_unknown_provider() { #[rstest] #[tokio::test] async fn rejects_a_model_with_no_resolvable_provider() { - assert!(matches!( - complete(request( - "claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - )) - .await, - Err(Error::InvalidProvider(_)) - )); + let error = complete(request( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .await + .unwrap_err(); + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.contains("unable to resolve custom_llm_provider"), + "{report}" + ); } #[rstest] -#[case::empty(json!([]), Some("chat completions requires at least one message"))] -#[case::malformed(json!("not a list"), None)] +#[case::empty(json!([]), "chat completions requires at least one message")] +#[case::malformed(json!("not a list"), "invalid chat completions messages")] #[tokio::test] -async fn rejects_an_empty_or_malformed_message_list( - #[case] messages: Value, - #[case] expected: Option<&str>, -) { +async fn rejects_an_empty_or_malformed_message_list(#[case] messages: Value, #[case] reason: &str) { let error = complete(request( "anthropic/claude-sonnet-4-5", None, @@ -261,11 +262,9 @@ async fn rejects_an_empty_or_malformed_message_list( )) .await .unwrap_err(); - if let Some(expected) = expected { - assert_eq!(error, Error::InvalidRequest(expected.to_string().into())); - } else { - assert!(matches!(error, Error::InvalidRequest(_))); - } + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(reason), "{report}"); } #[rstest] @@ -355,7 +354,7 @@ async fn rejects_a_forwarded_header_the_signer_computes(#[case] forwarded: &str) let error = complete(call) .await .expect_err("conflicting signing headers fail"); - assert!(matches!(error, Error::Unsupported(_)), "{error:?}"); + assert!(error.declined().is_some(), "{error:?}"); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index e310a0be9b1..694d9b4de31 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -26,6 +26,7 @@ serde_json = { workspace = true, features = ["preserve_order"] } tracing.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-framer = { workspace = true, features = ["sse"] } litellm-cache-memory.workspace = true litellm-host-native.workspace = true diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index 9be7ee6ffc3..a57eca1c6a9 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -1,6 +1,5 @@ use litellm_core_utils::get_llm_provider_logic::LlmProviders; use litellm_http::request::string_headers as shared_string_headers; -pub(super) use litellm_http::request::truncate_error_body; use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG; use litellm_llms::{ anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, @@ -77,7 +76,7 @@ mod tests { use rstest::rstest; - use super::{messages_provider, string_headers, truncate_error_body}; + use super::{messages_provider, string_headers}; use crate::Error; use litellm_core_utils::get_llm_provider_logic::LlmProviders; @@ -106,19 +105,6 @@ mod tests { assert!(messages_provider(provider, model).is_none()); } - #[test] - fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(400); - let truncated = truncate_error_body(&body); - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); - } - #[test] fn string_headers_rejects_non_string_values() { let headers = json!({"x-count": 3}).as_object().unwrap().clone(); diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 55c80d9bddf..fa802d20587 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -10,7 +10,7 @@ use litellm_core_utils::{ use litellm_host::interceptors::{ Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, WireRequest, }; -use litellm_http::{request::with_default_headers, transport::Error as TransportError}; +use litellm_http::request::with_default_headers; use litellm_llms::base_llm::{ auth::{Authenticated, ValidatedEnvironment, resolve_auth}, messages::{ @@ -26,7 +26,7 @@ use serde_json::Value; use super::{ Error, MessagesCall, MessagesCallResponse, MessagesRoute, - common_utils::{MessagesProvider, string_headers, truncate_error_body}, + common_utils::{MessagesProvider, string_headers}, constants::MESSAGES_TIMEOUT_SECS, route::Messages, types::invalid_request, @@ -188,7 +188,7 @@ async fn provider_call( None | Some(Value::Null) => false, Some(Value::Bool(stream)) => *stream, Some(value) => { - return Err(Error::InvalidRequest( + return Err(Error::invalid_request( litellm_llms::ErrorDetail::InvalidValue { field: "stream", expected: "a boolean", @@ -261,7 +261,7 @@ async fn call_provider( provider_name, )); } - let text = response.text().await.map_err(network)?; + let text = response.text().await.map_err(litellm_http::Error::from)?; log_response_body(&text); context.response_received(&text).await?; decode_response(config, &identity.model, &text) @@ -269,16 +269,12 @@ async fn call_provider( } fn serialize_failure(err: serde_json::Error) -> Error { - Error::InvalidRequest(litellm_llms::ErrorDetail::failed( + Error::invalid_request(litellm_llms::ErrorDetail::failed( "Anthropic messages request serialization", err, )) } -fn network(error: reqwest::Error) -> Error { - Error::Transport(TransportError::Network(error.to_string())) -} - async fn send( http: &litellm_http::Client, authenticated: Authenticated, @@ -288,13 +284,11 @@ async fn send( ) -> Result { let request = outbound_request( authenticated, - url.to_string(), + url, body, Some(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), )?; - litellm_inference::outbound::send(request, http) - .await - .map_err(network) + Ok(litellm_inference::outbound::send(request, http).await?) } async fn provider_error(response: reqwest::Response) -> Error { @@ -302,12 +296,9 @@ async fn provider_error(response: reqwest::Response) -> Error { match response.text().await { Ok(text) => { log_error_body(status, &text); - Error::Transport(TransportError::Http { - status, - body: truncate_error_body(&text), - }) + litellm_http::Error::upstream(status, text).into() } - Err(error) => network(error), + Err(error) => litellm_http::Error::from(error).into(), } } @@ -317,7 +308,7 @@ fn decode_response( text: &str, ) -> Result { let response = serde_json::from_str(text).map_err(|err| { - Error::InvalidResponse(litellm_llms::ErrorDetail::invalid( + Error::invalid_response(litellm_llms::ErrorDetail::invalid( "messages response JSON", err, )) @@ -325,6 +316,7 @@ fn decode_response( config .transform_anthropic_messages_response(model, response) .map_err(Error::from) + .map_err(Error::after_send) } fn streaming_response( @@ -339,7 +331,7 @@ fn streaming_response( .collect(); let chunks = match decoder { None => futures_util::stream::try_unfold(response, move |mut response| async move { - let chunk = response.chunk().await.map_err(network)?; + let chunk = response.chunk().await.map_err(litellm_http::Error::from)?; Ok(chunk.map(|chunk| { log_chunk(provider, "provider_response", &chunk); (chunk, response) @@ -362,7 +354,7 @@ fn decoded_chunks( let bytes: ByteStream = response .bytes_stream() .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)) - .map_err(std::io::Error::other) + .map_err(|error| std::io::Error::other(litellm_http::Error::from(error))) .boxed(); futures_util::stream::try_unfold(decode(bytes), move |mut events| async move { let Some(event) = events.try_next().await? else { @@ -541,14 +533,10 @@ mod tests { shaping, }) .await; - assert!( - matches!( - result, - Err(Error::Auth(litellm_auth::Error::MissingApiKey { .. })) - ), - "{:?}", - result.err() - ); + let Err(error) = result else { + panic!("a call without credentials fails"); + }; + assert!(error.is_missing_credential(), "{error:?}"); } #[rstest] @@ -771,19 +759,18 @@ mod tests { }, ..shaping }; - assert!(matches!( - wire_body( - json!({ - "model": "claude-test", - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": 16, - "metadata": {"user_id": 123} - }), - shaping, - ) - .await, - Err(Error::InvalidRequest(_)) - )); + let error = wire_body( + json!({ + "model": "claude-test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "metadata": {"user_id": 123} + }), + shaping, + ) + .await + .unwrap_err(); + assert!(error.is_request(), "{error:?}"); } #[rstest] @@ -800,7 +787,7 @@ mod tests { shaping, ) .await, - Err(Error::InvalidRequest( + Err(Error::invalid_request( litellm_llms::ErrorDetail::InvalidValue { field: "metadata.user_id", expected: "a string", @@ -859,12 +846,12 @@ mod tests { headers: call.wire.headers, signer: call.signer, }, - call.wire.url, + &call.wire.url, &call.wire.body, None, ) .unwrap(); - let authorization = signed.header("authorization").unwrap_or_default(); + let authorization = signed.headers()["authorization"].to_str().unwrap(); assert!( authorization.contains(&format!("/us-west-2/{service}/aws4_request")), "{authorization}" @@ -888,7 +875,7 @@ mod tests { .respond_with(ResponseTemplate::new(200).set_body_raw(body, "text/event-stream")) .mount(&upstream) .await; - let response = litellm_http::Client::plain_for_test() + let response = litellm_http::Client::for_test() .get(upstream.uri()) .send() .await @@ -902,7 +889,7 @@ mod tests { let chunk = chunks.next().await.unwrap(); match expected { Some(expected) => assert_eq!(chunk.unwrap().as_ref(), expected.as_bytes()), - None => assert!(matches!(chunk, Err(Error::InvalidResponse(_))), "{chunk:?}"), + None => assert!(chunk.unwrap_err().is_response()), } assert!(chunks.next().await.is_none()); } diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index 21eac0574ed..ae346259559 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -34,11 +34,11 @@ pub fn messages_body(body: Map) -> Result /// Python's `GenericLiteLLMParams(**kwargs)`. pub fn litellm_params(fields: Map) -> Result { serde_json::from_value(Value::Object(fields)) - .map_err(|err| Error::InvalidRequest(ErrorDetail::invalid("litellm params", err))) + .map_err(|err| Error::invalid_request(ErrorDetail::invalid("litellm params", err))) } pub(super) fn invalid_request(err: serde_json::Error) -> Error { - Error::InvalidRequest(ErrorDetail::invalid("Anthropic messages request", err)) + Error::invalid_request(ErrorDetail::invalid("Anthropic messages request", err)) } pub type MessagesCallResponse = diff --git a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs index 5420477c34a..f91d5c79f7c 100644 --- a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs +++ b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs @@ -309,7 +309,14 @@ async fn deployment_workspace_is_required_and_sent_as_a_header( }) .await; if !configured { - assert!(matches!(result, Err(Error::Auth(_)))); + let Err(error) = result else { + panic!("the call fails"); + }; + assert!(!error.dispatched(), "{error:?}"); + assert!( + error.is_request() || error.is_missing_credential(), + "{error:?}" + ); assert!(received(&upstream).await.is_empty()); return; } diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index ba20fb9eec4..c3d06027c32 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use litellm_host::lifecycle::ExecutionEvent; use std::sync::Mutex; @@ -23,6 +24,7 @@ struct RecordingHost { optional_params: Mutex>, facts: Mutex>, reject_result: bool, + reject_response: bool, } impl RecordingHost { @@ -34,6 +36,7 @@ impl RecordingHost { optional_params: Mutex::new(Vec::new()), facts: Mutex::new(Vec::new()), reject_result: false, + reject_response: false, } } @@ -81,7 +84,7 @@ impl litellm_host::interceptors::Interceptors< Result<(), Error> { self.facts.lock().unwrap().push(facts); if self.reject_result { - return Err(Error::Unsupported("result rejected")); + return Err(Error::Declined(Decline::new("result rejected"))); } Ok(()) } @@ -101,6 +104,9 @@ impl litellm_host::interceptors::Interceptors< Result<(), ::Error> { + if self.reject_response { + return Err(Error::Declined(Decline::new("response rejected"))); + } litellm_host::lifecycle::CallObserver::observe( self, litellm_host::lifecycle::CallEvent::Execution( @@ -121,6 +127,7 @@ async fn rejected_results_are_not_delivered_or_cached( call: MessagesCall, #[case] streaming: bool, #[case] hosted: bool, + #[values(false, true)] cached: bool, ) { use futures_util::TryStreamExt; use litellm_cache_memory::InMemoryCache; @@ -142,13 +149,8 @@ async fn rejected_results_are_not_delivered_or_cached( )))), CacheScope::Shared, )); - for (reject, expected_requests, cached) in [ - (true, 1, false), - (false, 2, false), - (true, 2, true), - (false, 2, true), - ] { - let request = authenticated( + let request = || { + authenticated( with_fields( MessagesCall { body: call.body.clone(), @@ -157,36 +159,41 @@ async fn rejected_results_are_not_delivered_or_cached( json!({"stream": streaming}), ), upstream.uri(), - ); - let host = RecordingHost { - reject_result: reject, - ..RecordingHost::passthrough(request) - }; - let result = if hosted { - litellm_host_native::in_process::run_hosted( - route.clone().machine(host.request().unwrap(), None), - host.runtime(), - ) + ) + }; + let host = RecordingHost { + reject_result: true, + ..RecordingHost::passthrough(request()) + }; + if cached { + let output = route.execute(request(), &(), None).await.unwrap(); + if let MessagesCallResponse::Stream { chunks, .. } = output { + chunks.try_collect::>().await.unwrap(); + } + } + let result = if hosted { + litellm_host_native::in_process::run_hosted( + route.clone().machine(host.request().unwrap(), None), + host.runtime(), + ) + .await + .map(|_| ()) + } else { + route + .execute(host.request().unwrap(), &host, None) .await .map(|_| ()) - } else { - match route.execute(host.request().unwrap(), &host, None).await { - Ok(MessagesCallResponse::Complete(_)) => Ok(()), - Ok(MessagesCallResponse::Stream { chunks, .. }) => { - chunks.try_collect::>().await.map(|_| ()) - } - Err(error) => Err(error), - } - }; - assert_eq!( - result, - if reject { - Err(Error::Unsupported("result rejected")) - } else { - Ok(()) - } - ); - assert_eq!(received(&upstream).await.len(), expected_requests); + }; + let error = result.unwrap_err(); + assert_eq!(error.declined().is_some(), cached); + assert_eq!(error.dispatched(), !cached); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("result rejected") + ); + assert_eq!(received(&upstream).await.len(), 1); + { let facts = host.facts.lock().unwrap(); assert_eq!(facts.len(), 1); assert_eq!( @@ -194,6 +201,46 @@ async fn rejected_results_are_not_delivered_or_cached( cached ); } + let output = route.execute(request(), &(), None).await.unwrap(); + if let MessagesCallResponse::Stream { chunks, .. } = output { + chunks.try_collect::>().await.unwrap(); + } + assert_eq!(received(&upstream).await.len(), if cached { 1 } else { 2 }); +} + +#[rstest] +#[case::native(false)] +#[case::hosted(true)] +#[tokio::test] +async fn response_hook_declines_are_terminal(call: MessagesCall, #[case] hosted: bool) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost { + reject_response: true, + ..RecordingHost::passthrough(authenticated(call, upstream.uri())) + }; + let route = messages_route(no_secrets()); + let result = if hosted { + litellm_host_native::in_process::run_hosted( + route.machine(host.request().unwrap(), None), + host.runtime(), + ) + .await + .map(|_| ()) + } else { + route + .execute(host.request().unwrap(), &host, None) + .await + .map(|_| ()) + }; + let error = result.unwrap_err(); + assert!(error.declined().is_none()); + assert!(error.dispatched()); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("response rejected") + ); + assert_eq!(received(&upstream).await.len(), 1); } async fn run_through(host: &RecordingHost) -> Result { @@ -301,7 +348,7 @@ async fn response_mode_follows_the_intercepted_request( let summaries = traces.summaries("litellm.route"); assert_eq!(summaries.len(), 1); let Some(expected_stream) = expected_stream else { - assert!(matches!(result, Err(Error::InvalidRequest(_)))); + assert!(result.unwrap_err().is_request()); assert!(received(&upstream).await.is_empty()); assert_eq!(summaries[0]["outcome"], "failure"); return; @@ -322,14 +369,17 @@ async fn a_before_send_failure_never_sends(call: MessagesCall) { let upstream = upstream([message_response()]).await; let host = RecordingHost::new( authenticated(call, upstream.uri()), - Box::new(|_| Err(Error::InvalidRequest("vetoed by the host".into()))), + Box::new(|_| Err(litellm_llms::Error::MissingField("vetoed by the host").into())), ); let error = run_through(&host) .await .expect_err("the host failure fails the call"); - assert_eq!(error, Error::InvalidRequest("vetoed by the host".into())); + assert_eq!( + error, + Error::from(litellm_llms::Error::MissingField("vetoed by the host")) + ); assert!(received(&upstream).await.is_empty()); assert!(host.raw_responses().is_empty()); } diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index fcc432acb76..a01e571ec75 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -144,7 +144,7 @@ impl LocalMessagesHost { .lock() .unwrap_or_else(|error| error.into_inner()) .take() - .ok_or_else(|| Error::InvalidRequest("messages request was already projected".into())) + .ok_or(litellm_llms::Error::MissingField("messages request").into()) } pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, (), Self, ()> { litellm_host_native::in_process::Host { diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 93bbbb071ba..e53bdb8d252 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -108,13 +108,7 @@ async fn a_call_without_credentials_fails_before_sending( .await .expect_err("a call without credentials fails"); - assert!( - matches!( - error, - Error::Auth(litellm_auth::Error::MissingApiKey { .. }) - ), - "{error:?}" - ); + assert!(error.is_missing_credential(), "{error:?}"); assert!(received(&upstream).await.is_empty()); } @@ -185,7 +179,9 @@ async fn unsupported_providers_are_rejected_before_sending( .await .expect_err("unsupported provider errors"); - assert_eq!(error, Error::InvalidProvider(reported.into())); + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(reported), "{report}"); } #[rstest] @@ -513,10 +509,9 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i let error = run(shaped(false)) .await .expect_err("an unsupported param is rejected without drop_params"); - assert!( - matches!(&error, Error::InvalidRequest(message) if message.to_string().contains(rejected_as)), - "{error:?}" - ); + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(rejected_as), "{report}"); assert!(received(&upstream).await.is_empty()); run_message(shaped(true)).await; @@ -957,11 +952,10 @@ async fn provider_validation_runs_before_caller_parameter_removal( )) .await; let error = result.expect_err("metadata is validated before removal"); + let report = litellm_error::Report::new(&error).to_string(); assert!( - error - .to_string() - .contains("metadata.user_id must be a string"), - "{error}" + report.contains("metadata.user_id must be a string"), + "{report}" ); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index b5d623a5582..7aff1286858 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -2,7 +2,6 @@ use litellm_host::{ interceptors::{ExecutionFacts, ResultSource}, lifecycle::ExecutionEvent, }; -use litellm_http::transport::Error as TransportError; use litellm_inference_messages::{MessagesCallResponse, messages_body}; use litellm_inference_testing::{ RecordingSecrets, http_config, no_secrets, provider_http, resources, @@ -166,17 +165,21 @@ async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) { .await .expect_err("upstream error propagates"); - let Error::Transport(TransportError::Http { status, body }) = error else { + let Some((status, body)) = error.upstream_response() else { panic!("{error:?}"); }; assert_eq!(status, 400); - assert_eq!(serde_json::from_str::(&body).unwrap(), envelope); + assert_eq!(serde_json::from_str::(body).unwrap(), envelope); } #[rstest] +#[case::ascii(format!("{}private-tail", "x".repeat(600)))] +#[case::unicode(format!("{}private-tail", "é".repeat(600)))] #[tokio::test] -async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall) { - let long = "x".repeat(600); +async fn a_long_error_body_is_preserved_with_bounded_diagnostics( + call: MessagesCall, + #[case] long: String, +) { let upstream = upstream([ResponseTemplate::new(500).set_body_string(long.clone())]).await; let error = run(MessagesCall { @@ -187,13 +190,13 @@ async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall .await .expect_err("upstream error propagates"); - assert_eq!( - error, - Error::Transport(TransportError::Http { - status: 500, - body: format!("{}... (truncated)", &long[..256]) - }) - ); + assert_eq!(error.upstream_response(), Some((500, long.as_str()))); + let diagnostic_body: String = long.chars().take(256).collect(); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.ends_with(&format!("{diagnostic_body}... (truncated)"))); + assert_eq!(report.matches(&diagnostic_body).count(), 1, "{report}"); + assert!(!report.contains("private-tail"), "{report}"); + assert!(!format!("{error:?}").contains("private-tail")); } #[rstest] @@ -216,11 +219,8 @@ async fn an_upstream_error_keeps_its_status_and_body(call: MessagesCall, #[case] .expect_err("upstream error propagates"); assert_eq!( - error, - Error::Transport(TransportError::Http { - status, - body: "upstream said no".into() - }) + error.upstream_response(), + Some((status, "upstream said no")) ); } @@ -242,7 +242,7 @@ async fn an_unreadable_success_body_is_an_invalid_response( .await .expect_err("an unreadable body fails"); - assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!(error.is_response(), "{error:?}"); } #[rstest] @@ -259,7 +259,23 @@ async fn a_provider_slower_than_the_timeout_fails_the_call(call: MessagesCall) { .await .expect_err("the call times out"); - assert!(matches!(error, Error::Transport(_)), "{error:?}"); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); +} + +#[rstest] +#[tokio::test] +async fn a_transport_error_does_not_echo_the_provider_url(call: MessagesCall) { + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(format!("{UNREACHABLE_BASE}/tenant-token-1234")), + ..call + }) + .await + .expect_err("nothing is listening"); + + assert!(!error.to_string().contains("tenant-token-1234"), "{error}"); } #[rstest] @@ -305,8 +321,11 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes fn a_body_that_does_not_parse_is_an_invalid_request(#[case] raw: Value) { let error = messages_body(object(raw)).expect_err("the body is rejected"); + assert!(error.is_request(), "{error:?}"); assert!( - matches!(&error, Error::InvalidRequest(message) if message.to_string().starts_with("invalid Anthropic messages request: ")), + litellm_error::Report::new(&error) + .to_string() + .contains("invalid Anthropic messages request"), "{error:?}" ); } diff --git a/litellm-rust/crates/inference-messages/tests/messages/secrets.rs b/litellm-rust/crates/inference-messages/tests/messages/secrets.rs index d4f7d2bf3a7..b4a3018ab93 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/secrets.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/secrets.rs @@ -109,7 +109,7 @@ async fn a_secret_manager_failure_fails_the_call_before_sending(call: MessagesCa .expect_err("a secret manager failure fails the call"); assert!( - matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)), + matches!(&error, Error::Secret(_)) && !error.is_request(), "{error:?}" ); assert!(received(&upstream).await.is_empty()); diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index cda688107b5..4a8a061f764 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -260,13 +260,7 @@ async fn an_upstream_error_fails_the_call_without_opening_the_stream( .await .expect_err("upstream error propagates"); - assert_eq!( - error, - Error::Transport(litellm_http::transport::Error::Http { - status: 429, - body: body.into() - }) - ); + assert_eq!(error.upstream_response(), Some((429, body))); assert!(host.seen.into_inner().unwrap().is_empty()); } @@ -332,7 +326,9 @@ async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { .expect("the stalled stream gives up within the timeout") .expect_err("a stalled body fails the call"); - assert!(matches!(error, Error::Transport(_)), "{error:?}"); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); let seen = host.seen.into_inner().unwrap(); assert!( matches!(seen.as_slice(), [Seen::Open(_), Seen::Deliver(chunk)] if chunk.as_ref() == b"event: message_start\ndata: {}\n\n"), @@ -390,13 +386,7 @@ async fn the_sdk_returns_http_errors_before_opening_a_stream(call: MessagesCall) .err() .expect("upstream failure is returned by messages()"); - assert_eq!( - error, - Error::Transport(litellm_http::transport::Error::Http { - status: 429, - body: "slow down".into(), - }) - ); + assert_eq!(error.upstream_response(), Some((429, "slow down"))); } #[rstest] @@ -470,7 +460,9 @@ async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesC .expect("the stalled body times out") .unwrap() .unwrap_err(); - assert!(matches!(error, Error::Transport(_)), "{error:?}"); + assert!(error.dispatched(), "{error:?}"); + assert!(!error.is_request(), "{error:?}"); + assert_eq!(error.upstream_response(), None); assert!(chunks.next().await.is_none()); tokio::time::timeout(Duration::from_secs(5), connection) .await diff --git a/litellm-rust/crates/inference-messages/tests/messages/tencent.rs b/litellm-rust/crates/inference-messages/tests/messages/tencent.rs index dce187f4ea2..4bab7de65e1 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/tencent.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/tencent.rs @@ -79,13 +79,9 @@ async fn missing_credentials_fail_before_sending(tencent_call: MessagesCall) { }) .await .unwrap_err(); - assert!(matches!( - error, - Error::Auth(litellm_auth::Error::MissingApiKey { - environment_variable: "TENCENT_API_KEY", - .. - }) - )); + assert!(error.is_missing_credential(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("TENCENT_API_KEY"), "{report}"); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-messages/tests/openrouter/main.rs b/litellm-rust/crates/inference-messages/tests/openrouter/main.rs index 74444142d37..80243718768 100644 --- a/litellm-rust/crates/inference-messages/tests/openrouter/main.rs +++ b/litellm-rust/crates/inference-messages/tests/openrouter/main.rs @@ -1,7 +1,7 @@ use std::{sync::Arc, time::Duration}; use futures_util::TryStreamExt; -use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping}; +use litellm_inference_messages::{MessagesCall, MessagesCallResponse, MessagesShaping}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; use rstest::{fixture, rstest}; use serde_json::json; @@ -92,13 +92,12 @@ async fn an_anthropic_key_alone_is_not_an_openrouter_credential(call: MessagesCa None, ) .await; - assert!(matches!( - result, - Err(Error::Auth(litellm_auth::Error::MissingApiKey { - provider: "OpenRouter", - .. - })) - )); + let Err(error) = result else { + panic!("the call fails"); + }; + assert!(error.is_missing_credential(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("OpenRouter"), "{report}"); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-ocr/Cargo.toml b/litellm-rust/crates/inference-ocr/Cargo.toml index 21ec013f873..bbbdaddb050 100644 --- a/litellm-rust/crates/inference-ocr/Cargo.toml +++ b/litellm-rust/crates/inference-ocr/Cargo.toml @@ -25,6 +25,7 @@ tokio.workspace = true tracing.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-auth-aws.workspace = true litellm-auth-gcp.workspace = true litellm-host-native.workspace = true @@ -34,4 +35,5 @@ litellm-llms = { workspace = true, features = ["test-support"] } rand.workspace = true rstest.workspace = true time.workspace = true +url.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/inference-ocr/src/document.rs b/litellm-rust/crates/inference-ocr/src/document.rs index 0d9a296d289..44c1049102a 100644 --- a/litellm-rust/crates/inference-ocr/src/document.rs +++ b/litellm-rust/crates/inference-ocr/src/document.rs @@ -46,9 +46,9 @@ pub fn read_path_document(path: &Path, mime_type: Option<&str>) -> Result OcrClient { - OcrClient::for_test( - litellm_http::Client::plain_for_test(), - litellm_http::Client::no_redirect_for_test(), - ) - .with_secrets(litellm_inference_testing::no_secrets()) + OcrClient::for_test().with_secrets(litellm_inference_testing::no_secrets()) } fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest { diff --git a/litellm-rust/crates/inference-ocr/src/provider_config.rs b/litellm-rust/crates/inference-ocr/src/provider_config.rs index 09041e9b72b..19e50af0dd1 100644 --- a/litellm-rust/crates/inference-ocr/src/provider_config.rs +++ b/litellm-rust/crates/inference-ocr/src/provider_config.rs @@ -1,6 +1,8 @@ use litellm_auth::{InputSource, SecretValue, Sourced}; -use litellm_core_utils::get_llm_provider_logic::{LlmProviders, get_custom_llm_provider}; -use litellm_inference::provider::{ResolvedProvider, resolve_llm_provider}; +use litellm_core_utils::get_llm_provider_logic::{ + CustomLlmProvider, LlmProviders, get_custom_llm_provider, +}; +use litellm_inference::provider::ResolvedProvider; use litellm_llms::{ aws_textract::ocr::{ analyze_transformation::TextractAnalyzeDocumentConfig, common_utils::TextractOperation, @@ -241,18 +243,20 @@ pub(crate) fn resolve_provider_config<'a>( model: &'a str, custom_llm_provider: Option<&'a str>, ) -> Result<(ResolvedProvider<'a>, OcrConfigKind), Error> { - let provider = match resolve_llm_provider(model, custom_llm_provider, "ocr") { - Ok(provider) => provider, - Err(_) if get_custom_llm_provider(model, custom_llm_provider).is_none() => { - ResolvedProvider { - model, - provider: LlmProviders::Mistral, - } - } - Err(litellm_inference::Error::InvalidProvider(provider)) => { - return Err(Error::InvalidProvider(provider)); - } - Err(error) => return Err(Error::InvalidProvider(error.to_string())), + let provider = match get_custom_llm_provider(model, custom_llm_provider) { + None => ResolvedProvider { + model, + provider: LlmProviders::Mistral, + }, + Some(CustomLlmProvider { + model, + custom_llm_provider, + }) => ResolvedProvider { + model, + provider: custom_llm_provider + .parse() + .map_err(|_| Error::InvalidProvider(custom_llm_provider.to_owned()))?, + }, }; let config = match provider.provider { LlmProviders::AwsTextract => match TextractOperation::from_model(provider.model)? { diff --git a/litellm-rust/crates/inference-ocr/src/route.rs b/litellm-rust/crates/inference-ocr/src/route.rs index 85125452bda..74ff829e8ad 100644 --- a/litellm-rust/crates/inference-ocr/src/route.rs +++ b/litellm-rust/crates/inference-ocr/src/route.rs @@ -43,7 +43,10 @@ fn caller_token_provider(services: HostServices) -> TokenProviderHandle { .call(OcrOp::AcquireAzureAdToken) .await .map_err(|error| { - litellm_auth::Error::CredentialAcquisition(error.to_string().into()) + litellm_auth::Error::CredentialAcquisition(litellm_auth::ErrorDetail::failed( + "host Azure AD token request", + error, + )) }) } }) diff --git a/litellm-rust/crates/inference-ocr/src/types.rs b/litellm-rust/crates/inference-ocr/src/types.rs index 1bd70c06116..98c609a158f 100644 --- a/litellm-rust/crates/inference-ocr/src/types.rs +++ b/litellm-rust/crates/inference-ocr/src/types.rs @@ -67,6 +67,7 @@ impl OcrConnectionInputs { .map(|value| (name.clone(), value.to_string())) .ok_or_else(|| Error::RequestField { path: format!("extra_headers.{name}"), + cause: None, }) }) .collect() @@ -105,6 +106,7 @@ impl LiteLLMOcrRequest { .filter(|value| *value > 0 && *value <= default_transport.max_response_bytes) .ok_or_else(|| Error::RequestField { path: "max_response_bytes".into(), + cause: None, }) }) .transpose()? @@ -313,7 +315,7 @@ mod tests { }; assert!(matches!( error, - Error::RequestField { ref path } if path == "extra_headers.x-a" + Error::RequestField { ref path, .. } if path == "extra_headers.x-a" )); } } diff --git a/litellm-rust/crates/inference-ocr/src/wire.rs b/litellm-rust/crates/inference-ocr/src/wire.rs index a8f5a77a1cd..56d22be6523 100644 --- a/litellm-rust/crates/inference-ocr/src/wire.rs +++ b/litellm-rust/crates/inference-ocr/src/wire.rs @@ -73,9 +73,8 @@ pub fn decode_request_input>( let timeout = wire .timeout_seconds .map(|seconds| { - Duration::try_from_secs_f64(seconds).map_err(|_| Error::RequestField { - path: "timeout_seconds".into(), - }) + Duration::try_from_secs_f64(seconds) + .map_err(|error| Error::request_field("timeout_seconds", error)) }) .transpose()?; LiteLLMOcrRequest::from_inputs( @@ -123,8 +122,8 @@ mod tests { #[case::non_object(json!([]), "document")] #[case::missing_type(json!({"document_url":"https://example.com/a.pdf"}), "document")] #[case::unsupported_type(json!({"type":"text"}), "type")] - #[case::missing_document_url(json!({"type":"document_url"}), "Document URL")] - #[case::missing_image_url(json!({"type":"image_url"}), "Document URL")] + #[case::missing_document_url(json!({"type":"document_url"}), "document URL")] + #[case::missing_image_url(json!({"type":"image_url"}), "document URL")] fn ocr_contract_malformed_document_is_bad_request( #[case] document: Value, #[case] field: &str, @@ -191,16 +190,13 @@ mod tests { assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); } - #[test] - fn missing_document_source_has_a_typed_public_error() { - for document in [ - serde_json::json!({"type": "document_url"}), - serde_json::json!({"type": "image_url"}), - ] { - assert!(matches!( - decode_document(document), - Err(Error::MissingDocumentUrl) - )); - } + #[rstest] + #[case::document_url(serde_json::json!({"type": "document_url"}))] + #[case::image_url(serde_json::json!({"type": "image_url"}))] + fn missing_document_source_has_a_typed_public_error(#[case] document: serde_json::Value) { + assert!(matches!( + decode_document(document), + Err(Error::MissingDocumentUrl) + )); } } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs index 5336d67895d..f0e72b19fc1 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs @@ -1,9 +1,12 @@ -use std::{collections::BTreeMap, time::SystemTime}; +use std::time::SystemTime; -use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; +use litellm_auth_aws::{Credentials, SigningScope, aws_signature_headers, sign_headers}; use rstest::rstest; use time::{PrimitiveDateTime, format_description}; -use wiremock::Request; +use wiremock::{ + Request, + http::{HeaderMap, HeaderName}, +}; use super::*; @@ -42,21 +45,26 @@ fn expected_authorization(url: &str, sent: &Request) -> String { .unwrap() .assume_utc() .into(); - let headers: BTreeMap = ["content-type", "x-amz-target"] + let headers: HeaderMap = ["content-type", "x-amz-target"] .into_iter() - .map(|name| (name.to_string(), sent.header(name).unwrap().to_string())) + .map(|name| (HeaderName::from_static(name), sent.headers[name].clone())) .collect(); - sign_post( + sign_headers( + "POST", url, &sent.body, &aws_signature_headers(&headers), - "eu-west-1", - "textract", - &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), - signed_at, + SigningScope { + region: "eu-west-1", + service: "textract", + credentials: &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + time: signed_at, + }, ) - .unwrap()["Authorization"] - .clone() + .unwrap()["authorization"] + .to_str() + .unwrap() + .to_string() } /// The recorded URL names wiremock's host, not the address the client signed for. @@ -183,12 +191,15 @@ async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit(#[ .await .unwrap_err(); - let Error::Provider { status, body, .. } = error else { - panic!("expected a provider error, got {error:?}"); - }; - assert_eq!(status, 400); + let message = litellm_error::Report::new(&error).to_string(); + assert!(error.is_request()); + assert_eq!(error.http_status_code(), Some(400)); assert!( - body.contains("multi-page documents are not supported"), - "{body}" + message.contains("Request has unsupported document format"), + "{message}" + ); + assert!( + message.contains("multi-page documents are not supported"), + "{message}" ); } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs index 9bff7490d58..6f591af73e0 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs @@ -222,24 +222,23 @@ async fn credential_precedence( false, json!({}), numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { - provider: "Azure AI", - guidance: "Set AZURE_AI_API_BASE environment variable or pass api_base parameter", - })), + |error: &Error| litellm_error::Report::new(error).to_string() + == "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter", 0 )] #[case::unsupported_oidc_reference( true, json!({"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}), numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::InvalidConfiguration(detail)) if detail.to_string() == "unsupported OIDC reference"), + |error: &Error| litellm_error::Report::new(error).to_string().contains("unsupported OIDC reference"), 0 )] #[case::empty_provider_token_ignores_static_token( true, json!({"azure_ad_token": "static-token"}), |_| String::new(), - |error: &Error| matches!(error, Error::MissingAzureAiCredentials), + |error: &Error| error.is_missing_credential() + && litellm_error::Report::new(error).to_string().contains("Azure AI credentials"), 1 )] #[tokio::test] diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs index cec43e2ed0e..98875cc777c 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs @@ -5,7 +5,7 @@ use std::{ use litellm_host::lifecycle::CallEvent; use litellm_host::lifecycle::ExecutionEvent; -use litellm_llms::base_llm::ocr::settings::OcrSettings; +use litellm_llms::base_llm::ocr::{error::PagesFault, settings::OcrSettings}; use rstest::rstest; @@ -64,10 +64,10 @@ async fn pages_features_and_extra_options_map_to_the_analyze_call() { } #[rstest] -#[case(json!({"pages": [true]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages": [1, "2"]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages": [-1]}), Error::Pages("negative page index".into()))] -#[case(json!({"pages": "1&&features=bad"}), Error::Pages("invalid native page range".into()))] +#[case(json!({"pages": [true]}), Error::Pages(PagesFault::Mixed))] +#[case(json!({"pages": [1, "2"]}), Error::Pages(PagesFault::Mixed))] +#[case(json!({"pages": [-1]}), Error::Pages(PagesFault::Negative))] +#[case(json!({"pages": "1&&features=bad"}), Error::Pages(PagesFault::Range))] #[case(json!({"features": "languages&pages=1"}), Error::Features)] #[case(json!({"req_format": "azure"}), Error::RequestFormat)] #[tokio::test] @@ -315,7 +315,11 @@ async fn polling_does_not_follow_redirects() { .await .unwrap_err(); - assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!( + error.upstream_response().map(|(status, ..)| status), + Some(302), + "{error}" + ); assert_eq!(received(&upstream).await.len(), 2); } @@ -389,12 +393,15 @@ async fn malformed_provider_pages_report_the_response_path( } #[rstest] -#[case::missing(None)] -#[case::relative(Some("/relative"))] -#[case::cross_origin(Some("http://example.com/operation"))] -#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"))] +#[case::missing(None, false)] +#[case::relative(Some("/relative"), true)] +#[case::cross_origin(Some("http://example.com/operation"), false)] +#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"), false)] #[tokio::test] -async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&str>) { +async fn an_unusable_operation_location_is_rejected( + #[case] location: Option<&str>, + #[case] unparseable: bool, +) { let response = location .into_iter() .fold(ResponseTemplate::new(202), |response, location| { @@ -407,6 +414,11 @@ async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&st .unwrap_err(); assert!(error.to_string().contains("operation-location"), "{error}"); + assert!(error.is_response()); + assert_eq!( + litellm_error::find::(&error).is_some(), + unparseable + ); assert_eq!(received(&upstream).await.len(), 1); } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs b/litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs index 007aa49a2fd..a5076cf0112 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs @@ -37,6 +37,15 @@ async fn a_non_image_document_is_rejected_before_sending( .await .unwrap_err(); - assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); + assert!( + matches!( + error, + Error::UnsupportedDocument { + provider: "Cohere Parse", + .. + } + ), + "{error:?}" + ); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/documents.rs b/litellm-rust/crates/inference-ocr/tests/ocr/documents.rs index 338a47c5f57..35de1326386 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/documents.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/documents.rs @@ -176,8 +176,8 @@ async fn a_missing_path_document_fails_before_sending() { assert!( matches!( &error, - Error::FileRead { path: failed, source } - if *failed == path && source.kind() == std::io::ErrorKind::NotFound + Error::FileRead { path: failed, cause } + if *failed == path && cause.kind() == std::io::ErrorKind::NotFound ), "{error:?}" ); @@ -216,7 +216,8 @@ async fn configured_client_preserves_document_url_policy(#[case] allowed: bool) .await; if !allowed { - assert!(matches!(result, Err(Error::BlockedDocumentUrl))); + let error = result.unwrap_err(); + assert!(error.is_request(), "{error:?}"); assert!(received(&documents).await.is_empty()); assert!(received(&upstream).await.is_empty()); return; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs index e413781f1a3..75a91092daf 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs @@ -38,7 +38,7 @@ fn recording_host( .with_before_send(move |wire, _| { interceptions.fetch_add(1, Ordering::SeqCst); match block { - true => Err(Error::InvalidRequest("blocked".into())), + true => Err(Error::MissingField("blocked")), false => Ok(wire), } }) @@ -95,10 +95,7 @@ async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() { .await .unwrap_err(); - assert!( - matches!(&error, Error::InvalidRequest(message) if message == "blocked"), - "{error:?}" - ); + assert!(matches!(error, Error::MissingField("blocked")), "{error:?}"); assert_eq!(interceptions.load(Ordering::SeqCst), 1); assert_eq!(*events.lock().unwrap(), ["started", "failure"]); assert!(received(&upstream).await.is_empty()); diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs index 4f9266c4a8d..c7a5d58ad3d 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs @@ -41,11 +41,10 @@ async fn drive_until( Err(error) => break Err(error), }; let answer = match op { - HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => host - .result_ready(facts) - .await - .map(|()| reply.send(())) - .map_err(HostFailure::Error), + HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => { + reply.send(host.result_ready(facts).await); + Ok(()) + } HostRequest::Stream(stream) => match stream { litellm_host::protocol::StreamDelivery::Open(head, _) => match head {}, litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {}, @@ -64,10 +63,8 @@ async fn drive_until( } HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { ops.push("response"); - host.after_provider_response(raw) - .await - .map(|()| reply.send(())) - .map_err(HostFailure::Error) + reply.send(host.after_provider_response(raw).await); + Ok(()) } }; if let Err(failure) = answer { @@ -85,10 +82,10 @@ async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, sto _ = stop.notified() => break, step = machine.resume() => { match step.unwrap() { - MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::ResultReady { reply, .. })) => reply.send(()), + MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::ResultReady { reply, .. })) => reply.send(Ok(())), MachineStep::Suspended(HostRequest::HostCall(op)) => host.handle_host_call(op).await.unwrap(), MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, reply, .. })) => reply.send(*wire), - MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(()), + MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(Ok(())), MachineStep::Suspended(HostRequest::Stream(stream)) => match stream { litellm_host::protocol::StreamDelivery::Open(head, _) => match head {}, litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {}, @@ -117,10 +114,14 @@ async fn a_hand_driven_machine_performs_the_same_call() { assert_eq!(outcome.unwrap().pages[0].markdown, "native"); assert_eq!(received(&upstream).await.len(), 1); assert_eq!(ops, ["BeforeSend", "response"]); - assert!(matches!( - machine.resume().await, - Err(Error::InvalidRequest(_)) - )); + let Err(error) = machine.resume().await else { + panic!("resuming a finished call is a host protocol fault"); + }; + assert!(!error.is_request(), "{error:?}"); + assert!( + litellm_error::find::(&error).is_some(), + "{error:?}" + ); } #[rstest::rstest] @@ -153,8 +154,11 @@ async fn a_path_document_is_read_by_core_without_a_host_operation() { } #[rstest] -#[case::failed(HostFailure::Error(Error::InvalidRequest("before_provider_request failed".into())), "before_provider_request failed")] -#[case::cancelled(HostFailure::Cancelled(Error::InvalidRequest("cancelled".into())), "cancelled")] +#[case::failed( + HostFailure::Error(Error::MissingField("before_provider_request failed")), + "before_provider_request failed" +)] +#[case::cancelled(HostFailure::Cancelled(Error::MissingField("cancelled")), "cancelled")] #[tokio::test] async fn a_before_send_failure_ends_the_call_without_reaching_transport( #[case] failure: HostFailure, @@ -174,7 +178,7 @@ async fn a_before_send_failure_ends_the_call_without_reaching_transport( .await; assert!( - matches!(&outcome, Err(Error::InvalidRequest(actual)) if actual == message), + matches!(&outcome, Err(Error::MissingField(actual)) if *actual == message), "{outcome:?}" ); assert_eq!(ops, ["BeforeSend"]); @@ -255,17 +259,17 @@ async fn interrupt_drops_provider_captures_before_returning() { drive_until_notified(&mut machine, &host, &entered).await; assert!(!dropped.load(Ordering::SeqCst)); - let acknowledgement = machine.interrupt(HostFailure::Cancelled(Error::InvalidRequest( - "cancelled".into(), - ))); + let acknowledgement = + machine.interrupt(HostFailure::Cancelled(Error::MissingField("cancelled"))); assert!( dropped.load(Ordering::SeqCst), "interrupt returned while provider captures were still alive" ); - assert!( - matches!(acknowledgement.await, Err(Error::InvalidRequest(message)) if message == "cancelled") - ); + assert!(matches!( + acknowledgement.await, + Err(Error::MissingField("cancelled")) + )); } #[rstest::rstest] @@ -290,7 +294,7 @@ async fn interrupting_an_in_flight_provider_request_closes_its_connection() { let mut machine = ocr_route().machine(host.request().unwrap(), None); drive_until_notified(&mut machine, &host, &received).await; - let cancelled = Error::InvalidRequest("cancelled".into()); + let cancelled = Error::MissingField("cancelled"); assert!( machine diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs index 5a905f907d2..5a48bba9e1a 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs @@ -209,7 +209,7 @@ impl LocalOcrHost { request, caller_token: false, }) - .ok_or_else(|| Error::InvalidRequest("OCR request was already projected".into())) + .ok_or(Error::MissingField("OCR request")) } pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, Self, Self, ()> { litellm_host_native::in_process::Host { @@ -223,11 +223,11 @@ impl LocalOcrHost { impl litellm_host_native::services::HostCallHandler for LocalOcrHost { async fn handle_host_call(&self, op: OcrOp) -> Result<(), Error> { match op { - OcrOp::AcquireAzureAdToken(_) => { - Err(Error::Auth(litellm_auth::Error::CredentialAcquisition( - "OCR host has no Azure AD token provider".into(), - ))) - } + OcrOp::AcquireAzureAdToken(_) => Err(Error::Auth( + litellm_auth::Error::CredentialAcquisition(litellm_auth::ErrorDetail::Unresolved { + subject: "OCR host Azure AD token provider", + }), + )), } } } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs b/litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs index 528adc7feaf..21cf4ce455d 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs @@ -92,28 +92,13 @@ async fn an_upstream_error_keeps_its_status_whole_body_and_headers( .unwrap_err(); assert_eq!(received(&upstream).await.len(), 1); - let Error::Provider { - status, - body, - headers, - } = error - else { - panic!("expected provider error, got {error:?}"); - }; + let (status, body, headers) = error.upstream_response().expect("the provider answered"); assert_eq!(status, 422); - for (name, value) in [ - ("retry-after", "17"), - ("x-request-id", "request-123"), - ("x-future-header", "retained"), - ] { - assert!( - headers - .iter() - .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value), - "{name} missing from {headers:?}" - ); - } assert_eq!(body, expected_body); + let header = |name: &str| headers.get(name).and_then(|value| value.to_str().ok()); + assert_eq!(header("retry-after"), Some("17")); + assert_eq!(header("x-request-id"), Some("request-123")); + assert_eq!(header("x-future-header"), Some("retained")); } #[rstest] diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs b/litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs index 649dca3d690..f6f458a4916 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs @@ -175,8 +175,11 @@ async fn a_failed_upload_stops_before_parse(#[case] upload: ResponseTemplate) { } #[rstest] -#[case::remote_url("https://example.com/a.pdf", Error::ReductoSource)] -#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into() })] +#[case::remote_url( + "https://example.com/a.pdf", + Error::UnsupportedDocument { provider: "Reducto", guidance: "" } +)] +#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into(), cause: None })] #[case::data_uri_without_payload("data:application/pdf;base64", Error::InvalidDataUri)] #[case::invalid_base64("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] #[tokio::test] @@ -202,7 +205,6 @@ async fn invalid_document_sources_are_rejected_before_sending( std::mem::discriminant(&expected) ); assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); } #[tokio::test] diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs index 81b52d076a5..dba12de60f6 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/vertex_ai.rs @@ -71,7 +71,8 @@ async fn invalid_credentials_fail_before_sending() { .await .unwrap_err(); - assert!(error.to_string().contains("vertex_credentials"), "{error}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("vertex_credentials"), "{report}"); } #[rstest] @@ -91,11 +92,10 @@ async fn a_request_controlled_api_base_is_rejected_before_vertex_auth( let error = perform(request).await.unwrap_err(); + let report = litellm_error::Report::new(&error).to_string(); assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint"), - "{error}" + report.contains("request-controlled Vertex AI endpoint"), + "{report}" ); } diff --git a/litellm-rust/crates/inference-responses/Cargo.toml b/litellm-rust/crates/inference-responses/Cargo.toml index 793483ee76d..083179c50c8 100644 --- a/litellm-rust/crates/inference-responses/Cargo.toml +++ b/litellm-rust/crates/inference-responses/Cargo.toml @@ -25,8 +25,10 @@ tokio-tungstenite.workspace = true tracing.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-cache.workspace = true litellm-cache-memory.workspace = true +litellm-error.workspace = true litellm-host-native.workspace = true litellm-inference-testing.workspace = true litellm-tracing.workspace = true diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index 02182d23ae3..19092a63a14 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -72,28 +72,29 @@ pub(super) async fn execute( let stream = match wire.body.get("stream") { None => false, Some(serde_json::Value::Bool(value)) => *value, - Some(_) => return Err(Error::InvalidRequest("stream must be a boolean".into())), + Some(_) => { + return Err(Error::invalid_request( + litellm_llms::ErrorDetail::Expected { + subject: "stream", + expected: "a boolean", + }, + )); + } }; let outbound = litellm_inference::outbound::outbound_request( Authenticated { headers: wire.headers, signer: authenticated.signer, }, - wire.url, + &wire.url, &wire.body, Some(call.timeout.unwrap_or(Duration::from_secs(600))), )?; - let response = litellm_inference::outbound::send(outbound, &route.http) - .await - .map_err(network)?; + let response = litellm_inference::outbound::send(outbound, &route.http).await?; let status = response.status().as_u16(); if !response.status().is_success() { - let body = response.text().await.map_err(network)?; - return Err(litellm_http::transport::Error::Http { - status, - body: litellm_http::request::truncate_error_body(&body), - } - .into()); + let body = response.text().await.map_err(litellm_http::Error::from)?; + return Err(litellm_http::Error::upstream(status, body).into()); } if stream { let headers = response @@ -105,14 +106,14 @@ pub(super) async fn execute( .collect(); let chunks = response .bytes_stream() - .map(|chunk| chunk.map_err(network)) + .map(|chunk| chunk.map_err(|error| litellm_http::Error::from(error).into())) .boxed(); return Ok(ResponsesOutput::Stream { head: ResponsesStreamHead { headers }, chunks, }); } - let body = response.text().await.map_err(network)?; + let body = response.text().await.map_err(litellm_http::Error::from)?; let raw = RawResponse { body: body.clone() }; if let Some(observers) = observers { observers.emit(litellm_host::lifecycle::CallEvent::Execution( @@ -123,17 +124,18 @@ pub(super) async fn execute( .after_provider_response(raw) .await .map_err(Error::post_call)?; - let value = serde_json::from_str(&body) - .map_err(|error| Error::InvalidResponse(error.to_string().into()))?; + let value = serde_json::from_str(&body).map_err(|error| { + Error::invalid_response(litellm_llms::ErrorDetail::invalid( + "Responses API response JSON", + error, + )) + })?; config .transform_response_api_response(value) .map(ResponsesOutput::Complete) .map_err(Error::from) + .map_err(Error::after_send) }, ) .await } - -fn network(error: reqwest::Error) -> Error { - litellm_http::transport::Error::Network(error.to_string()).into() -} diff --git a/litellm-rust/crates/inference-responses/src/lib.rs b/litellm-rust/crates/inference-responses/src/lib.rs index b6e151ecec3..9af5dea1fe8 100644 --- a/litellm-rust/crates/inference-responses/src/lib.rs +++ b/litellm-rust/crates/inference-responses/src/lib.rs @@ -11,6 +11,7 @@ mod handler; use std::sync::Arc; use litellm_auth::AuthServices; +use litellm_core_utils::decline::Decline; use litellm_core_utils::get_llm_provider_logic::{ CustomLlmProvider, LlmProviders, get_custom_llm_provider, }; @@ -24,6 +25,8 @@ use litellm_llms::{ use litellm_secrets::source::SecretSource; use types::{ResponsesCall, ResponsesOutput, ResponsesProvider}; +const NOT_NATIVE: Decline = Decline::new("native HTTP responses provider"); + fn responses_provider(provider: LlmProviders) -> Option<&'static dyn BaseResponsesApiConfig> { match provider { LlmProviders::Openai => Some(&OpenAiResponsesApiConfig), @@ -111,9 +114,8 @@ impl ResponsesRoute { } let provider = provider_name .parse::() - .map_err(|_| Error::Unsupported("native HTTP responses provider"))?; - let config = responses_provider(provider) - .ok_or(Error::Unsupported("native HTTP responses provider"))?; + .map_err(|_| Error::Declined(NOT_NATIVE))?; + let config = responses_provider(provider).ok_or(Error::Declined(NOT_NATIVE))?; let wire_model = match provider { LlmProviders::Openai if resolved_model.contains('/') => { return Err(Error::InvalidProvider(call.model)); @@ -123,7 +125,7 @@ impl ResponsesRoute { format!("openrouter/{resolved_model}") } LlmProviders::Openrouter => resolved_model.into(), - _ => return Err(Error::Unsupported("native HTTP responses provider")), + _ => return Err(Error::Declined(NOT_NATIVE)), }; let provider = ResponsesProvider { identity: ProviderIdentity { diff --git a/litellm-rust/crates/inference-responses/src/websocket.rs b/litellm-rust/crates/inference-responses/src/websocket.rs index 213a9a95831..2afcde0a7a4 100644 --- a/litellm-rust/crates/inference-responses/src/websocket.rs +++ b/litellm-rust/crates/inference-responses/src/websocket.rs @@ -1,14 +1,18 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use futures_util::{SinkExt, StreamExt}; -use litellm_http::websocket::{UpstreamWebSocket, connect_upstream}; +use litellm_http::{ + Client, + websocket::{self, UpstreamWebSocket}, +}; use tokio::sync::Mutex; use tokio_tungstenite::tungstenite::{ Message, - client::IntoClientRequest, - http::{HeaderName, HeaderValue}, + http::{HeaderMap, HeaderName, HeaderValue}, }; +use litellm_llms::ErrorDetail; + use super::Error; #[derive(Clone)] @@ -24,42 +28,32 @@ impl ResponsesWebSocketConnection { fields(outcome) )] pub async fn connect_url( + client: &Client, url: &str, headers: &HashMap, timeout: Option, ) -> Result { litellm_inference::diagnostic::operation("litellm.websocket.connect_url", async { - let mut request = url.into_client_request().map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; - for (name, value) in headers { - let header_name = name - .parse::() - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; - let header_value = HeaderValue::from_str(value) - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; - request.headers_mut().insert(header_name, header_value); - } - let connect = connect_upstream(request); - let result = match timeout { - Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { - Error::Transport(litellm_http::transport::Error::Network( - "Responses WebSocket connection timed out".into(), + let headers = headers + .iter() + .map(|(name, value)| { + Ok(( + name.parse::().map_err(|error| { + Error::invalid_request(ErrorDetail::invalid( + "WebSocket header name", + error, + )) + })?, + HeaderValue::from_str(value).map_err(|error| { + Error::invalid_request(ErrorDetail::invalid( + "WebSocket header value", + error, + )) + })?, )) - })?, - None => connect.await, - }; - let (socket, _) = result.map_err(|error| match *error { - tokio_tungstenite::tungstenite::Error::Http(response) => { - Error::Transport(litellm_http::transport::Error::Http { - status: response.status().as_u16(), - body: String::new(), - }) - } - other => { - Error::Transport(litellm_http::transport::Error::Network(other.to_string())) - } - })?; + }) + .collect::>()?; + let socket = websocket::connect(client, url, headers, timeout).await?; Ok(Self { socket: Arc::new(Mutex::new(Some(socket))), }) @@ -77,13 +71,12 @@ impl ResponsesWebSocketConnection { litellm_inference::diagnostic::operation("litellm.websocket.send_text", async { let mut socket = self.socket.lock().await; let Some(socket) = socket.as_mut() else { - return Err(Error::Transport(litellm_http::transport::Error::Network( - "Responses WebSocket is closed".into(), - ))); + return Err(litellm_http::Error::websocket_closed().into()); }; - socket.send(Message::Text(text)).await.map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - }) + socket + .send(Message::Text(text)) + .await + .map_err(|error| litellm_http::Error::from(error).into()) }) .await } @@ -104,12 +97,12 @@ impl ResponsesWebSocketConnection { Some(Ok(Message::Text(text))) => Ok(Some(text)), Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) .map(Some) - .map_err(|error| Error::InvalidResponse(error.to_string().into())), + .map_err(|error| { + Error::invalid_response(ErrorDetail::invalid("WebSocket text frame", error)) + }), Some(Ok(Message::Close(_))) | None => Ok(None), Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network( - error.to_string(), - ))), + Some(Err(error)) => Err(litellm_http::Error::from(error).into()), } }) .await @@ -125,9 +118,10 @@ impl ResponsesWebSocketConnection { litellm_inference::diagnostic::operation("litellm.websocket.close", async { let mut socket = self.socket.lock().await; if let Some(socket) = socket.as_mut() { - socket.close(None).await.map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; + socket + .close(None) + .await + .map_err(litellm_http::Error::from)?; } *socket = None; Ok(()) diff --git a/litellm-rust/crates/inference-responses/tests/openrouter.rs b/litellm-rust/crates/inference-responses/tests/openrouter.rs index e52cdba8ce8..8cdb68c376f 100644 --- a/litellm-rust/crates/inference-responses/tests/openrouter.rs +++ b/litellm-rust/crates/inference-responses/tests/openrouter.rs @@ -1,10 +1,7 @@ use std::sync::Arc; use futures_util::TryStreamExt; -use litellm_inference_responses::{ - Error, - types::{ResponsesCall, ResponsesOutput}, -}; +use litellm_inference_responses::types::{ResponsesCall, ResponsesOutput}; use litellm_inference_testing::{RecordingSecrets, no_secrets}; use rstest::{fixture, rstest}; use serde_json::json; @@ -108,13 +105,12 @@ async fn a_missing_key_fails_before_transport(call: ResponsesCall) { None, ) .await; - assert!(matches!( - result, - Err(Error::Auth(litellm_auth::Error::MissingApiKey { - provider: "OpenRouter", - .. - })) - )); + let Err(error) = result else { + panic!("the call fails"); + }; + assert!(error.is_missing_credential(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("OpenRouter"), "{report}"); assert!(received(&upstream).await.is_empty()); } diff --git a/litellm-rust/crates/inference-responses/tests/responses.rs b/litellm-rust/crates/inference-responses/tests/responses.rs index ae07a71dabd..89368a941ef 100644 --- a/litellm-rust/crates/inference-responses/tests/responses.rs +++ b/litellm-rust/crates/inference-responses/tests/responses.rs @@ -103,10 +103,7 @@ async fn unsupported_providers_fail_before_transport( ) .await; - assert!(matches!( - result, - Err(Error::Unsupported("native HTTP responses provider")) - )); + assert!(matches!(result, Err(ref error) if error.declined().is_some())); assert!(received(&upstream).await.is_empty()); } @@ -304,6 +301,25 @@ async fn provider_failures_emit_failure_once( ); } +#[rstest] +#[tokio::test] +async fn a_transport_error_does_not_echo_the_provider_url(call: ResponsesCall) { + let host = RecordingCall::::new(ResponsesCall { + api_base: Some(format!("{UNREACHABLE_BASE}/tenant-token-1234")), + ..call + }); + let call = host.request.lock().unwrap().take().unwrap(); + + let Err(error) = responses_route(no_secrets()) + .execute(call, &host, None) + .await + else { + panic!("nothing is listening"); + }; + + assert!(!error.to_string().contains("tenant-token-1234"), "{error}"); +} + #[rstest] #[case::explicit(true)] #[case::from_secrets(false)] @@ -538,6 +554,7 @@ async fn websocket_operations_trace_outcomes_without_capturing_frames_or_credent .logger() .instrument(async { let connection = ResponsesWebSocketConnection::connect_url( + &litellm_http::Client::for_test(), &format!("ws://{address}/responses"), &std::collections::HashMap::from([( "authorization".into(), @@ -556,12 +573,11 @@ async fn websocket_operations_trace_outcomes_without_capturing_frames_or_credent Some("private-frame-sentinel") ); connection.close().await.unwrap(); - assert!( - connection - .send_text("private-frame-sentinel".into()) - .await - .is_err() - ); + let error = connection + .send_text("private-frame-sentinel".into()) + .await + .expect_err("a closed connection sends nothing"); + assert!(!error.dispatched(), "{error:?}"); }) .await; server.await.unwrap(); diff --git a/litellm-rust/crates/inference-responses/tests/support/mod.rs b/litellm-rust/crates/inference-responses/tests/support/mod.rs index 81f133a0471..984b878eb21 100644 --- a/litellm-rust/crates/inference-responses/tests/support/mod.rs +++ b/litellm-rust/crates/inference-responses/tests/support/mod.rs @@ -13,6 +13,9 @@ use litellm_secrets::source::SecretSource; use serde_json::Value; use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any}; +/// A port nothing listens on, for calls that must fail before any request is sent. +pub const UNREACHABLE_BASE: &str = "http://127.0.0.1:1"; + pub fn responses_route( secrets: Arc, ) -> litellm_inference_responses::ResponsesRoute { diff --git a/litellm-rust/crates/inference-transcription/Cargo.toml b/litellm-rust/crates/inference-transcription/Cargo.toml index c501d329852..cf50b3f84fd 100644 --- a/litellm-rust/crates/inference-transcription/Cargo.toml +++ b/litellm-rust/crates/inference-transcription/Cargo.toml @@ -17,6 +17,7 @@ serde_json.workspace = true tracing.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-inference-testing.workspace = true litellm-tracing.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/inference-transcription/src/handler.rs b/litellm-rust/crates/inference-transcription/src/handler.rs index 5a7238c5e36..2ed3b1eba7b 100644 --- a/litellm-rust/crates/inference-transcription/src/handler.rs +++ b/litellm-rust/crates/inference-transcription/src/handler.rs @@ -2,7 +2,7 @@ use std::time::Duration; use litellm_http::{ Client, - request::{string_headers, truncate_error_body, with_default_headers}, + request::{string_headers, with_default_headers}, }; use litellm_llms::base_llm::{ audio_transcription::transformation::BaseAudioTranscriptionConfig, @@ -48,35 +48,27 @@ pub(super) async fn execute( let authenticated = resolve_auth(auth, environment, &env_lookup).await?; let outbound = litellm_inference::outbound::outbound_request( authenticated, - url, + &url, &body, Some( call.timeout .unwrap_or(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)), ), )?; - let response = litellm_inference::outbound::send(outbound, http) - .await - .map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; + let response = litellm_inference::outbound::send(outbound, http).await?; let status = response.status(); - let text = response.text().await.map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; + let text = response.text().await.map_err(litellm_http::Error::from)?; if !status.is_success() { - return Err(Error::Transport(litellm_http::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - })); + return Err(litellm_http::Error::upstream(status.as_u16(), text).into()); } let response_json = serde_json::from_str(&text).map_err(|error| { - Error::InvalidResponse(litellm_llms::ErrorDetail::invalid( + Error::invalid_response(litellm_llms::ErrorDetail::invalid( "audio response JSON", error, )) })?; - Ok(config - .transform_audio_transcription_response(provider.model, response_json)? - .into_json()) + config + .transform_audio_transcription_response(provider.model, response_json) + .map(|response| response.into_json()) + .map_err(|error| Error::from(error).after_send()) } diff --git a/litellm-rust/crates/inference-transcription/tests/audio_transcription.rs b/litellm-rust/crates/inference-transcription/tests/audio_transcription.rs index eb252bdff7e..9e48edce0c0 100644 --- a/litellm-rust/crates/inference-transcription/tests/audio_transcription.rs +++ b/litellm-rust/crates/inference-transcription/tests/audio_transcription.rs @@ -152,13 +152,7 @@ async fn invalid_audio_is_rejected_before_sending( .await .expect_err("invalid audio is rejected"); - assert!( - matches!( - error, - Error::InvalidRequest(_) | Error::MissingField(_) | Error::InvalidType { .. } - ), - "{error:?}" - ); + assert!(error.is_request(), "{error:?}"); assert!(received(&upstream).await.is_empty()); } @@ -185,7 +179,9 @@ async fn unsupported_providers_are_rejected_before_sending( .await .expect_err("unsupported provider errors"); - assert_eq!(error, Error::InvalidProvider(reported.into())); + assert!(error.is_request(), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(reported), "{report}"); } #[rstest] @@ -222,14 +218,28 @@ async fn an_upstream_error_keeps_its_status_and_body( .expect_err("upstream error propagates"); assert_eq!( - error, - Error::Transport(litellm_http::transport::Error::Http { - status, - body: "upstream said no".into() - }) + error.upstream_response(), + Some((status, "upstream said no")) ); } +#[rstest] +#[tokio::test] +async fn a_transport_error_does_not_echo_the_provider_url( + request: AudioTranscriptionRequest<'static>, +) { + let base = format!("{UNREACHABLE_BASE}/tenant-token-1234"); + + let error = transcribe(AudioTranscriptionRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("nothing is listening"); + + assert!(!error.to_string().contains("tenant-token-1234"), "{error}"); +} + #[rstest] #[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))] #[case::no_output(json_response(json!({"unexpected": true})))] @@ -248,7 +258,7 @@ async fn an_unreadable_success_body_is_an_invalid_response( .await .expect_err("an unreadable body fails"); - assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!(error.is_response(), "{error:?}"); } #[rstest] diff --git a/litellm-rust/crates/inference/AGENTS.md b/litellm-rust/crates/inference/AGENTS.md index 7f2bd0387be..bcbc98a209e 100644 --- a/litellm-rust/crates/inference/AGENTS.md +++ b/litellm-rust/crates/inference/AGENTS.md @@ -53,10 +53,10 @@ ## Error placement -- The workspace `Error definitions` rules shape each crate's error. This section decides which crate and module a failure belongs to +- The [rust-error skill](../../.agents/skills/rust-error/SKILL.md) shapes each crate's error. This section decides which crate and module a failure belongs to - A failure is declared once, by the lowest crate that raises it - - every crate above nests it unchanged (`#[error(transparent)] Auth(#[from] litellm_auth::Error)`) or maps it once at its boundary, as `src/error.rs` does for `litellm_llms::Error` - - `Error` collects route failures for every format crate and never re-declares a variant a lower crate raises + - every crate above nests it in one variant and delegates its predicates, as `Error::Llm`, `Error::Auth` and `Error::Http` do, or translates it once in the shared `From` when its meaning changes, as `From` does for `Unsupported` and `From` does for a header the signer computes + - `Error` collects route failures for every format crate and never re-declares a variant a lower crate raises. Routes build transformation failures with `Error::invalid_request` and `Error::invalid_response` and collapse anything raised after the send with `Error::after_send`, and callers ask `declined()`, `is_request()`, `is_response()`, `is_missing_credential()`, `dispatched()`, `upstream_response()` and `rejected_request()` instead of matching nested variants - Scope follows the concept, not the first caller - an error type under `litellm-llms`'s `/` is private to that provider: no other provider and nothing in `base_llm` may import it - a failure two providers or two routes can hit (wire framing, stream event decoding, a malformed provider response) belongs to the crate that owns the concept: `litellm-framer` for framing, `litellm_llms::Error` for the transformation layer diff --git a/litellm-rust/crates/inference/Cargo.toml b/litellm-rust/crates/inference/Cargo.toml index acc211f214b..807a3ea85d3 100644 --- a/litellm-rust/crates/inference/Cargo.toml +++ b/litellm-rust/crates/inference/Cargo.toml @@ -29,6 +29,7 @@ time.workspace = true url.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-cache-memory.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-llms = { workspace = true, features = ["test-support"] } diff --git a/litellm-rust/crates/inference/src/caching.rs b/litellm-rust/crates/inference/src/caching.rs index 09d03c06215..d9b5b891015 100644 --- a/litellm-rust/crates/inference/src/caching.rs +++ b/litellm-rust/crates/inference/src/caching.rs @@ -377,5 +377,66 @@ async fn publish( facts: facts.clone(), })); } - interceptors.result_ready(facts).await + let from_provider = facts.source == ResultSource::Provider; + interceptors.result_ready(facts).await.map_err(|error| { + if from_provider { + Error::post_call(error) + } else { + error + } + }) +} + +#[cfg(test)] +mod tests { + use litellm_core_utils::decline::Decline; + use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest}; + + use super::*; + + struct RejectResult; + + impl Interceptors for RejectResult { + async fn before_provider_request( + &self, + wire: WireRequest, + _: RequestContext, + ) -> Result { + Ok(wire) + } + + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { + Ok(()) + } + + async fn result_ready(&self, _: ExecutionFacts) -> Result<(), Error> { + Err(Error::Declined(Decline::new("result rejected"))) + } + } + + #[rstest::rstest] + #[case::provider(ResultSource::Provider, false)] + #[case::cache(ResultSource::Cache { key: "cached".into() }, true)] + #[tokio::test] + async fn result_rejection_only_declines_without_a_provider_call( + #[case] source: ResultSource, + #[case] may_decline: bool, + ) { + let facts = ExecutionFacts { + provider: ProviderIdentity { + model: "model".into(), + provider: "provider".into(), + }, + cost: Cost::settle(&source, None), + source, + }; + let error = publish(facts, &RejectResult, None).await.unwrap_err(); + assert_eq!(error.declined().is_some(), may_decline); + assert_eq!(error.dispatched(), !may_decline); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("result rejected") + ); + } } diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index 655b768dde1..963aa7b90c5 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -1,6 +1,6 @@ use litellm_cache_response::CachePolicy; use litellm_host::{ - interceptors::{ExecutionFacts, Interceptors, RawResponse}, + interceptors::{ExecutionFacts, Interceptors, RawResponse, ResultSource}, lifecycle::{CallEvent, ExecutionEvent}, observation::ObservationSender, }; @@ -28,7 +28,17 @@ impl<'a, I: Interceptors> CallContext<'a, I> { facts: facts.clone(), })); } - self.interceptors.result_ready(facts).await + let from_provider = facts.source == ResultSource::Provider; + self.interceptors + .result_ready(facts) + .await + .map_err(|error| { + if from_provider { + Error::post_call(error) + } else { + error + } + }) } pub async fn response_received(&self, body: &str) -> Result<(), Error> { diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index 895c8c55e01..a8e6fb4e1b3 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -1,57 +1,48 @@ -//! One error for every route in this crate. OCR still carries its own, richer enum. -//! -//! A variant is declared by the layer that produces it and nested here as is: -//! credentials by `litellm_auth` (AWS folds into it at that crate's boundary), the wire by -//! `litellm_http`, secrets by `litellm_secrets`. The transformation layer's [`LlmError`] -//! maps onto the same-named variants once, here, so no route re-declares them. - use std::sync::Arc; -use litellm_http::transport::Error as TransportError; -use litellm_llms::Error as LlmError; +use litellm_core_utils::decline::Decline; +use litellm_host::machine::MachineFault; +use litellm_llms::{Error as LlmError, ErrorDetail}; #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum Error { - #[error("expected {expected}, got {actual}")] - InvalidType { - expected: &'static str, - actual: &'static str, - }, - #[error("missing required field: {0}")] - MissingField(&'static str), #[error("invalid provider: {0}")] InvalidProvider(String), - #[error("invalid request: {0}")] - InvalidRequest(#[source] litellm_llms::ErrorDetail), - #[error("invalid response: {0}")] - InvalidResponse(#[source] litellm_llms::ErrorDetail), - #[error("unsupported by the rust path: {0}")] - Unsupported(&'static str), + #[error("invalid provider: unable to resolve custom_llm_provider for {route} request")] + UnresolvedProvider { route: &'static str }, + #[error(transparent)] + Declined(Decline), + #[error(transparent)] + Llm(LlmError), #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] - Transport(#[from] TransportError), - #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] - Http(#[from] litellm_http::Error), + Http(litellm_http::Error), #[error(transparent)] Secret(#[from] SecretError), - #[error("post-call hook failed: {0}")] + #[error("the host driver failed")] + Host(#[source] MachineFault), + #[error("post-call hook failed")] PostCallHook(#[source] Arc), } -impl From for Error { - fn from(fault: litellm_host::machine::MachineFault) -> Self { - use litellm_host::machine::MachineFault; - Self::InvalidRequest(match fault { - MachineFault::Abandoned => "host driver was abandoned".into(), - MachineFault::Protocol(message) => format!("host {message}").into(), - }) +impl From for Error { + fn from(fault: MachineFault) -> Self { + Self::Host(fault) } } impl Error { + pub fn invalid_request(detail: ErrorDetail) -> Self { + Self::Llm(LlmError::InvalidRequest(detail)) + } + + pub fn invalid_response(detail: ErrorDetail) -> Self { + Self::Llm(LlmError::InvalidResponse(detail)) + } + pub fn post_call(error: Self) -> Self { Self::PostCallHook(Arc::new(error)) } @@ -60,31 +51,138 @@ impl Error { /// the provider's answer. pub fn is_request(&self) -> bool { match self { - Self::InvalidType { .. } - | Self::MissingField(_) - | Self::InvalidProvider(_) - | Self::InvalidRequest(_) - | Self::Unsupported(_) - | Self::Headers(_) => true, - Self::Auth(error) => !matches!(error, litellm_auth::Error::MissingApiKey { .. }), - Self::InvalidResponse(_) - | Self::Transport(_) - | Self::Http(_) + Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } | Self::Headers(_) => true, + Self::Llm(error) => error.is_request(), + Self::Auth(error) => error.is_request(), + Self::Http(error) => error.is_request(), + Self::Declined(_) | Self::Secret(_) | Self::Host(_) | Self::PostCallHook(_) => false, + } + } + + pub fn is_response(&self) -> bool { + match self { + Self::Llm(error) => error.is_response(), + Self::Http(error) => error.is_response(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Declined(_) + | Self::Auth(_) + | Self::Headers(_) | Self::Secret(_) + | Self::Host(_) | Self::PostCallHook(_) => false, } } + + pub fn is_missing_credential(&self) -> bool { + match self { + Self::Llm(error) => error.is_missing_credential(), + Self::Auth(error) => error.is_missing_credential(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Declined(_) + | Self::Headers(_) + | Self::Http(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) => false, + } + } + + pub fn declined(&self) -> Option<&Decline> { + match self { + Self::Declined(decline) => Some(decline), + Self::Llm(error) => error.declined(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Auth(_) + | Self::Headers(_) + | Self::Http(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) => None, + } + } + + /// Re-tag an error raised while handling an answer the provider already returned. + /// + /// A config reports the same variants on either side of the call, so a missing field or + /// a decline can mean "this request cannot be translated" before the send and "this + /// response cannot be normalized" after it. Only the second kind was billed, and a host + /// that keeps a reference implementation must not retry it. + pub fn after_send(self) -> Self { + if self.is_response() || self.upstream_response().is_some() { + return self; + } + Self::invalid_response(ErrorDetail::invalid("provider response", self)) + } + + pub fn dispatched(&self) -> bool { + match self { + Self::Llm(error) => error.dispatched(), + Self::Http(error) => error.dispatched(), + Self::Host(_) | Self::PostCallHook(_) => true, + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Declined(_) + | Self::Auth(_) + | Self::Headers(_) + | Self::Secret(_) => false, + } + } + + pub fn rejected_request(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Llm(error) => error.rejected_request(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Declined(_) + | Self::Auth(_) + | Self::Headers(_) + | Self::Http(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) => None, + } + } + + pub fn upstream_response(&self) -> Option<(u16, &str)> { + match self { + Self::Llm(error) => error.upstream_response(), + Self::Http(error) => error.upstream_response(), + Self::InvalidProvider(_) + | Self::UnresolvedProvider { .. } + | Self::Declined(_) + | Self::Auth(_) + | Self::Headers(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) => None, + } + } +} + +impl From for Error { + fn from(error: litellm_http::Error) -> Self { + // Python drops the caller's copy and prefers a forwarded Authorization + // over the signature, so leave the request to it. + if error.is_computed_header() { + return Self::Declined(Decline::new("request forwards a header AWS SigV4 computes")); + } + Self::Http(error) + } } impl From for Error { fn from(error: LlmError) -> Self { match error { - LlmError::InvalidType { expected, actual } => Self::InvalidType { expected, actual }, - LlmError::MissingField(field) => Self::MissingField(field), - LlmError::InvalidRequest(message) => Self::InvalidRequest(message), - LlmError::InvalidResponse(message) => Self::InvalidResponse(message), - LlmError::Unsupported(reason) => Self::Unsupported(reason), + LlmError::Declined(decline) => Self::Declined(decline), + LlmError::Http(error) => Self::from(error), LlmError::Auth(error) => Self::Auth(error), + error @ (LlmError::InvalidType { .. } + | LlmError::MissingField(_) + | LlmError::InvalidRequest(_) + | LlmError::InvalidResponse(_)) => Self::Llm(error), } } } @@ -116,26 +214,205 @@ impl Eq for SecretError {} #[cfg(test)] mod tests { use super::Error; + use litellm_core_utils::decline::Decline; + use litellm_host::machine::MachineFault; use litellm_llms::{Error as LlmError, ErrorDetail}; use rstest::rstest; - #[test] - fn a_missing_api_key_is_the_environment_not_the_request() { - assert!( - !Error::Auth(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: "ANTHROPIC_API_KEY", - }) - .is_request() - ); - assert!(Error::Auth(litellm_auth::Error::InvalidHeader).is_request()); - assert!(Error::InvalidRequest("top_k".into()).is_request()); - assert!(!Error::InvalidResponse("bad json".into()).is_request()); + #[rstest] + #[case::missing_api_key( + litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + } + .into(), + false + )] + #[case::invalid_auth_header(litellm_auth::Error::InvalidHeader.into(), true)] + #[case::transformation_missing_api_key( + LlmError::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + }) + .into(), + false + )] + #[case::invalid_request( + LlmError::InvalidRequest(ErrorDetail::Expected { + subject: "top_k", + expected: "an integer", + }) + .into(), + true + )] + #[case::invalid_response( + LlmError::InvalidResponse(ErrorDetail::MissingField { + subject: "response", + field: "usage", + }) + .into(), + false + )] + #[case::host_driver(MachineFault::Abandoned.into(), false)] + #[case::post_call_hook_over_a_request_fault( + Error::post_call(LlmError::MissingField("usage").into()), + false + )] + fn only_the_callers_input_is_a_request_fault(#[case] error: Error, #[case] is_request: bool) { + assert_eq!(error.is_request(), is_request); } + + #[rstest] + #[case::own_auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + } + .into())] + #[case::transformation_auth( + LlmError::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + }) + .into() + )] + fn a_missing_api_key_is_a_missing_credential_wherever_it_is_raised(#[case] error: Error) { + assert!(error.is_missing_credential()); + } + + #[rstest] + #[case::direct(litellm_http::Error::computed_header("authorization").into())] + #[case::through_a_transformation( + LlmError::Http(litellm_http::Error::computed_header("authorization")).into() + )] + fn a_header_the_signer_computes_is_left_to_the_host_on_every_route(#[case] error: Error) { + assert!(error.declined().is_some(), "{error:?}"); + assert!(!error.dispatched()); + } + + #[test] + fn a_transformation_decline_survives_into_the_route_error() { + let decline = Decline::new("unrecognized request parameter"); + let error = Error::from(LlmError::Declined(decline)); + assert_eq!(error.declined(), Some(&decline)); + assert!(!error.is_request()); + assert!(!error.dispatched()); + } + + #[test] + fn a_post_call_hook_failure_is_never_a_decline() { + let hook = Error::post_call(Error::Declined(Decline::new("hook"))); + assert_eq!(hook.declined(), None); + } + + #[rstest] + #[case::declined(Error::Declined(Decline::new("non-text response content block")))] + #[case::transformation_decline(LlmError::Declined(Decline::new("streaming")).into())] + #[case::missing_field(LlmError::MissingField("usage").into())] + #[case::invalid_request(Error::invalid_request(ErrorDetail::Expected { + subject: "whatever", + expected: "anything else", + }))] + #[case::auth(Error::Auth(litellm_auth::Error::InvalidHeader))] + #[case::host_driver(MachineFault::Abandoned.into())] + fn errors_after_the_send_are_response_faults_and_never_a_decline(#[case] error: Error) { + let report = litellm_error::Report::new(&error).to_string(); + let collapsed = error.after_send(); + assert!(collapsed.is_response(), "{collapsed:?}"); + assert!(!collapsed.is_request(), "{collapsed:?}"); + assert!(collapsed.dispatched(), "{collapsed:?}"); + assert_eq!(collapsed.declined(), None); + let collapsed_report = litellm_error::Report::new(&collapsed).to_string(); + assert!(collapsed_report.ends_with(&report), "{collapsed_report}"); + } + + #[test] + fn a_response_fault_is_not_wrapped_twice_after_the_send() { + let original = Error::invalid_response(ErrorDetail::MissingField { + subject: "response", + field: "usage", + }); + let report = litellm_error::Report::new(&original).to_string(); + assert_eq!( + litellm_error::Report::new(&original.after_send()).to_string(), + report + ); + } + + #[test] + fn a_provider_answer_stays_as_is_after_the_send() { + let upstream = Error::from(litellm_http::Error::upstream(429, "slow down".into())); + assert_eq!( + upstream.after_send().upstream_response(), + Some((429, "slow down")) + ); + } + + #[test] + fn an_upstream_failure_through_a_transformation_keeps_the_provider_answer() { + let error = Error::from(LlmError::Http(litellm_http::Error::upstream( + 429, + "slow down".into(), + ))); + assert_eq!(error.upstream_response(), Some((429, "slow down"))); + assert!(error.dispatched()); + } + + #[rstest] + #[case::invalid_request( + Error::invalid_request(ErrorDetail::Expected { + subject: "messages", + expected: "non-empty", + }), + Some("messages must be non-empty") + )] + #[case::missing_field( + LlmError::MissingField("max_tokens").into(), + Some("missing required field: max_tokens") + )] + #[case::invalid_response( + Error::invalid_response(ErrorDetail::MissingField { + subject: "response", + field: "usage", + }), + None + )] + #[case::provider_rejection(litellm_http::Error::upstream(400, "bad".into()).into(), None)] + #[case::unresolved_provider(Error::InvalidProvider("unknown".into()), None)] + fn a_rejected_request_names_only_what_the_caller_got_wrong( + #[case] error: Error, + #[case] expected: Option<&str>, + ) { + let rejection = error + .rejected_request() + .map(|rejection| litellm_error::Report::new(rejection).to_string()); + assert_eq!(rejection.as_deref(), expected); + } + + #[test] + fn a_post_call_hook_failure_is_not_the_providers_answer() { + let error = + Error::post_call(litellm_http::Error::upstream(500, "hook backend down".into()).into()); + assert_eq!(error.upstream_response(), None); + assert!(!error.is_request()); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.starts_with("post-call hook failed: "), "{report}"); + assert!(report.contains("hook backend down"), "{report}"); + } + + #[test] + fn a_host_driver_fault_keeps_the_fault_as_its_source() { + let error = Error::from(MachineFault::Abandoned); + assert_eq!( + litellm_error::find::(&error), + Some(&MachineFault::Abandoned) + ); + assert!(error.dispatched()); + } + #[rstest] #[case::request(true)] #[case::response(false)] - fn contextual_errors_preserve_sources_and_route_classification(#[case] request: bool) { + fn contextual_errors_keep_their_cause_and_route_classification(#[case] request: bool) { let source = serde_json::from_str::("{").unwrap_err(); let source_message = source.to_string(); let detail = ErrorDetail::invalid("test payload", source); @@ -145,16 +422,10 @@ mod tests { LlmError::InvalidResponse(detail) }); assert_eq!(error.is_request(), request); - let category = if request { "request" } else { "response" }; - assert_eq!( - error.to_string(), - format!("invalid {category}: invalid test payload: {source_message}") - ); - let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| { - error.source() - }) - .find_map(|error| error.downcast_ref::()) - .expect("the original JSON error remains available"); - assert_eq!(source.to_string(), source_message); + assert_eq!(error.is_response(), !request); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("test payload"), "{report}"); + assert_eq!(report.matches(&source_message).count(), 1, "{report}"); + assert!(litellm_error::find::(&error).is_some()); } } diff --git a/litellm-rust/crates/inference/src/outbound.rs b/litellm-rust/crates/inference/src/outbound.rs index 41fe3aa62b8..b0cc71d425d 100644 --- a/litellm-rust/crates/inference/src/outbound.rs +++ b/litellm-rust/crates/inference/src/outbound.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_http::outbound::OutboundRequest; +use litellm_http::outbound::{json, signed}; use litellm_llms::base_llm::auth::Authenticated; use serde_json::Value; @@ -11,25 +11,26 @@ use serde_json::Value; fields(status) )] pub async fn send( - request: OutboundRequest, + request: reqwest::Request, client: &litellm_http::Client, -) -> Result { - request.send(client).await.inspect(|response| { - tracing::Span::current().record("status", response.status().as_u16()); - }) +) -> Result { + let response = client.execute(request).await?; + tracing::Span::current().record("status", response.status().as_u16()); + Ok(response) } /// Header credentials are already in `headers`; SigV4 is applied here, over the /// bytes that are sent. pub fn outbound_request( authenticated: Authenticated, - url: String, + url: &str, body: &Value, timeout: Option, -) -> Result { +) -> Result { let Authenticated { headers, signer } = authenticated; + let request = json(url, &headers, body, timeout)?; match signer { - None => OutboundRequest::json(url, headers, body, timeout), - Some(signer) => OutboundRequest::signed_json(url, headers, body, timeout, &signer), + None => Ok(request), + Some(signer) => signed(request, &signer), } } diff --git a/litellm-rust/crates/inference/src/provider.rs b/litellm-rust/crates/inference/src/provider.rs index 305aca85f97..fa226a4ef8d 100644 --- a/litellm-rust/crates/inference/src/provider.rs +++ b/litellm-rust/crates/inference/src/provider.rs @@ -25,11 +25,7 @@ pub fn resolve_llm_provider<'a>( custom_llm_provider: provider, }) }) - .ok_or_else(|| { - Error::InvalidProvider(format!( - "unable to resolve custom_llm_provider for {route} request" - )) - })?; + .ok_or(Error::UnresolvedProvider { route })?; let provider = custom_llm_provider .parse() .map_err(|_| Error::InvalidProvider(custom_llm_provider.to_string()))?; diff --git a/litellm-rust/crates/inference/tests/caching.rs b/litellm-rust/crates/inference/tests/caching.rs index a0a6a757565..4e8b2349ac8 100644 --- a/litellm-rust/crates/inference/tests/caching.rs +++ b/litellm-rust/crates/inference/tests/caching.rs @@ -14,6 +14,7 @@ use litellm_cache_memory::InMemoryCache; use litellm_cache_response::{ CacheOptions, CachePolicy, CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, }; +use litellm_core_utils::decline::Decline; use litellm_host::{ call::{CallOutput, OutputOf}, interceptors::{ @@ -204,7 +205,8 @@ async fn streamed( .chunks(3) .map(|bytes| Ok(Bytes::copy_from_slice(bytes))) .collect::>(); - let ending = fail.then_some(Err(Error::Unsupported("test transport failure"))); + let ending = + fail.then_some(Err(Error::Declined(Decline::new("test transport failure")))); Ok(CallOutput::Stream { head: (), chunks: stream::iter(chunks.into_iter().chain(ending)).boxed(), @@ -315,7 +317,7 @@ async fn a_provider_failure_never_populates_the_cache(cache: Arc Result, litellm_cache::Error> { - Err(litellm_cache::Error::Unavailable) + Err(litellm_cache::Error::Unavailable(None)) } fn set_cache( @@ -426,7 +428,7 @@ impl litellm_cache::BaseCache for UnavailableCache { _: Self::Value, _: &Self::Context, ) -> Result<(), litellm_cache::Error> { - Err(litellm_cache::Error::Unavailable) + Err(litellm_cache::Error::Unavailable(None)) } } @@ -482,7 +484,7 @@ impl Interceptors for CacheHitAccounting { ); self.key.set(key).unwrap(); if self.reject { - return Err(Error::Unsupported("cache accounting rejected")); + return Err(Error::Declined(Decline::new("cache accounting rejected"))); } Ok(()) } @@ -569,10 +571,10 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure( .await }; if reject { - assert!(matches!( - result, - Err(Error::Unsupported("cache accounting rejected")) - )); + assert_eq!( + result.unwrap_err().declined(), + Some(&Decline::new("cache accounting rejected")) + ); } else { assert_eq!(result.unwrap(), expected); } diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 2d3ea85d1fc..962f61ef4ec 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -11,6 +11,7 @@ test-support = ["litellm-http/test-support"] [dependencies] litellm-llms-types.workspace = true litellm-core-utils.workspace = true +litellm-error.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-auth-aws.workspace = true litellm-auth-azure.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs index b8ea701af99..0513f88038d 100644 --- a/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/batches/transformation.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use litellm_llms_types::formats::batches::{BatchRequestCounts, BatchResponse, BatchStatus}; use litellm_llms_types::formats::messages::MessagesResponse; use litellm_llms_types::providers::anthropic::{BATCHES_PATH, MESSAGES_PATH}; @@ -127,7 +128,9 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation { } fn transform_create_batch_request(&self) -> Result { - Err(Error::Unsupported("Anthropic message batch creation")) + Err(Error::Declined(Decline { + reason: "Anthropic message batch creation", + })) } fn transform_create_batch_response( @@ -135,7 +138,9 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation { _response: AnthropicMessageBatch, _now: i64, ) -> Result { - Err(Error::Unsupported("Anthropic message batch creation")) + Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "Anthropic message batch creation", + ))) } fn retrieve_batch_url( @@ -149,7 +154,12 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation { } let mut url = batches_base_url(api_base, env_lookup)?; url.path_segments_mut() - .map_err(|_| Error::InvalidRequest("Anthropic API base cannot be a base URL".into()))? + .map_err(|()| { + Error::InvalidRequest(crate::ErrorDetail::Expected { + subject: "Anthropic API base", + expected: "a base URL", + }) + })? .push(batch_id); Ok(url.into()) } @@ -217,11 +227,11 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation { .map(|(index, line)| { let record: BatchResultRecord = serde_json::from_str(line.trim()).map_err(|error| { - Error::InvalidResponse(crate::ErrorDetail::InvalidLine { - subject: "Anthropic batch result", - line: index + 1, - source: crate::ErrorSource::new(error), - }) + Error::InvalidResponse(crate::ErrorDetail::invalid_line( + "Anthropic batch result", + index + 1, + error, + )) })?; match record.result { BatchResult::Succeeded { message } => Ok(*message), @@ -314,30 +324,40 @@ mod tests { assert_eq!(messages[0].id, "msg_1"); } - #[test] - fn reports_malformed_and_unsuccessful_batch_results() { - assert!(matches!( - ANTHROPIC_BATCHES_TRANSFORMATION.transform_batch_results("not-json"), - Err(Error::InvalidResponse(message)) if message.to_string().contains("line 1") - )); - assert!(matches!( - ANTHROPIC_BATCHES_TRANSFORMATION.transform_batch_results( - r#"{"result":{"type":"errored","error":{"type":"invalid_request_error"}}}"# - ), - Err(Error::InvalidResponse(message)) if message.to_string().contains("request failed") - )); + #[rstest::rstest] + #[case::malformed("not-json", "line 1")] + #[case::errored( + r#"{"result":{"type":"errored","error":{"type":"invalid_request_error"}}}"#, + "request failed" + )] + fn reports_malformed_and_unsuccessful_batch_results(#[case] body: &str, #[case] fact: &str) { + let error = ANTHROPIC_BATCHES_TRANSFORMATION + .transform_batch_results(body) + .unwrap_err(); + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains(fact) + ); } - #[test] - fn preserves_python_placeholder_for_batch_creation() { - assert!(matches!( - ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_request(), - Err(Error::Unsupported("Anthropic message batch creation")) - )); + #[rstest::rstest] + fn batch_creation_only_declines_before_a_provider_response() { + let error = ANTHROPIC_BATCHES_TRANSFORMATION + .transform_create_batch_request() + .unwrap_err(); + assert!(!error.dispatched()); + assert_eq!( + error.declined().map(|decline| decline.reason), + Some("Anthropic message batch creation") + ); let response: AnthropicMessageBatch = serde_json::from_value(json!({})).unwrap(); - assert!(matches!( - ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_response(response, 0), - Err(Error::Unsupported("Anthropic message batch creation")) - )); + let error = ANTHROPIC_BATCHES_TRANSFORMATION + .transform_create_batch_response(response, 0) + .unwrap_err(); + assert!(error.dispatched()); + assert!(error.is_response()); + assert!(error.declined().is_none()); } } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/handler.rs b/litellm-rust/crates/llms/src/anthropic/chat/handler.rs deleted file mode 100644 index eaee006f228..00000000000 --- a/litellm-rust/crates/llms/src/anthropic/chat/handler.rs +++ /dev/null @@ -1,163 +0,0 @@ -use std::collections::HashMap; - -use litellm_llms_types::formats::{ - chat_completions::{ - ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk, - ChatCompletionsUsage, - }, - messages::streaming::{ - MessagesContentBlock, MessagesContentBlockDelta, MessagesStreamEvent, MessagesStreamUsage, - }, -}; -use serde_json::Value; - -use crate::{ - Error, - base_llm::{base_model_iterator::StreamTransformer, chat::streaming::StreamShape}, -}; - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum AnthropicJsonChunkType { - ValidJson, - AccumulatedJson, -} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum AnthropicContentBlockType { - Text, - ToolUse, - ServerToolUse, - Thinking, - RedactedThinking, - Compaction, - ToolResult(String), - Other(String), -} - -#[derive(Clone, Debug, PartialEq)] -pub struct AnthropicContentBlockDeltaEvent { - pub index: u64, - pub delta: MessagesContentBlockDelta, -} - -pub struct ModelResponseIterator { - pub content_blocks: Vec, - pub tool_index: i64, - pub json_mode: bool, - pub speed: Option, - pub tool_name_reverse_map: HashMap, - pub response_id: String, - pub served_model: Option, - pub is_response_format_tool: bool, - pub converted_response_format_tool: bool, - pub accumulated_json: String, - pub chunk_type: AnthropicJsonChunkType, - pub current_content_block_type: Option, - pub web_search_results: Vec, - pub web_search_calls: HashMap, - pub compaction_blocks: Vec, - pub reasoning_content_chunks: Vec, - pub server_tool_inputs: HashMap, - pub tool_results: Vec, - pub current_server_tool_id: Option, - pub container_id: Option, -} - -impl ModelResponseIterator { - pub fn new(_shape: StreamShape) -> Self { - todo!() - } - - pub fn check_empty_tool_call_args(&self) -> bool { - todo!() - } - - pub fn handle_usage(&mut self, _usage: MessagesStreamUsage) -> ChatCompletionsUsage { - todo!() - } - - pub fn handle_content_block_delta( - &mut self, - _index: u64, - _delta: MessagesContentBlockDelta, - ) -> ( - String, - Option, - Vec, - Option, - Option, - ) { - todo!() - } - - pub fn handle_content_block_start( - &mut self, - _index: u64, - _content_block: MessagesContentBlock, - ) -> Result { - todo!() - } - - pub fn handle_json_mode_chunk( - &mut self, - _text: String, - _tool_use: Option, - ) -> (String, Option) { - todo!() - } - - pub fn handle_accumulated_json_chunk( - &mut self, - _data: &str, - _is_final: bool, - ) -> Result, Error> { - todo!() - } - - pub fn handle_redacted_thinking_content( - &mut self, - _content_block: &MessagesContentBlock, - ) -> Vec { - todo!() - } - - pub fn web_search_call_snapshot(&self) -> HashMap { - todo!() - } - - pub fn complete_web_search_call(&mut self, _result: Value) { - todo!() - } - - pub fn build_code_interpreter_results(&self) -> Vec { - todo!() - } - - pub fn handle_message_delta( - &mut self, - _event: MessagesStreamEvent, - ) -> (Option, Option, Option) { - todo!() - } - - pub fn chunk_parser( - &mut self, - _event: MessagesStreamEvent, - ) -> Result { - todo!() - } -} - -impl StreamTransformer for ModelResponseIterator { - type Input = MessagesStreamEvent; - type Output = ChatCompletionChunk; - type Error = Error; - - fn transform(&mut self, _input: Self::Input) -> Result, Self::Error> { - todo!() - } - - fn finish(&mut self) -> Result, Self::Error> { - todo!() - } -} diff --git a/litellm-rust/crates/llms/src/anthropic/chat/mod.rs b/litellm-rust/crates/llms/src/anthropic/chat/mod.rs index f0050b7dc71..f239b6921fa 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/mod.rs @@ -1,2 +1 @@ -pub mod handler; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 91baa247ac8..eb796a900f5 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,4 +1,5 @@ use litellm_auth::SecretValue; +use litellm_core_utils::decline::Decline; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, prompt_templates::factory::{Conversation, build_conversation}, @@ -14,23 +15,19 @@ use serde_json::{Map, Value, json}; use crate::{ Error, - anthropic::{ - chat::handler::ModelResponseIterator, - common_utils::{ - API_KEY_PLACEMENT, complete_anthropic_url, forwarded_oauth_bearer, - resolve_anthropic_api_key, - }, + anthropic::common_utils::{ + API_KEY_PLACEMENT, complete_anthropic_url, forwarded_oauth_bearer, + resolve_anthropic_api_key, }, base_llm::{ auth::AuthScheme, chat::{ - streaming::{ChatStream, StreamShape}, + streaming::StreamShape, transformation::{ BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData, - Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param, + ValidatedEnvironment, unsupported_message, unsupported_param, }, }, - messages::streaming::event_stream as messages_event_stream, }, }; @@ -130,15 +127,14 @@ impl BaseConfig for AnthropicConfig { serde_json::from_value(response.body).map_err(|error| { Error::InvalidResponse(crate::ErrorDetail::invalid("messages response", error)) })?; - // The route declines tool and thinking requests, so a non-text block - // means the response carries something this path never asked for. - // Decline rather than silently dropping it; the host falls back. if body .content .iter() .any(|block| matches!(block, TextResponseBlock::Other)) { - return Err(Error::Unsupported("non-text response content block")); + return Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "non-text response content block", + ))); } let text: String = body .content @@ -194,13 +190,6 @@ impl BaseConfig for AnthropicConfig { Ok(ValidatedEnvironment { headers, auth }) } - fn model_response_iterator(&self, shape: StreamShape) -> Option { - Some(ChatStream::new( - messages_event_stream, - ModelResponseIterator::new(shape), - )) - } - fn default_headers(&self) -> &'static [(&'static str, &'static str)] { DEFAULT_HEADERS } @@ -214,7 +203,7 @@ impl BaseConfig for AnthropicConfig { &self, messages: &[ChatMessage], optional_params: &Map, - ) -> Option { + ) -> Option { unsupported_param(self.supported_openai_param_mappings(), &[], optional_params) .or_else(|| messages.iter().find_map(unsupported_message)) // Anthropic rejects a request whose first turn is not a user turn. @@ -222,7 +211,7 @@ impl BaseConfig for AnthropicConfig { // core cannot observe, so decline instead of guessing. .or_else(|| { (!build_conversation(messages).opens_on_user_turn()) - .then_some(Unsupported("conversation does not open on a user turn")) + .then_some(Decline::new("conversation does not open on a user turn")) }) } } @@ -271,9 +260,10 @@ mod tests { anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG, base_llm::{ auth::AuthScheme, - chat::transformation::{BaseConfig, ProviderChatResponseData, Unsupported}, + chat::transformation::{BaseConfig, ProviderChatResponseData}, }, }; + use litellm_core_utils::decline::Decline; use litellm_llms_types::formats::chat_completions::{ChatCompletionsResponse, ChatMessage}; use rstest::rstest; use serde_json::{Map, Value, json}; @@ -301,7 +291,7 @@ mod tests { .transform_response("claude-sonnet-4-5", ProviderChatResponseData { body }) } - fn reason(msgs: Value, opts: Value) -> Option { + fn reason(msgs: Value, opts: Value) -> Option { ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) } @@ -409,7 +399,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"top_k": 40}) ), - Some(Unsupported("unrecognized request parameter")) + Some(Decline::new("unrecognized request parameter")) ); } @@ -420,7 +410,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"stream": true, "max_tokens": 16}) ), - Some(Unsupported("streaming")) + Some(Decline::new("streaming")) ); } @@ -445,7 +435,7 @@ mod tests { fn declines_any_param_outside_the_allowlist(#[case] param: Value) { assert_eq!( reason(json!([{"role": "user", "content": "hi"}]), param.clone()), - Some(Unsupported("unrecognized request parameter")), + Some(Decline::new("unrecognized request parameter")), "expected {param} to decline" ); } @@ -463,7 +453,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported("unrecognized message field")) + Some(Decline::new("unrecognized message field")) ); assert_eq!( reason( @@ -473,7 +463,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported("unrecognized message field")) + Some(Decline::new("unrecognized message field")) ); assert_eq!( reason( @@ -483,7 +473,7 @@ mod tests { ]}]), json!({}) ), - Some(Unsupported("non-text message content")) + Some(Decline::new("non-text message content")) ); assert_eq!( reason( @@ -493,7 +483,7 @@ mod tests { ]}]), json!({}) ), - Some(Unsupported("non-text message content")) + Some(Decline::new("non-text message content")) ); } @@ -504,7 +494,7 @@ mod tests { // than declining to Python before it. assert_eq!( reason(json!([{"role": "user", "content": []}]), json!({})), - Some(Unsupported("message without content")) + Some(Decline::new("message without content")) ); assert_eq!( reason( @@ -525,7 +515,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported("conversation does not open on a user turn")) + Some(Decline::new("conversation does not open on a user turn")) ); } @@ -658,7 +648,7 @@ mod tests { } #[test] - fn declines_a_response_carrying_a_non_text_block() { + fn a_response_carrying_a_non_text_block_is_invalid() { let err = transform_response(json!({ "model": "claude-sonnet-4-5", "content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}], @@ -666,7 +656,12 @@ mod tests { "usage": {"input_tokens": 1, "output_tokens": 1} })) .expect_err("non-text block"); - assert_eq!(err, Error::Unsupported("non-text response content block")); + assert_eq!( + err, + Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "non-text response content block" + )) + ); } #[rstest::rstest] @@ -676,16 +671,11 @@ mod tests { #[case::missing_model(json!({"content": [], "usage": {"input_tokens": 1, "output_tokens": 1}}))] fn errors_on_a_response_missing_required_fields(#[case] body: Value) { let error = transform_response(body).expect_err("invalid response"); - assert!(matches!(error, Error::InvalidResponse(_))); - let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| { - error.source() - }) - .find_map(|error| error.downcast_ref::()) - .expect("the JSON decoding source is preserved"); - assert_eq!( - error.to_string(), - format!("invalid response: invalid messages response: {source}") - ); + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + let cause = litellm_error::find::(&error).expect("JSON cause"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("invalid messages response"), "{report}"); + assert_eq!(report.matches(&cause.to_string()).count(), 1, "{report}"); } #[test] diff --git a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs index 6e8de65d6fa..54f924ddec6 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs @@ -208,24 +208,29 @@ mod tests { #[case::empty(json!({}), Ok(json!({})))] #[case::numeric_user_id( json!({"user_id": 123}), - Err(Error::InvalidRequest("metadata.user_id must be a string, got 123".to_string().into())), + Err(invalid_value("metadata.user_id", "a string", json!(123))), )] #[case::boolean_user_id( json!({"user_id": true}), - Err(Error::InvalidRequest("metadata.user_id must be a string, got true".to_string().into())), + Err(invalid_value("metadata.user_id", "a string", json!(true))), )] #[case::not_an_object( json!(["u-1"]), - Err(Error::InvalidRequest(r#"metadata must be an object, got ["u-1"]"#.to_string().into())), + Err(invalid_value("metadata", "an object", json!(["u-1"]))), )] fn validate_anthropic_api_metadata_passes_only_a_string_user_id( #[case] metadata: Value, #[case] expected: Result, ) { - assert_eq!( - validate_anthropic_api_metadata(&metadata).map_err(|error| error.to_string()), - expected.map_err(|error| error.to_string()), - ); + assert_eq!(validate_anthropic_api_metadata(&metadata), expected); + } + + fn invalid_value(field: &'static str, expected: &'static str, actual: Value) -> Error { + Error::InvalidRequest(crate::ErrorDetail::InvalidValue { + field, + expected, + actual, + }) } #[rstest] diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs index e813d4ac8cc..82119cecfa5 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -123,7 +123,7 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { &self, error_message: String, status_code: u16, - headers: Vec<(String, String)>, + headers: reqwest::header::HeaderMap, ) -> Error { error_class(error_message, status_code, headers) } diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs index bc199f40938..e3a4f947433 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs @@ -17,6 +17,7 @@ const AWS_JSON_CONTENT_TYPE: &str = "application/x-amz-json-1.1"; const TARGET_HEADER: &str = "X-Amz-Target"; const CONTENT_TYPE_HEADER: &str = "Content-Type"; const UNSUPPORTED_DOCUMENT: &str = "UnsupportedDocumentException"; +const SYNC_API_GUIDANCE: &str = "aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported"; const SYNC_DOCUMENT_MAX_BYTES: usize = 10 * 1024 * 1024; const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; @@ -321,20 +322,18 @@ struct AwsError { /// Textract answers both an unsupported format and a multi-page PDF or TIFF /// with a bare "unsupported document format", which reads like a corrupt file. /// Say what the synchronous API accepts. -pub(super) fn error_class(body: String, status: u16, headers: Vec<(String, String)>) -> Error { - let unsupported = serde_json::from_str::(&body) - .ok() - .filter(|error| error.kind.ends_with(UNSUPPORTED_DOCUMENT)); - Error::Provider { - status, - body: match unsupported { - Some(error) => format!( - "{UNSUPPORTED_DOCUMENT}: {}. aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported", - error.message - ), - None => body, +pub(super) fn error_class(body: String, status: u16, headers: reqwest::header::HeaderMap) -> Error { + match serde_json::from_str::(&body) { + Ok(error) if error.kind.ends_with(UNSUPPORTED_DOCUMENT) => Error::RejectedDocument { + provider: "AWS Textract", + message: error.message, + guidance: SYNC_API_GUIDANCE, + }, + Ok(_) | Err(_) => Error::Provider { + status, + body, + headers: Box::new(headers), }, - headers, } } @@ -408,11 +407,11 @@ mod tests { fn a_model_outside_the_operations_is_refused_with_the_supported_names(#[case] model: &str) { let error = TextractOperation::from_model(model).unwrap_err(); - assert_eq!( - error.to_string(), - format!( - "invalid model: aws_textract has no model {model:?} - use one of: detect-document-text, analyze-document" - ) + let message = error.to_string(); + assert!(message.contains(&format!("{model:?}")), "{message}"); + assert!( + message.contains("detect-document-text, analyze-document"), + "{message}" ); assert_eq!(error.http_status_code(), Some(400)); } @@ -542,43 +541,46 @@ mod tests { #[rstest] #[case::bare_type( r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#, - Some("Request has unsupported document format") + "Request has unsupported document format" )] #[case::namespaced_type( r#"{"__type":"com.amazonaws.textract#UnsupportedDocumentException","Message":"bad"}"#, - Some("bad") + "bad" )] #[case::lowercase_message( r#"{"__type":"UnsupportedDocumentException","message":"bad"}"#, - Some("bad") + "bad" )] - #[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#, None)] - #[case::json_without_a_type(r#"{"Message":"no"}"#, None)] - #[case::not_json("bad gateway", None)] - fn only_an_unsupported_document_gains_the_sync_api_hint( + fn an_unsupported_document_is_the_callers_fault_with_the_sync_api_hint( #[case] body: &str, - #[case] hinted_message: Option<&str>, + #[case] provider_message: &str, ) { - let response_headers = vec![("x-amzn-requestid".to_string(), "abc".to_string())]; + let error = error_class(body.into(), 400, reqwest::header::HeaderMap::new()); - let Error::Provider { - status, - body: reported, - headers, - } = error_class(body.into(), 400, response_headers.clone()) - else { - panic!("expected a provider error"); - }; + let message = error.to_string(); + assert!(error.is_request()); + assert_eq!(error.http_status_code(), Some(400)); + assert!(error.upstream_response().is_none()); + assert!(message.contains(provider_message), "{message}"); + assert!(message.contains(HINT), "{message}"); + } - assert_eq!(status, 400); - assert_eq!(headers, response_headers); - match hinted_message { - Some(message) => { - assert!(reported.contains(message), "{reported}"); - assert!(reported.contains(HINT), "{reported}"); - } - None => assert_eq!(reported, body), - } + #[rstest] + #[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#)] + #[case::json_without_a_type(r#"{"Message":"no"}"#)] + #[case::not_json("bad gateway")] + fn any_other_rejection_passes_through_untouched(#[case] body: &str) { + let response_headers = reqwest::header::HeaderMap::from_iter([( + reqwest::header::HeaderName::from_static("x-amzn-requestid"), + reqwest::header::HeaderValue::from_static("abc"), + )]); + + let error = error_class(body.into(), 403, response_headers.clone()); + + assert_eq!( + error.upstream_response(), + Some((403, body, &response_headers)) + ); } #[rstest] diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs index 4b27a191734..34213926942 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -104,7 +104,7 @@ impl BaseOcrConfig for TextractDetectTextConfig { &self, error_message: String, status_code: u16, - headers: Vec<(String, String)>, + headers: reqwest::header::HeaderMap, ) -> Error { error_class(error_message, status_code, headers) } @@ -237,7 +237,7 @@ mod tests { let error = TextractDetectTextConfig.get_error_class( r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#.into(), 400, - Vec::new(), + reqwest::header::HeaderMap::new(), ); assert!( diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs index 72e321547e3..bec8bee3f34 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs @@ -120,7 +120,8 @@ impl BaseOcrConfig for AzureAICohereParseConfig { impl AzureAICohereParseConfig { fn get_complete_url(&self, base: &str) -> Result { - let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; + let mut url = + reqwest::Url::parse(base).map_err(|error| Error::request_field("api_base", error))?; if !matches!(url.scheme(), "http" | "https") { return Err(invalid_api_base()); } @@ -133,13 +134,14 @@ impl AzureAICohereParseConfig { ApiUrl::parse(url.as_str()) .and_then(|url| url.complete_path(&AZURE_COHERE_PARSE_PATH)) .map(|url| url.into_string()) - .map_err(|_| invalid_api_base()) + .map_err(|error| Error::request_field("api_base", error)) } } fn invalid_api_base() -> Error { Error::RequestField { path: "api_base".into(), + cause: None, } } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs index 8cbf2455e01..7da685febf9 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs @@ -26,7 +26,7 @@ pub(super) async fn resolve_entra( .get_azure_ad_token(config, env_lookup) .await .or_else(|error| match error { - litellm_auth::Error::EmptyCallerCredential(_) => Ok(None), + litellm_auth::Error::EmptyCallerCredential { .. } => Ok(None), other => Err(other), }) .map(|credential| { @@ -48,7 +48,7 @@ pub(super) fn validate_destination( && credential_source != InputSource::Request { return Err(litellm_auth::Error::InvalidConfiguration( - "host credentials cannot be sent to a request-controlled Azure endpoint".into(), + litellm_auth::ErrorDetail::RequestControlledEndpoint { provider: "Azure" }, ) .into()); } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 2a50989f4e4..d7233db0a76 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -13,7 +13,7 @@ use tokio::time::Instant; use crate::base_llm::ocr::{ document::InlineDocument, - error::Error, + error::{Error, PagesFault, PollUrl}, handler::{CallHooks, OcrClient, read_json_response}, settings::OcrSettings, transformation::{ @@ -185,7 +185,10 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { ) -> Result { let endpoint = nonblank(request.connection.api_base.clone()) .or_else(|| nonblank(request.connection.secret(AZURE_DI_ENDPOINT_ENV))) - .ok_or_else(|| Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?; + .ok_or(Error::Auth(litellm_auth::Error::MissingApiBase { + provider: "Azure Document Intelligence", + guidance: "Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base", + }))?; self.build_ocr_url( &endpoint, &request.model, @@ -256,14 +259,12 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result, Error> Some(Value::Array(pages)) if pages.iter().all(Value::is_number) => pages .iter() .map(|page| { - let page = page - .as_i64() - .ok_or_else(|| Error::Pages("page index is out of range".into()))?; + let page = page.as_i64().ok_or(Error::Pages(PagesFault::OutOfRange))?; if page < 0 { - return Err(Error::Pages("negative page index".into())); + return Err(Error::Pages(PagesFault::Negative)); } page.checked_add(1) - .ok_or_else(|| Error::Pages("page index is out of range".into())) + .ok_or(Error::Pages(PagesFault::OutOfRange)) }) .collect::, _>>()? .into_iter() @@ -276,7 +277,7 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result, Error> token .as_str() .map(str::trim) - .ok_or_else(|| Error::Pages("expected only integers or only strings".into())) + .ok_or(Error::Pages(PagesFault::Mixed)) }) .collect::, _>>()? .join(","), @@ -286,13 +287,11 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result, Error> .collect::>() .join(","), Some(_) => { - return Err(Error::Pages( - "expected an array of integers or strings, or a native page range".into(), - )); + return Err(Error::Pages(PagesFault::Shape)); } }; if !normalized.split(',').all(valid_page_token) { - return Err(Error::Pages("invalid native page range".into())); + return Err(Error::Pages(PagesFault::Range)); } Ok(Some(normalized)) } @@ -458,8 +457,14 @@ async fn read_operation_response( .and_then(|value| value.to_str().ok()) .ok_or(Error::PollLocation)? .to_string(); - let original = Url::parse(original_url).map_err(|_| Error::PollOrigin)?; - let operation = Url::parse(&location).map_err(|_| Error::PollOrigin)?; + let original = Url::parse(original_url).map_err(|source| Error::InvalidPollUrl { + url: PollUrl::Submission, + source, + })?; + let operation = Url::parse(&location).map_err(|source| Error::InvalidPollUrl { + url: PollUrl::OperationLocation, + source, + })?; if original.origin() != operation.origin() || !operation.username().is_empty() || operation.password().is_some() @@ -504,8 +509,7 @@ async fn poll_operation( let response = tokio::time::timeout_at(deadline, litellm_http::request::http_request(builder)) .await - .map_err(|_| Error::PollTimeout)? - .map_err(litellm_http::transport::Error::from)?; + .map_err(|_| Error::PollTimeout)??; let retry = response .headers() .get(reqwest::header::RETRY_AFTER) @@ -578,9 +582,7 @@ impl AzureDocumentIntelligenceOcrConfig { ) .into_string() }) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } async fn resolve_headers( @@ -623,7 +625,10 @@ impl AzureDocumentIntelligenceOcrConfig { } let token = super::super::common_utils::resolve_entra(auth, config, env_lookup) .await? - .ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?; + .ok_or(Error::MissingCredential { + credential: "Azure Document Intelligence credentials", + guidance: "set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID", + })?; super::super::common_utils::validate_destination(connection, token.source())?; Ok( std::iter::once(("Authorization".into(), format!("Bearer {}", token.value()))) @@ -842,7 +847,7 @@ mod tests { .unwrap_err(); assert!( - error + litellm_error::Report::new(&error) .to_string() .contains("request-controlled Azure endpoint") ); diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 83d88bbb3cb..2f7957fe56e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -169,7 +169,10 @@ impl AzureAiOcrConfig { } let key = super::common_utils::resolve_entra(auth, config, env_lookup) .await? - .ok_or(Error::MissingAzureAiCredentials)?; + .ok_or(Error::MissingCredential { + credential: "Azure AI credentials", + guidance: "set AZURE_AI_API_KEY or configure Entra ID", + })?; super::common_utils::validate_destination(connection, key.source())?; Ok(bearer_headers(connection, key.value())) } @@ -183,9 +186,7 @@ impl AzureAiOcrConfig { ApiUrl::parse(&base) .and_then(|url| url.complete_path(&AZURE_AI_OCR_PATH)) .map(|url| url.into_string()) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } } @@ -302,7 +303,7 @@ mod tests { .unwrap_err(); assert!( - error + litellm_error::Report::new(&error) .to_string() .contains("request-controlled Azure endpoint") ); diff --git a/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs b/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs index a283ca84089..6df321b4b16 100644 --- a/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs +++ b/litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs @@ -5,6 +5,15 @@ use futures_util::{Stream, StreamExt, stream, stream::BoxStream}; pub type ByteStream = BoxStream<'static, Result>; +/// The transport error a body stream carried as an `io::Error`, however many times framing +/// wrapped it on the way. +pub(crate) fn tunnelled_http(error: &io::Error) -> Option<&litellm_http::Error> { + let inner = error.get_ref()?; + inner + .downcast_ref::() + .or_else(|| tunnelled_http(inner.downcast_ref::()?)) +} + pub trait StreamTransformer { type Input; type Output; diff --git a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs index 30bff76651e..8ba4030540c 100644 --- a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use litellm_llms_types::formats::chat_completions::{ ChatCompletionsResponse, ChatMessage, ChatMessageContent, }; @@ -29,16 +30,6 @@ const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; -/// Why a request cannot be served by the Rust path. -/// -/// The core declines rather than guessing: the host turns this into a -/// transparent fallback to the Python implementation, which covers the full -/// surface. Acceptance is an allowlist, so a parameter or message shape the -/// core has never seen declines by construction instead of being translated -/// wrong. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Unsupported(pub &'static str); - pub trait BaseConfig: Sync { fn secret_names(&self) -> Vec<&'static str>; @@ -97,7 +88,7 @@ pub trait BaseConfig: Sync { &self, messages: &[ChatMessage], optional_params: &Map, - ) -> Option { + ) -> Option { unsupported_param( self.supported_openai_param_mappings(), self.config_params(), @@ -111,13 +102,13 @@ pub fn unsupported_param( supported: &'static [(&'static str, &'static str)], config: &'static [&'static str], optional_params: &Map, -) -> Option { +) -> Option { if optional_params .get(STREAM_PARAM) .and_then(Value::as_bool) .unwrap_or(false) { - return Some(Unsupported("streaming")); + return Some(Decline::new("streaming")); } optional_params .keys() @@ -128,29 +119,29 @@ pub fn unsupported_param( .any(|(_, provider_name)| *provider_name == key) && !config.contains(&key.as_str()) }) - .then_some(Unsupported("unrecognized request parameter")) + .then_some(Decline::new("unrecognized request parameter")) } /// Message shapes the core can translate faithfully: text content, either a /// plain string or a non-empty list of parts that are all /// `{"type": "text", "text": ...}`. Tool calls, tool results, and multimodal /// parts decline so Python's fuller translation handles them. -pub fn unsupported_message(message: &ChatMessage) -> Option { +pub fn unsupported_message(message: &ChatMessage) -> Option { if message .extra .keys() .any(|key| !IGNORABLE_MESSAGE_FIELDS.contains(&key.as_str())) { - return Some(Unsupported("unrecognized message field")); + return Some(Decline::new("unrecognized message field")); } if !matches!(message.role.as_str(), "system" | "user" | "assistant") { - return Some(Unsupported("unrecognized message role")); + return Some(Decline::new("unrecognized message role")); } match &message.content { - None => Some(Unsupported("message without content")), + None => Some(Decline::new("message without content")), Some(ChatMessageContent::Text(_)) => None, Some(ChatMessageContent::Parts(parts)) if parts.is_empty() => { - Some(Unsupported("message without content")) + Some(Decline::new("message without content")) } Some(ChatMessageContent::Parts(parts)) => parts .iter() @@ -159,6 +150,6 @@ pub fn unsupported_message(message: &ChatMessage) -> Option { || part.get("text").and_then(Value::as_str).is_none() || part.as_object().is_some_and(|object| object.len() != 2) }) - .then_some(Unsupported("non-text message content")), + .then_some(Decline::new("non-text message content")), } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs index b32ea4bf73a..898c24369c3 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs @@ -1,9 +1,6 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError, mime::Mime}; -use litellm_http::{ - media::{DownloadPolicy, Error as MediaError, MediaFetcher}, - transport::Error as TransportError, -}; +use litellm_http::media::{DownloadPolicy, MediaFetcher}; use reqwest::Url; use crate::base_llm::ocr::{ @@ -61,9 +58,7 @@ pub async fn inline_remote_document( validate_inline_document(&document)?; return Ok(document); } - let url = Url::parse(source).map_err(|_| Error::RequestField { - path: "document URL".into(), - })?; + let url = Url::parse(source).map_err(|error| Error::request_field("document URL", error))?; let downloaded = fetcher .fetch( url, @@ -74,7 +69,7 @@ pub async fn inline_remote_document( }, ) .await - .map_err(map_media_error)?; + .map_err(Error::DocumentDownload)?; let result = document.with_source(format!( "data:{};base64,{}", downloaded.content_type, @@ -84,28 +79,6 @@ pub async fn inline_remote_document( Ok(result) } -fn map_media_error(error: MediaError) -> Error { - match error { - MediaError::BlockedUrl => Error::BlockedDocumentUrl, - MediaError::DownloadDisabled => Error::DownloadDisabled, - MediaError::DownloadTooLarge => Error::DownloadTooLarge, - MediaError::TooManyRedirects => Error::TooManyRedirects, - MediaError::MissingRedirectLocation => Error::MissingRedirectLocation, - MediaError::InvalidRedirect => Error::InvalidRedirect, - MediaError::Http(status) => TransportError::Http { - status, - body: "OCR document download failed".into(), - } - .into(), - MediaError::Timeout => TransportError::Http { - status: 408, - body: "OCR document download timed out".into(), - } - .into(), - MediaError::Transport(error) => error.into(), - } -} - #[cfg(test)] mod tests { use std::collections::BTreeMap as Map; @@ -119,20 +92,17 @@ mod tests { } } - #[test] - fn decodes_data_urls_and_limits_decoded_size() { - for (source, expected) in [ - ("data:application/pdf;base64,YWJj", b"abc".as_slice()), - ("DATA:application/pdf;BASE64,YWI", b"ab".as_slice()), - ("data:,a%20b%00%FF", b"a b\0\xff".as_slice()), - ] { - let inline = InlineDocument::parse(source).unwrap().unwrap(); - assert_eq!(inline.decode(expected.len()).unwrap(), expected); - assert!(matches!( - inline.decode(expected.len() - 1), - Err(Error::InlineDocumentTooLarge) - )); - } + #[rstest::rstest] + #[case::base64("data:application/pdf;base64,YWJj", b"abc")] + #[case::uppercase_unpadded("DATA:application/pdf;BASE64,YWI", b"ab")] + #[case::percent_encoded("data:,a%20b%00%FF", b"a b\0\xff")] + fn decodes_data_urls_and_limits_decoded_size(#[case] source: &str, #[case] expected: &[u8]) { + let inline = InlineDocument::parse(source).unwrap().unwrap(); + assert_eq!(inline.decode(expected.len()).unwrap(), expected); + assert!(matches!( + inline.decode(expected.len() - 1), + Err(Error::InlineDocumentTooLarge) + )); } #[test] @@ -150,15 +120,12 @@ mod tests { ); } - #[test] - fn rejects_invalid_inline_documents() { - for source in [ - "https://example.com/document.pdf", - "data:application/pdf;base64", - "data:application/pdf;base64,INVALID!", - ] { - assert!(validate_inline_document(&document(source)).is_err()); - } + #[rstest::rstest] + #[case::remote("https://example.com/document.pdf")] + #[case::no_payload("data:application/pdf;base64")] + #[case::invalid_base64("data:application/pdf;base64,INVALID!")] + fn rejects_invalid_inline_documents(#[case] source: &str) { + assert!(validate_inline_document(&document(source)).is_err()); } #[tokio::test] @@ -190,16 +157,14 @@ mod tests { clippy::disallowed_types, reason = "the pool has no default-header setting to stand in for provider credentials" )] - let provider_http = litellm_http::Client::for_test( + let provider_http = litellm_http::Client::wrap_for_test( reqwest::Client::builder() .default_headers(provider_headers) .build() .unwrap(), ); - let client = crate::base_llm::ocr::handler::OcrClient::for_test( - provider_http, - litellm_http::Client::no_redirect_for_test(), - ); + let client = + crate::base_llm::ocr::handler::OcrClient::for_test().with_provider_http(provider_http); let converted = inline_remote_document( client.document_fetcher(), OcrDocument::ImageUrl { diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 392dd56412c..d51fe996dc4 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -1,53 +1,62 @@ +use std::{io, path::PathBuf, sync::Arc}; + +use litellm_core_utils::call_arguments::ArgumentError; +use litellm_error::Source; +use reqwest::header::HeaderMap; + +static NO_HEADERS: std::sync::LazyLock = std::sync::LazyLock::new(HeaderMap::new); + #[derive(Clone, Debug, thiserror::Error)] pub enum Error { #[error("upstream OCR error ({status}): {body}")] Provider { status: u16, body: String, - headers: Vec<(String, String)>, + headers: Box, }, #[error("File is empty or could not be read")] EmptyFile, - #[error("Failed to read OCR file {}: {source}", path.display())] + #[error("could not read OCR file {}", path.display())] FileRead { - path: std::path::PathBuf, + path: PathBuf, #[source] - source: std::sync::Arc, + cause: Arc, }, - #[error("OCR document preparation task failed: {0}")] - DocumentTask(#[source] std::sync::Arc), - #[error("Invalid MIME type: {0}")] + #[error("OCR document preparation task failed")] + DocumentTask(#[source] Arc), + #[error("invalid MIME type {0:?}")] InvalidMimeType(String), - #[error( - "Cohere Parse only accepts `image_url` documents; document_url and PDF inputs are not supported" - )] - CohereImageOnly, + #[error("{provider} cannot read this document - {guidance}")] + UnsupportedDocument { + provider: &'static str, + guidance: &'static str, + }, + #[error("{provider} cannot read this document: {message} - {guidance}")] + RejectedDocument { + provider: &'static str, + message: String, + guidance: &'static str, + }, #[error("Invalid `req_format`. Expected 'native' or 'litellm'.")] RequestFormat, #[error("invalid OCR request field: {path}")] - RequestField { path: String }, + RequestField { + path: String, + #[source] + cause: Option, + }, + #[error("invalid OCR optional params")] + Arguments(#[from] ArgumentError), #[error("missing required field: {0}")] MissingField(&'static str), - #[error("Document URL is required")] + #[error("document URL is required")] MissingDocumentUrl, #[error("invalid OCR document data URI")] InvalidDataUri, - #[error( - "Reducto requires a reducto:// id or a data URI; plain HTTP URLs are not supported, upload the file first" - )] - ReductoSource, #[error("inline OCR document exceeds the size limit")] InlineDocumentTooLarge, - #[error("OCR document URL is blocked by network policy")] - BlockedDocumentUrl, - #[error("OCR document downloads are disabled")] - DownloadDisabled, - #[error("OCR document download exceeds the size limit")] - DownloadTooLarge, - #[error("OCR document download exceeded the redirect limit")] - TooManyRedirects, #[error("invalid OCR pages: {0}")] - Pages(String), + Pages(PagesFault), #[error("invalid OCR features")] Features, #[error("OCR model cannot be a dot segment")] @@ -55,25 +64,29 @@ pub enum Error { #[error("OCR response exceeds the size limit of {limit} bytes")] TooLarge { limit: usize }, #[error("invalid OCR response field: {path}")] - ResponseField { path: String }, + ResponseField { + path: String, + #[source] + cause: Option, + }, #[error("OCR response is missing non-empty content")] EmptyContent, - #[error("OCR document redirect is missing a location")] - MissingRedirectLocation, - #[error("OCR document redirect location is invalid")] - InvalidRedirect, #[error("OCR operation ended with status {0}")] OperationStatus(String), #[error("OCR response numeric value is out of range: {0}")] NumericRange(&'static str), #[error("OCR accepted response is missing a valid operation-location")] PollLocation, + #[error("OCR {url} is not a valid URL")] + InvalidPollUrl { + url: PollUrl, + #[source] + source: url::ParseError, + }, #[error("OCR operation-location must use the submission origin without credentials")] PollOrigin, #[error("OCR polling timed out")] PollTimeout, - #[error("unsupported by the rust path: {0}")] - Unsupported(&'static str), #[error("invalid provider: {0}")] InvalidProvider(String), #[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))] @@ -82,108 +95,266 @@ pub enum Error { model: String, supported: &'static [&'static str], }, - #[error("invalid request: {0}")] - InvalidRequest(String), - #[error("invalid response: {0}")] - InvalidResponse(String), - #[error( - "invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID" - )] - MissingAzureAiCredentials, - #[error( - "invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID" - )] - MissingAzureDocumentIntelligenceCredentials, - #[error( - "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" - )] - MissingReductoApiKey, - #[error("secret resolution failed: {0}")] - Secret(#[source] std::sync::Arc), + #[error("Missing {credential} - {guidance}")] + MissingCredential { + credential: &'static str, + guidance: &'static str, + }, + #[error("secret resolution failed")] + Secret(#[source] Arc), + #[error("the OCR host driver failed")] + Host(#[source] litellm_host::machine::MachineFault), + #[error("the OCR host call raised an exception")] + HostCall, #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] - Transport(#[from] litellm_http::transport::Error), - #[error(transparent)] Params(#[from] litellm_core_utils::params::Error), #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] Http(#[from] litellm_http::Error), + #[error(transparent)] + DocumentDownload(litellm_http::Error), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum PagesFault { + #[error("page index is out of range")] + OutOfRange, + #[error("negative page index")] + Negative, + #[error("expected only integers or only strings")] + Mixed, + #[error("expected an array of integers or strings, or a native page range")] + Shape, + #[error("invalid native page range")] + Range, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum PollUrl { + #[error("submission URL")] + Submission, + #[error("operation-location")] + OperationLocation, } impl From for Error { fn from(fault: litellm_host::machine::MachineFault) -> Self { - use litellm_host::machine::MachineFault; - Self::InvalidRequest(match fault { - MachineFault::Abandoned => "OCR host driver was abandoned".into(), - MachineFault::Protocol(message) => format!("OCR {message}"), - }) - } -} - -impl From for Error { - fn from(error: litellm_core_utils::call_arguments::ArgumentError) -> Self { - Self::RequestField { - path: format!("optional_params.{}", error.path), - } + Self::Host(fault) } } impl Error { - pub fn http_status_code(&self) -> Option { - match self { - Self::Provider { status, .. } - | Self::Transport(litellm_http::transport::Error::Http { status, .. }) => Some(*status), - error if error.is_request() => Some(400), - _ => None, + pub fn request_field( + path: impl Into, + cause: impl std::error::Error + Send + Sync + 'static, + ) -> Self { + Self::RequestField { + path: path.into(), + cause: Some(Source::new(cause)), } } - pub fn is_request(&self) -> bool { - matches!( - self, + pub fn response_field( + path: impl Into, + cause: impl std::error::Error + Send + Sync + 'static, + ) -> Self { + Self::ResponseField { + path: path.into(), + cause: Some(Source::new(cause)), + } + } + + pub fn upstream_response(&self) -> Option<(u16, &str, &HeaderMap)> { + match self { + Self::Provider { + status, + body, + headers, + } => Some((*status, body, headers)), + Self::DocumentDownload(error) => error.status().map(|status| { + let body = if error.is_timeout() { + "OCR document download timed out" + } else { + "OCR document download failed" + }; + (status, body, &*NO_HEADERS) + }), + Self::Http(error) => error + .upstream_response() + .map(|(status, body)| (status, body, &*NO_HEADERS)), Self::EmptyFile - | Self::InvalidMimeType(_) - | Self::CohereImageOnly - | Self::RequestFormat - | Self::RequestField { .. } - | Self::MissingField(_) - | Self::MissingDocumentUrl - | Self::InvalidDataUri - | Self::ReductoSource - | Self::InlineDocumentTooLarge - | Self::BlockedDocumentUrl - | Self::DownloadDisabled - | Self::DownloadTooLarge - | Self::TooManyRedirects - | Self::Pages(_) - | Self::Features - | Self::DotModel - | Self::InvalidRequest(_) - | Self::InvalidProvider(_) - | Self::InvalidModel { .. } - | Self::Params(_) - | Self::Headers(_) - | Self::Http(_) - | Self::Auth(litellm_auth::Error::MissingParam { .. }) - ) + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::TooLarge { .. } + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::MissingCredential { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall + | Self::Auth(_) + | Self::Params(_) + | Self::Headers(_) => None, + } + } + + pub fn is_missing_credential(&self) -> bool { + match self { + Self::MissingCredential { .. } => true, + Self::Auth(error) => error.is_missing_credential(), + Self::Provider { .. } + | Self::EmptyFile + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::TooLarge { .. } + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall + | Self::Params(_) + | Self::Headers(_) + | Self::Http(_) + | Self::DocumentDownload(_) => false, + } + } + + pub fn http_status_code(&self) -> Option { + if let Some((status, ..)) = self.upstream_response() { + return Some(status); + } + if let Self::Http(error) = self + && let Some(status) = error.status() + { + return Some(status); + } + self.is_request().then_some(400) + } + + pub fn is_request(&self) -> bool { + match self { + Self::EmptyFile + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::Params(_) + | Self::Headers(_) => true, + Self::Auth(error) => error.is_request(), + Self::Http(error) | Self::DocumentDownload(error) => error.is_request(), + Self::Provider { .. } + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::TooLarge { .. } + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout + | Self::MissingCredential { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall => false, + } } pub fn is_response(&self) -> bool { - matches!( - self, + match self { Self::TooLarge { .. } - | Self::ResponseField { .. } - | Self::EmptyContent - | Self::MissingRedirectLocation - | Self::InvalidRedirect - | Self::OperationStatus(_) - | Self::NumericRange(_) - | Self::PollLocation - | Self::PollOrigin - | Self::PollTimeout - | Self::InvalidResponse(_) - ) + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::InvalidPollUrl { .. } + | Self::PollOrigin + | Self::PollTimeout => true, + Self::Http(error) | Self::DocumentDownload(error) => error.is_response(), + Self::Provider { .. } + | Self::EmptyFile + | Self::FileRead { .. } + | Self::DocumentTask(_) + | Self::InvalidMimeType(_) + | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } + | Self::RequestFormat + | Self::RequestField { .. } + | Self::Arguments(_) + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::InlineDocumentTooLarge + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::InvalidProvider(_) + | Self::InvalidModel { .. } + | Self::MissingCredential { .. } + | Self::Secret(_) + | Self::Host(_) + | Self::HostCall + | Self::Auth(_) + | Self::Params(_) + | Self::Headers(_) => false, + } } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 2cf955ff0c6..7bbc751c614 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -7,8 +7,8 @@ use litellm_host::interceptors::WireRequest; use litellm_http::{ Client, ClientVariant, HttpClientConfig, HttpClientPool, media::{MediaFetcher, UrlPolicy}, - outbound::{OutboundRequest, RequestSigner}, - transport, + outbound::{RequestSigner, json, signed}, + request::header_pairs, }; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; @@ -86,21 +86,27 @@ impl OcrClient { } #[cfg(any(test, feature = "test-support"))] - pub fn for_test(provider_http: Client, no_redirect_http: Client) -> Self { + pub fn for_test() -> Self { Self { secrets: Arc::new( - litellm_secrets::source::EnvironmentSecrets::python_compatible( - provider_http.clone(), - ), + litellm_secrets::source::EnvironmentSecrets::python_compatible(Client::for_test()), ), - provider_http, - polling_http: no_redirect_http.clone(), - document_fetcher: MediaFetcher::for_test(no_redirect_http), + provider_http: Client::for_test(), + polling_http: Client::for_test(), + document_fetcher: MediaFetcher::for_test(Client::for_test()), auth: Arc::new(AuthServices::default()), settings: OcrSettings::default(), } } + #[cfg(any(test, feature = "test-support"))] + pub fn with_provider_http(self, provider_http: Client) -> Self { + Self { + provider_http, + ..self + } + } + #[cfg(any(test, feature = "test-support"))] pub fn with_settings(self, settings: OcrSettings) -> Self { Self { settings, ..self } @@ -122,29 +128,17 @@ pub async fn ocr( ) -> Result { let http = config.prepare_request(request, client, hooks).await?; let url = http.url().to_string(); - let headers = http.headers().to_vec(); - let response = http - .send(client.provider_http()) + let headers = header_pairs(http.headers()); + let response = client + .provider_http() + .execute(http) .await - .map_err(transport_error)?; + .map_err(litellm_http::Error::from)?; if !response.status().is_success() { - let headers = response - .headers() - .iter() - .filter_map(|(name, value)| { - value - .to_str() - .ok() - .map(|value| (name.to_string(), value.to_string())) - }) - .collect(); - return match read_response_bytes(response, request.connection.max_response_bytes).await { - Err(Error::Transport(transport::Error::Http { status, body })) => { - Err(config.get_error_class(body, status, headers)) - } - Err(error) => Err(error), - Ok(_) => unreachable!("non-success response produces an HTTP error"), - }; + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = read_error_body(response, request.connection.max_response_bytes).await?; + return Err(config.get_error_class(body, status, headers)); } let context = OcrResponseContext { client, @@ -172,43 +166,41 @@ pub async fn read_response_bytes( mut response: reqwest::Response, limit: usize, ) -> Result { - let status = response.status(); - if status.is_success() - && response - .content_length() - .is_some_and(|length| length > limit as u64) + if !response.status().is_success() { + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = read_error_body(response, limit).await?; + return Err(Error::Provider { + status, + body, + headers: Box::new(headers), + }); + } + if response + .content_length() + .is_some_and(|length| length > limit as u64) { return Err(Error::TooLarge { limit }); } let mut bytes = BytesMut::new(); - while let Some(chunk) = response.chunk().await.map_err(transport_error)? { - let remaining = limit.saturating_sub(bytes.len()); - if status.is_success() && chunk.len() > remaining { + while let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? { + if chunk.len() > limit.saturating_sub(bytes.len()) { return Err(Error::TooLarge { limit }); } - bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]); - if !status.is_success() && bytes.len() == limit { - break; - } - } - if !status.is_success() { - return Err(transport::Error::Http { - status: status.as_u16(), - body: String::from_utf8_lossy(&bytes).into_owned(), - } - .into()); + bytes.extend_from_slice(&chunk); } Ok(bytes.freeze()) } -pub fn transport_error(error: reqwest::Error) -> Error { - if error.is_timeout() { - return Error::Transport(transport::Error::Http { - status: 408, - body: "OCR request timed out".into(), - }); +async fn read_error_body(mut response: reqwest::Response, limit: usize) -> Result { + let mut bytes = BytesMut::new(); + while bytes.len() < limit + && let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? + { + let remaining = limit - bytes.len(); + bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]); } - transport::Error::from(error).into() + Ok(String::from_utf8_lossy(&bytes).into_owned()) } pub async fn transform_request_body( @@ -219,7 +211,7 @@ pub async fn transform_request_body( body: B, signer: Option<&dyn RequestSigner>, hooks: &dyn CallHooks, -) -> Result { +) -> Result { let composed = litellm_core_utils::call_arguments::compose_body( &request.optional_params, &body, @@ -232,20 +224,20 @@ pub async fn transform_request_body( if !changed.body.is_object() { return Err(Error::RequestField { path: "guardrail.body".into(), + cause: None, }); } config.validate_request_body(&changed.body)?; - let timeout = Some(request.connection.timeout); + let built = json( + url, + &changed.headers, + &changed.body, + Some(request.connection.timeout), + )?; Ok(match signer { - Some(signer) => OutboundRequest::signed_json( - url.into(), - changed.headers, - &changed.body, - timeout, - signer, - ), - None => OutboundRequest::json(url.into(), changed.headers, &changed.body, timeout), - }?) + Some(signer) => signed(built, signer)?, + None => built, + }) } fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest { @@ -261,10 +253,10 @@ pub fn build_http_request( url: String, headers: Vec<(String, String)>, body: &impl Serialize, -) -> Result { - Ok(OutboundRequest::json( - url, - headers, +) -> Result { + Ok(json( + &url, + &headers, body, Some(request.connection.timeout), )?) @@ -276,9 +268,8 @@ pub async fn guardrail_document( headers: &[(String, String)], hooks: &dyn CallHooks, ) -> Result<(OcrDocument, Vec<(String, String)>), Error> { - let body = serde_json::to_value(&request.document).map_err(|_| Error::RequestField { - path: "document".into(), - })?; + let body = serde_json::to_value(&request.document) + .map_err(|error| Error::request_field("document", error))?; let changed = hooks .before_provider_request(wire_request(url, headers, body)) .await?; @@ -292,6 +283,7 @@ pub fn body_document(body: &Value) -> Result { .and_then(Value::as_object) .ok_or_else(|| Error::RequestField { path: "body.document".into(), + cause: None, })?; let source = document .iter() @@ -311,8 +303,6 @@ mod tests { net::TcpListener, }; - use super::*; - #[rstest::rstest] #[tokio::test] async fn request_timeout_has_an_http_408_status() { @@ -322,16 +312,16 @@ mod tests { let _connection = listener.accept().await.unwrap(); tokio::time::sleep(Duration::from_secs(1)).await; }); - let error = litellm_http::Client::plain_for_test() + let error = litellm_http::Client::for_test() .get(format!("http://{address}")) .timeout(Duration::from_millis(10)) .send() .await .unwrap_err(); - assert!(matches!( - transport_error(error), - Error::Transport(transport::Error::Http { status: 408, .. }) - )); + assert_eq!( + Error::from(litellm_http::Error::from(error)).http_status_code(), + Some(408) + ); server.abort(); } @@ -347,7 +337,7 @@ mod tests { socket.write_all(response.as_bytes()).await.unwrap(); std::future::pending::<()>().await; }); - let response = litellm_http::Client::plain_for_test() + let response = litellm_http::Client::for_test() .get(format!("http://{address}")) .send() .await @@ -402,10 +392,7 @@ mod tests { .await .unwrap_err(); - let Error::Transport(litellm_http::transport::Error::Http { status, body }) = error else { - panic!("unexpected error: {error}"); - }; - assert_eq!(status, 429); - assert_eq!(body, prefix); + let (status, body, _) = error.upstream_response().expect("the provider answered"); + assert_eq!((status, body), (429, prefix.as_str())); } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index 2da0e1abbfc..19dd8e24226 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,15 +1,15 @@ -use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; -use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; - -use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; -use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; -use litellm_http::outbound::{OutboundRequest, RequestSigner}; -use litellm_secrets::source::Secrets; use serde::{ Serialize, de::{DeserializeOwned, IntoDeserializer}, }; use serde_json::{Map, Value}; +use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; + +use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; +use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; +use litellm_http::outbound::RequestSigner; +use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; +use litellm_secrets::source::Secrets; use crate::base_llm::ocr::{ error::Error, @@ -194,17 +194,13 @@ pub struct DecodedOcrResponse { pub fn decode_request_value(value: Value, prefix: &str) -> Result { serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| { - Error::RequestField { - path: format!("{prefix}.{}", error.path()), - } + Error::request_field(format!("{prefix}.{}", error.path()), error.into_inner()) }) } pub fn decode_response_value(value: Value, prefix: &str) -> Result { serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| { - Error::ResponseField { - path: format!("{prefix}.{}", error.path()), - } + Error::response_field(format!("{prefix}.{}", error.path()), error.into_inner()) }) } @@ -213,19 +209,15 @@ pub fn decode_response( native: bool, ) -> Result, Error> { let mut deserializer = serde_json::Deserializer::from_slice(bytes); - let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| { - Error::ResponseField { - path: error.path().to_string(), - } - })?; - deserializer.end().map_err(|_| Error::ResponseField { - path: "response".into(), - })?; + let data = serde_path_to_error::deserialize(&mut deserializer) + .map_err(|error| Error::response_field(error.path().to_string(), error.into_inner()))?; + deserializer + .end() + .map_err(|error| Error::response_field("response", error))?; let native = if native { Some( - serde_json::from_slice(bytes).map_err(|_| Error::ResponseField { - path: "response".into(), - })?, + serde_json::from_slice(bytes) + .map_err(|error| Error::response_field("response", error))?, ) } else { None @@ -369,12 +361,12 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { &self, error_message: String, status_code: u16, - headers: Vec<(String, String)>, + headers: reqwest::header::HeaderMap, ) -> Error { Error::Provider { status: status_code, body: error_message, - headers, + headers: Box::new(headers), } } @@ -391,7 +383,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> impl Future> + Send { + ) -> impl Future> + Send { async move { let params = self.map_ocr_params(&request.optional_params, &request.model)?; let environment = self.validate_environment(request, client).await?; diff --git a/litellm-rust/crates/llms/src/base_llm/sse.rs b/litellm-rust/crates/llms/src/base_llm/sse.rs index 5ff2e3adddd..23f0c1b9389 100644 --- a/litellm-rust/crates/llms/src/base_llm/sse.rs +++ b/litellm-rust/crates/llms/src/base_llm/sse.rs @@ -15,12 +15,15 @@ use bytes::Bytes; use futures_util::{StreamExt, TryStreamExt, future::ready, stream::BoxStream}; use litellm_framer::{ - frames, + SseError, frames, sse::{SseCodec, SseEvent, decode_all}, }; use serde::{Serialize, de::DeserializeOwned}; -use crate::{Error, ErrorDetail, base_llm::base_model_iterator::ByteStream}; +use crate::{ + Error, ErrorDetail, + base_llm::base_model_iterator::{ByteStream, tunnelled_http}, +}; pub type EventStream = BoxStream<'static, Result>; @@ -70,8 +73,13 @@ pub fn invalid(error: serde_json::Error) -> Error { Error::InvalidResponse(ErrorDetail::invalid("stream event", error)) } -fn framing(error: litellm_framer::SseError) -> Error { - Error::InvalidResponse(ErrorDetail::failed("stream framing", error)) +fn framing(error: SseError) -> Error { + match &error { + SseError::Body(body) if let Some(http) = tunnelled_http(body) => Error::Http(http.clone()), + SseError::Body(_) | SseError::InvalidUtf8(_) => { + Error::InvalidResponse(ErrorDetail::failed("stream framing", error)) + } + } } #[cfg(test)] @@ -94,6 +102,45 @@ mod tests { assert_eq!(is_done(&event), done); } + #[rstest] + #[tokio::test] + async fn a_transport_failure_tunnelled_through_the_body_comes_back_as_itself() { + use futures_util::StreamExt; + + let failure = litellm_http::Error::upstream(502, "gateway went away".into()); + let bytes: ByteStream = futures_util::stream::iter([ + Ok(Bytes::from_static(b"data: {}\n\n")), + Err(std::io::Error::other(failure.clone())), + ]) + .boxed(); + + let mut events = event_stream::(bytes, is_done); + assert!(events.next().await.expect("first event").is_ok()); + let error = events.next().await.expect("failure").unwrap_err(); + assert_eq!(error, Error::Http(failure)); + } + + #[rstest] + #[tokio::test] + async fn a_body_failure_that_is_not_transport_is_an_invalid_response() { + use futures_util::StreamExt; + + let bytes: ByteStream = + futures_util::stream::iter([Err(std::io::Error::other("disk full"))]).boxed(); + + let error = event_stream::(bytes, is_done) + .next() + .await + .expect("failure") + .unwrap_err(); + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + assert!( + litellm_error::Report::new(&error) + .to_string() + .contains("disk full") + ); + } + #[rstest] fn the_done_frame_we_send_decodes_as_a_control_frame() { let events: Vec = decode(DONE_FRAME, is_done).unwrap(); diff --git a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs index 61dd264336e..0f8703f223b 100644 --- a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs @@ -14,7 +14,7 @@ use serde_json::{Map, Value, json}; use strum::IntoStaticStr; use crate::{ - Error, + Error, ErrorDetail, base_llm::{ audio_transcription::transformation::{ AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, Headers, @@ -59,20 +59,12 @@ fn audio_fields(audio: Value) -> Result { .and_then(Value::as_str) .filter(|value| !value.is_empty()) .ok_or(Error::MissingField("audio.data"))?; - let format = object.get("format").cloned().ok_or_else(|| { - Error::InvalidRequest( - "audio.format must be wav, mp3, flac, or ogg" - .to_string() - .into(), - ) - })?; - let format: AudioFormat = serde_json::from_value(format).map_err(|_| { - Error::InvalidRequest( - "audio.format must be wav, mp3, flac, or ogg" - .to_string() - .into(), - ) - })?; + let format = object + .get("format") + .cloned() + .ok_or(Error::MissingField("audio.format"))?; + let format: AudioFormat = serde_json::from_value(format) + .map_err(|error| Error::InvalidRequest(ErrorDetail::invalid("audio.format", error)))?; Ok(AudioInput { data: data.to_string(), format, @@ -135,14 +127,13 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { ) -> Result { let response: ConverseResponse = serde_json::from_value(response_json).map_err(|error| { - Error::InvalidResponse(format!("invalid Converse response: {error}").into()) + Error::InvalidResponse(ErrorDetail::invalid("Converse response", error)) })?; if response.message_content_is_non_text() { - return Err(Error::InvalidResponse( - "Bedrock response contains non-text transcript content" - .to_string() - .into(), - )); + return Err(Error::InvalidResponse(ErrorDetail::Expected { + subject: "Bedrock transcript content", + expected: "text", + })); } Ok(AudioTranscriptionResponseData { text: response.content_text(), diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index caa345a7024..976c9161dd2 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -5,6 +5,7 @@ use litellm_auth_aws::{ constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, resolve_bedrock_region, }; +use litellm_core_utils::decline::Decline; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, prompt_templates::factory::{Conversation, TurnRole, build_conversation}, @@ -19,14 +20,14 @@ use litellm_llms_types::{ use serde_json::{Map, Value, json}; use crate::{ - Error, + Error, ErrorDetail, base_llm::{ auth::AuthScheme, chat::{ streaming::StreamShape, transformation::{ BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData, - Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param, + ValidatedEnvironment, unsupported_message, unsupported_param, }, }, }, @@ -154,25 +155,20 @@ impl BaseConfig for AmazonConverseConfig { response: ProviderChatResponseData, ) -> Result { let body = response.body; - if !body.is_object() { - return Err(Error::InvalidResponse( - "converse response is not an object".into(), - )); - } - for field in ["output", "usage"] { - if body.get(field).is_none() { - return Err(Error::InvalidResponse( - format!("invalid Converse response: missing field `{field}`").into(), - )); - } - } + let has_usage = body.get("usage").is_some(); let response: ConverseResponse = serde_json::from_value(body).map_err(|error| { - Error::InvalidResponse(format!("invalid Converse response: {error}").into()) + Error::InvalidResponse(ErrorDetail::invalid("Converse response", error)) })?; - // The route declines tool requests, so anything other than a text block - // is something this path never asked for. Decline; the host falls back. + if !has_usage { + return Err(Error::InvalidResponse(ErrorDetail::MissingField { + subject: "Converse response", + field: "usage", + })); + } if response.message_content_is_non_text() { - return Err(Error::Unsupported("non-text response content block")); + return Err(Error::InvalidResponse(ErrorDetail::Unhandled( + "non-text response content block", + ))); } let text = response.content_text(); let computed = usage_from_parts( @@ -269,7 +265,7 @@ impl BaseConfig for AmazonConverseConfig { &self, messages: &[ChatMessage], optional_params: &Map, - ) -> Option { + ) -> Option { unsupported_param( self.supported_openai_param_mappings(), CONFIG_PARAMS, @@ -283,7 +279,7 @@ impl BaseConfig for AmazonConverseConfig { messages .iter() .any(has_blank_text) - .then_some(Unsupported("blank message text")) + .then_some(Decline::new("blank message text")) }) // Converse has no assistant prefill: Python inserts a continue turn // when a conversation opens or closes on an assistant message, and @@ -296,7 +292,7 @@ impl BaseConfig for AmazonConverseConfig { .turns .last() .is_some_and(|turn| turn.role == TurnRole::Assistant); - (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( + (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Decline::new( "conversation does not run user turn to user turn", )) }) @@ -358,11 +354,12 @@ mod tests { Error, base_llm::{ auth::AuthScheme, - chat::transformation::{BaseConfig, ProviderChatResponseData, Unsupported}, + chat::transformation::{BaseConfig, ProviderChatResponseData}, }, bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, }; use litellm_auth::CredentialPlacement; + use litellm_core_utils::decline::Decline; use litellm_llms_types::formats::chat_completions::{ChatCompletionsResponse, ChatMessage}; use rstest::rstest; use serde_json::{Map, Value, json}; @@ -396,7 +393,7 @@ mod tests { ) } - fn reason(msgs: Value, opts: Value) -> Option { + fn reason(msgs: Value, opts: Value) -> Option { BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) } @@ -472,7 +469,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"stream": true}) ), - Some(Unsupported("streaming")) + Some(Decline::new("streaming")) ); } @@ -483,7 +480,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"topK": 40}) ), - Some(Unsupported("unrecognized request parameter")) + Some(Decline::new("unrecognized request parameter")) ); } @@ -497,7 +494,7 @@ mod tests { fn declines_tools_and_other_params_outside_the_allowlist(#[case] param: Value) { assert_eq!( reason(json!([{"role": "user", "content": "hi"}]), param.clone()), - Some(Unsupported("unrecognized request parameter")), + Some(Decline::new("unrecognized request parameter")), "expected {param} to decline" ); } @@ -514,7 +511,7 @@ mod tests { json!([{"role": "user", "content": content}, {"role": "user", "content": "hi"}]), json!({}) ), - Some(Unsupported("blank message text")), + Some(Decline::new("blank message text")), "expected blank content {content} to decline" ); } @@ -526,7 +523,7 @@ mod tests { // host before the call rather than an error after it. assert_eq!( reason(json!([{"role": "user", "content": []}]), json!({})), - Some(Unsupported("message without content")) + Some(Decline::new("message without content")) ); assert_eq!( reason( @@ -547,7 +544,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported( + Some(Decline::new( "conversation does not run user turn to user turn" )) ); @@ -559,7 +556,7 @@ mod tests { ]), json!({}) ), - Some(Unsupported( + Some(Decline::new( "conversation does not run user turn to user turn" )) ); @@ -819,7 +816,7 @@ mod tests { ]}]), json!({}) ), - Some(Unsupported("non-text message content")) + Some(Decline::new("non-text message content")) ); } @@ -846,7 +843,7 @@ mod tests { } #[test] - fn declines_a_response_carrying_a_tool_use_block() { + fn a_response_carrying_a_tool_use_block_is_invalid() { let err = transform_response(json!({ "output": {"message": {"content": [ {"toolUse": {"toolUseId": "t1", "name": "f", "input": {}}} @@ -855,37 +852,32 @@ mod tests { "usage": {"inputTokens": 1, "outputTokens": 1} })) .expect_err("tool use block"); - assert_eq!(err, Error::Unsupported("non-text response content block")); - } - - #[test] - fn errors_on_a_response_missing_required_fields() { assert_eq!( - transform_response(json!("nope")).expect_err("not an object"), - Error::InvalidResponse("converse response is not an object".to_string().into()) - ); - assert_eq!( - transform_response(json!({"usage": {}})).expect_err("no output"), - Error::InvalidResponse("invalid Converse response: missing field `output`".into()) - ); - assert_eq!( - transform_response(json!({"output": {"message": {"content": []}}})) - .expect_err("no usage"), - Error::InvalidResponse("invalid Converse response: missing field `usage`".into()) + err, + Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "non-text response content block" + )) ); } - #[test] - fn rejects_malformed_text_and_token_counts() { - for response in [ - json!({"output": {"message": {"content": [{"text": 123}]}}, "usage": {"inputTokens": 1, "outputTokens": 1}}), - json!({"output": {"message": {"content": [{"text": "x"}]}}, "usage": {"inputTokens": "1", "outputTokens": 1}}), - ] { - assert!(matches!( - transform_response(response), - Err(Error::InvalidResponse(message)) if message.to_string().starts_with("invalid Converse response:") - )); - } + #[rstest::rstest] + #[case::not_an_object(json!("nope"), "invalid type")] + #[case::no_output(json!({"usage": {"inputTokens": 1, "outputTokens": 1}}), "missing field `output`")] + #[case::no_usage(json!({"output": {"message": {"content": []}}}), "missing usage")] + #[case::non_string_text( + json!({"output": {"message": {"content": [{"text": 123}]}}, "usage": {"inputTokens": 1, "outputTokens": 1}}), + "invalid type" + )] + #[case::non_integer_tokens( + json!({"output": {"message": {"content": [{"text": "x"}]}}, "usage": {"inputTokens": "1", "outputTokens": 1}}), + "invalid type" + )] + fn rejects_a_malformed_response(#[case] response: Value, #[case] fact: &str) { + let error = transform_response(response).expect_err("malformed response"); + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains("Converse response"), "{report}"); + assert!(report.contains(fact), "{report}"); } #[test] diff --git a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs index 878c424ae2d..15a6ff38822 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs @@ -2,6 +2,7 @@ use base64::Engine; use bytes::Buf; use futures_util::{Stream, StreamExt}; use litellm_framer::{ + EventStreamError, aws_event_stream::{AwsEventStreamCodec, Message}, frames, }; @@ -11,9 +12,8 @@ use serde_json::Value; use crate::{ Error, - anthropic::chat::handler::ModelResponseIterator, base_llm::{ - chat::streaming::{ChatStream, StreamShape}, + base_model_iterator::tunnelled_http, messages::streaming::{ByteStream, EventStream}, }, }; @@ -47,11 +47,21 @@ where B: Buf + Send, E: std::error::Error + Send + Sync + 'static, { - frames(input, AwsEventStreamCodec).map(|message| { - decode_invoke_chunk(message.map_err(|error| { + frames(input, AwsEventStreamCodec).map(|message| decode_invoke_chunk(message.map_err(framing)?)) +} + +fn framing(error: EventStreamError) -> Error { + match &error { + EventStreamError::Body(body) if let Some(http) = tunnelled_http(body) => { + Error::Http(http.clone()) + } + EventStreamError::Body(_) + | EventStreamError::InvalidLength(_) + | EventStreamError::Truncated + | EventStreamError::Malformed(_) => { Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error)) - })?) - }) + } + } } pub fn decode_invoke_anthropic_chunk(chunk: Value) -> Result { @@ -64,34 +74,6 @@ pub fn invoke_anthropic_event_stream(bytes: ByteStream) -> EventStream { Box::pin(invoke_chunk_stream(bytes).map(|chunk| decode_invoke_anthropic_chunk(chunk?))) } -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString)] -enum InvokeProvider { - #[strum(serialize = "anthropic")] - Anthropic, - #[strum(serialize = "deepseek_r1")] - DeepseekR1, - #[strum(serialize = "moonshot")] - Moonshot, - #[strum(disabled)] - Unsupported, -} - -pub fn invoke_chat_stream(invoke_provider: &str, shape: StreamShape) -> Result { - match invoke_provider - .parse() - .unwrap_or(InvokeProvider::Unsupported) - { - InvokeProvider::Anthropic => Ok(ChatStream::new( - invoke_anthropic_event_stream, - ModelResponseIterator::new(shape), - )), - InvokeProvider::DeepseekR1 | InvokeProvider::Moonshot => Err(Error::Unsupported( - "Bedrock invoke streaming for this model family", - )), - InvokeProvider::Unsupported => Err(Error::Unsupported("Bedrock invoke streaming")), - } -} - #[cfg(test)] mod tests { use aws_smithy_eventstream::frame::write_message_to; @@ -112,22 +94,6 @@ mod tests { const TEXT_DELTA: &str = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#; - #[rstest::rstest] - #[case::anthropic("anthropic", InvokeProvider::Anthropic)] - #[case::deepseek_r1("deepseek_r1", InvokeProvider::DeepseekR1)] - #[case::moonshot("moonshot", InvokeProvider::Moonshot)] - #[case::unknown("qwen", InvokeProvider::Unsupported)] - #[case::case_sensitive("Anthropic", InvokeProvider::Unsupported)] - fn invoke_provider_maps_each_model_family_name( - #[case] name: &str, - #[case] expected: InvokeProvider, - ) { - assert_eq!( - name.parse().unwrap_or(InvokeProvider::Unsupported), - expected - ); - } - fn aws_wire(chunk: &str) -> Vec { let payload = serde_json::json!({"bytes": STANDARD.encode(chunk)}); let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header( diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 72958fd6cb9..79c30df9b19 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -1,6 +1,6 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, ParamSpec, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec, SecretValue}; use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region}; use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{header_value, with_header, without_headers}; @@ -36,6 +36,7 @@ const WORKSPACE_FIELDS: &[&str] = &[ const WORKSPACE_HEADER: &str = "anthropic-workspace-id"; static WORKSPACE: ParamSpec = ParamSpec { setting: "workspace_id", + kind: ParamKind::Setting, wire: WORKSPACE_FIELDS, module_global: None, env: WORKSPACE_ENVS, diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 1f6b6edf28d..3ec60f12021 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use std::convert::Infallible; use crate::{ @@ -130,9 +131,9 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig { _request: MessagesRequest, _context: &MessagesTransformContext, ) -> Result { - Err(Error::Unsupported( + Err(Error::Declined(Decline::new( "Bedrock invoke messages request shaping", - )) + ))) } fn secret_names(&self) -> &'static [&'static str] { diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index dadab413fe4..76ff27b56fc 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -195,11 +195,10 @@ impl CohereParseConfig { .and_then(env_lookup) .filter(|key| !key.trim().is_empty()) }) - .ok_or_else(|| { - Error::Auth(litellm_auth::Error::ProviderAuthentication( - "Missing COHERE_API_KEY - set it in the environment or pass api_key".into(), - )) - })?; + .ok_or(Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "Cohere", + environment_variable: COHERE_API_KEY_ENV, + }))?; Ok( std::iter::once(("Authorization".into(), format!("Bearer {key}"))) .chain(connection.extra_headers.clone()) @@ -208,27 +207,28 @@ impl CohereParseConfig { } fn build_ocr_url(&self, api_base: &str) -> Result { - let parsed = reqwest::Url::parse(api_base).map_err(|_| invalid_api_base())?; + let parsed = reqwest::Url::parse(api_base) + .map_err(|error| Error::request_field("api_base", error))?; if !matches!(parsed.scheme(), "http" | "https") { return Err(invalid_api_base()); } ApiUrl::parse(api_base) .and_then(|url| url.complete_path(&["v2", "parse"])) .map(|url| url.into_string()) - .map_err(|_| invalid_api_base()) + .map_err(|error| Error::request_field("api_base", error)) } } pub fn validate_document(document: &OcrDocument) -> Result<(), Error> { let OcrDocument::ImageUrl { image_url, .. } = document else { - return Err(Error::CohereImageOnly); + return Err(image_only()); }; if image_url.is_empty() { - return Err(Error::CohereImageOnly); + return Err(image_only()); } if let Some(inline) = InlineDocument::parse(image_url)? { if !inline.mime_type().type_.eq_ignore_ascii_case("image") { - return Err(Error::CohereImageOnly); + return Err(image_only()); } inline.decode(OCR_INLINE_MAX_BYTES)?; } @@ -260,7 +260,7 @@ pub fn normalize_response( fn image_url(document: OcrDocument) -> Result { validate_document(&document)?; let OcrDocument::ImageUrl { image_url, .. } = document else { - return Err(Error::CohereImageOnly); + return Err(image_only()); }; Ok(image_url) } @@ -329,9 +329,17 @@ fn billed_pages(response: &CohereResponse) -> Option { response.meta.as_ref()?.billed_units.as_ref()?.pages } +fn image_only() -> Error { + Error::UnsupportedDocument { + provider: "Cohere Parse", + guidance: "it only accepts `image_url` documents, not document_url or PDF inputs", + } +} + fn invalid_api_base() -> Error { Error::RequestField { path: "api_base".into(), + cause: None, } } @@ -368,11 +376,14 @@ mod tests { #[test] fn options_reject_invalid_output_format() { let invalid = serde_json::from_value(json!({"output_format":"html"})).unwrap(); - assert!(matches!( - CohereParseConfig.map_ocr_params(&invalid, "parse"), - Err(Error::RequestField { path }) - if path == "optional_params.output_format" - )); + let Err(error) = CohereParseConfig.map_ocr_params(&invalid, "parse") else { + panic!("html is not a Cohere output format"); + }; + assert!(error.is_request()); + let argument = + litellm_error::find::(&error) + .expect("the argument failure is kept"); + assert_eq!(argument.path, "output_format"); } #[test] @@ -432,7 +443,7 @@ mod tests { .unwrap(); assert!(matches!( normalize_response("parse", response).unwrap_err(), - Error::ResponseField { path } + Error::ResponseField { path, .. } if path == "pages[0].markdown.images[0].image_base64" )); } @@ -657,7 +668,10 @@ mod tests { fn request_requires_image(#[case] value: Value) { assert!(matches!( validate_document(&serde_json::from_value(value).unwrap()), - Err(Error::CohereImageOnly) + Err(Error::UnsupportedDocument { + provider: "Cohere Parse", + .. + }) )); } diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index 10a4f696406..b3d1421083a 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -1,3 +1,6 @@ +use litellm_core_utils::decline::Decline; +use litellm_error::Source; + #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum Error { #[error("expected {expected}, got {actual}")] @@ -7,31 +10,126 @@ pub enum Error { }, #[error("missing required field: {0}")] MissingField(&'static str), - #[error("invalid request: {0}")] + #[error("invalid request")] InvalidRequest(#[source] ErrorDetail), - #[error("invalid response: {0}")] + #[error("invalid response")] InvalidResponse(#[source] ErrorDetail), - #[error("unsupported: {0}")] - Unsupported(&'static str), + #[error(transparent)] + Declined(Decline), #[error(transparent)] Auth(#[from] litellm_auth::Error), + #[error(transparent)] + Http(#[from] litellm_http::Error), +} + +impl Error { + pub fn is_request(&self) -> bool { + match self { + Self::InvalidType { .. } | Self::MissingField(_) | Self::InvalidRequest(_) => true, + Self::Auth(error) => error.is_request(), + Self::Http(error) => error.is_request(), + Self::InvalidResponse(_) | Self::Declined(_) => false, + } + } + + pub fn is_response(&self) -> bool { + match self { + Self::InvalidResponse(_) => true, + Self::Http(error) => error.is_response(), + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::Declined(_) + | Self::Auth(_) => false, + } + } + + pub fn is_missing_credential(&self) -> bool { + match self { + Self::Auth(error) => error.is_missing_credential(), + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::InvalidResponse(_) + | Self::Declined(_) + | Self::Http(_) => false, + } + } + + /// A transformation runs on both sides of the send, so only the transport and credential + /// resolution, which always runs before it, can tell whether the request reached the provider. + pub fn dispatched(&self) -> bool { + match self { + Self::Http(error) => error.dispatched(), + Self::Auth(_) => false, + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::InvalidResponse(_) => true, + Self::Declined(_) => false, + } + } + + /// Why the transformation rejected the caller's request, rendered as Python's own + /// validation renders it, without this layer's "invalid request" prefix. + pub fn rejected_request(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::InvalidRequest(detail) => Some(detail), + Self::InvalidType { .. } | Self::MissingField(_) => Some(self), + Self::InvalidResponse(_) | Self::Declined(_) | Self::Auth(_) | Self::Http(_) => None, + } + } + + pub fn declined(&self) -> Option<&Decline> { + match self { + Self::Declined(decline) => Some(decline), + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::InvalidResponse(_) + | Self::Auth(_) + | Self::Http(_) => None, + } + } + + pub fn upstream_response(&self) -> Option<(u16, &str)> { + match self { + Self::Http(error) => error.upstream_response(), + Self::InvalidType { .. } + | Self::MissingField(_) + | Self::InvalidRequest(_) + | Self::InvalidResponse(_) + | Self::Declined(_) + | Self::Auth(_) => None, + } + } } #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ErrorDetail { - #[error("{0}")] - Message(String), - #[error("invalid {subject}: {source}")] + #[error("invalid {subject}")] Invalid { subject: &'static str, #[source] - source: ErrorSource, + cause: Source, }, - #[error("{operation} failed: {source}")] + #[error("{operation} failed")] Failed { operation: &'static str, #[source] - source: ErrorSource, + cause: Source, + }, + #[error("{route} requires at least one message")] + NoMessages { route: &'static str }, + #[error("{subject} must be {expected}")] + Expected { + subject: &'static str, + expected: &'static str, + }, + #[error("{subject} is missing {field}")] + MissingField { + subject: &'static str, + field: &'static str, }, #[error("{field} must be {expected}, got {actual}")] InvalidValue { @@ -60,12 +158,20 @@ pub enum ErrorDetail { value: String, hint: String, }, - #[error("invalid {subject} on line {line}: {source}")] + #[error( + "Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({{LLM_PROVIDER}}_API_KEY) or via params" + )] + MissingApiBase, + #[error("the host call raised an exception")] + HostCall, + #[error("{0} is not handled by the rust path")] + Unhandled(&'static str), + #[error("invalid {subject} on line {line}")] InvalidLine { subject: &'static str, line: usize, #[source] - source: ErrorSource, + cause: Source, }, #[error("{operation} failed: {detail}")] RemoteFailure { @@ -77,64 +183,33 @@ pub enum ErrorDetail { impl ErrorDetail { pub fn invalid( subject: &'static str, - source: impl std::error::Error + Send + Sync + 'static, + cause: impl std::error::Error + Send + Sync + 'static, ) -> Self { Self::Invalid { subject, - source: ErrorSource::new(source), + cause: Source::new(cause), } } pub fn failed( operation: &'static str, - source: impl std::error::Error + Send + Sync + 'static, + cause: impl std::error::Error + Send + Sync + 'static, ) -> Self { Self::Failed { operation, - source: ErrorSource::new(source), + cause: Source::new(cause), + } + } + + pub fn invalid_line( + subject: &'static str, + line: usize, + cause: impl std::error::Error + Send + Sync + 'static, + ) -> Self { + Self::InvalidLine { + subject, + line, + cause: Source::new(cause), } } } - -impl From for ErrorDetail { - fn from(message: String) -> Self { - Self::Message(message) - } -} - -impl From<&str> for ErrorDetail { - fn from(message: &str) -> Self { - Self::Message(message.into()) - } -} - -#[derive(Clone, Debug)] -pub struct ErrorSource(std::sync::Arc); - -impl std::ops::Deref for ErrorSource { - type Target = dyn std::error::Error + Send + Sync; - - fn deref(&self) -> &Self::Target { - self.0.as_ref() - } -} - -impl std::fmt::Display for ErrorSource { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(&self.0, formatter) - } -} - -impl ErrorSource { - pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self { - Self(std::sync::Arc::new(error)) - } -} - -impl PartialEq for ErrorSource { - fn eq(&self, other: &Self) -> bool { - std::sync::Arc::ptr_eq(&self.0, &other.0) - } -} - -impl Eq for ErrorSource {} diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index eb232404d93..68e6e4010c2 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -17,4 +17,4 @@ pub mod reducto; pub mod tencent; pub mod vertex_ai; -pub use error::{Error, ErrorDetail, ErrorSource}; +pub use error::{Error, ErrorDetail}; diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index 29d4d5c1610..fa4e2254e86 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -171,9 +171,7 @@ impl MistralOcrConfig { ApiUrl::parse(base) .and_then(|url| url.complete_path(&["v1", "ocr"])) .map(|url| url.into_string()) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } } @@ -186,6 +184,7 @@ pub fn normalize_response( Some(None) => { return Err(Error::ResponseField { path: "model".into(), + cause: None, }); } None => model.to_string(), @@ -231,7 +230,7 @@ mod tests { let response = serde_json::from_value(json!({"model":null})).unwrap(); assert!(matches!( normalize_response("fallback", response).unwrap_err(), - Error::ResponseField { path } if path == "model" + Error::ResponseField { path, .. } if path == "model" )); } @@ -263,7 +262,7 @@ mod tests { .unwrap_err(); assert!(matches!( error, - Error::ResponseField { path: actual } if actual == path + Error::ResponseField { path: actual, .. } if actual == path )); } diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index 99cc3559514..cf74871894b 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -6,7 +6,7 @@ use serde_json::{Map, Value}; use litellm_auth::{CredentialPlacement, SecretValue}; use crate::{ - Error, + Error, ErrorDetail, base_llm::{ auth::{AuthScheme, ValidatedEnvironment}, responses::transformation::{ @@ -180,15 +180,19 @@ pub fn responses_request( params: Map, ) -> Result { if !input.is_string() && !input.is_array() { - return Err(Error::InvalidRequest( - "responses input must be a string or an array".into(), - )); + return Err(Error::InvalidRequest(ErrorDetail::Expected { + subject: "responses input", + expected: "a string or an array", + })); } if params .get("stream") .is_some_and(|stream| !stream.is_boolean()) { - return Err(Error::InvalidRequest("stream must be a boolean".into())); + return Err(Error::InvalidRequest(ErrorDetail::Expected { + subject: "stream", + expected: "a boolean", + })); } Ok(Value::Object( params @@ -202,7 +206,9 @@ pub fn responses_request( } pub fn responses_response(body: Value) -> Result { - serde_json::from_value(body).map_err(|error| Error::InvalidResponse(error.to_string().into())) + serde_json::from_value(body).map_err(|error| { + Error::InvalidResponse(ErrorDetail::invalid("Responses API response", error)) + }) } #[cfg(test)] diff --git a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs index 0a89adf90f7..e998f8785cb 100644 --- a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs @@ -123,9 +123,12 @@ impl BaseConfig for OpenAILikeChatConfig { ) -> Result { let mut body = response.body; sanitize_usage(&mut body); - let body = body - .as_object() - .ok_or_else(|| Error::InvalidResponse("chat response is not an object".into()))?; + let body = + body.as_object() + .ok_or(Error::InvalidResponse(crate::ErrorDetail::Expected { + subject: "chat response", + expected: "an object", + }))?; let choices = body .get("choices") @@ -243,14 +246,20 @@ fn normalize_choice(position: usize, choice: &Value) -> Result Option { + fn reason(msgs: Value, opts: Value) -> Option { OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) } @@ -414,9 +424,13 @@ mod tests { #[rstest] fn a_missing_api_base_is_an_error() { assert!(matches!( - OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG - .get_complete_url(None, "my-model", ¶ms(json!({})), &no_env), - Err(Error::InvalidRequest(message)) if message.to_string().starts_with("Missing API Base") + OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG.get_complete_url( + None, + "my-model", + ¶ms(json!({})), + &no_env + ), + Err(Error::InvalidRequest(crate::ErrorDetail::MissingApiBase)) )); } @@ -528,9 +542,7 @@ mod tests { } #[rstest] - fn a_tool_call_response_declines_instead_of_dropping_the_calls() { - // The `json_mode` rewrite needs a request flag the route does not carry, so - // a tool-call answer falls back to Python rather than losing the calls. + fn a_tool_call_response_fails_instead_of_dropping_the_calls() { assert_eq!( transform_response(json!({ "model": "m", @@ -547,12 +559,14 @@ mod tests { "finish_reason": "tool_calls", }], })), - Err(Error::Unsupported("tool call response")) + Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "tool call response" + ))) ); } #[rstest] - fn a_refusal_declines_instead_of_returning_an_empty_reply() { + fn a_refusal_fails_instead_of_returning_an_empty_reply() { assert_eq!( transform_response(json!({ "model": "m", @@ -561,12 +575,14 @@ mod tests { "finish_reason": "stop", }], })), - Err(Error::Unsupported("refusal response")) + Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "refusal response" + ))) ); } #[rstest] - fn a_non_text_response_content_declines() { + fn a_non_text_response_content_is_invalid() { assert_eq!( transform_response(json!({ "model": "m", @@ -575,7 +591,9 @@ mod tests { "finish_reason": "stop", }], })), - Err(Error::Unsupported("non-text response content")) + Err(Error::InvalidResponse(crate::ErrorDetail::Unhandled( + "non-text response content" + ))) ); } @@ -585,7 +603,7 @@ mod tests { fn declines(#[case] opts: Value, #[case] expected: &'static str) { assert_eq!( reason(json!([{"role": "user", "content": "hi"}]), opts), - Some(Unsupported(expected)) + Some(Decline::new(expected)) ); } @@ -615,7 +633,7 @@ mod tests { json!([{"role": "user", "content": "hi"}]), json!({"tools": [{"type": "function", "function": {"name": "f"}}]}), ), - Some(Unsupported("unrecognized request parameter")) + Some(Decline::new("unrecognized request parameter")) ); } } diff --git a/litellm-rust/crates/llms/src/openai_like/common_utils.rs b/litellm-rust/crates/llms/src/openai_like/common_utils.rs index b0f5adafe00..d798f851d79 100644 --- a/litellm-rust/crates/llms/src/openai_like/common_utils.rs +++ b/litellm-rust/crates/llms/src/openai_like/common_utils.rs @@ -1,7 +1,7 @@ //! Shared OpenAI-like credential and endpoint resolution, mirroring //! `litellm/llms/openai_like/common_utils.py`. -use crate::Error; +use crate::{Error, ErrorDetail}; /// `OpenAILikeChatConfig._get_openai_compatible_provider_info`: the deployment's /// `api_base` wins over `OPENAI_LIKE_API_BASE`, and the deployment key over @@ -31,12 +31,7 @@ pub fn complete_openai_like_url( env_lookup: &dyn Fn(&str) -> Option, ) -> Result { let (api_base, _) = openai_compatible_provider_info(api_base, None, env_lookup); - let api_base = api_base.ok_or_else(|| { - Error::InvalidRequest( - "Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({LLM_PROVIDER}_API_KEY) or via params" - .into(), - ) - })?; + let api_base = api_base.ok_or(Error::InvalidRequest(ErrorDetail::MissingApiBase))?; if custom_endpoint { return Ok(api_base); } diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 1979b442936..bf0ddc1c7a8 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -1,11 +1,5 @@ use std::collections::BTreeMap; -use litellm_core_utils::{ - call_arguments::{CallArguments, compose_body}, - params::OpaqueParams, - url_utils::ApiUrl, -}; -use litellm_http::outbound::OutboundRequest; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::{Map, Value, json}; @@ -18,6 +12,11 @@ use crate::base_llm::ocr::{ decode_and_normalize_response, }, }; +use litellm_core_utils::{ + call_arguments::{CallArguments, compose_body}, + params::OpaqueParams, + url_utils::ApiUrl, +}; use litellm_llms_types::formats::ocr::{ LiteLLMOcrResponse, OcrDocument, OcrPage, OcrResponseFormat, OcrUsageInfo, }; @@ -173,7 +172,7 @@ impl BaseOcrConfig for ReductoParseV3Config { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -262,7 +261,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -275,7 +274,7 @@ async fn prepare_upload_request, -) -> Result { +) -> Result { let params = config.map_ocr_params(&request.optional_params, &request.model)?; let headers = config.validate_environment(request, client).await?; let url = config.get_complete_url(request, ¶ms, &headers)?; @@ -300,9 +299,16 @@ async fn prepare_upload_request Error { + Error::UnsupportedDocument { + provider: "Reducto", + guidance: "pass a reducto:// id or a data URI, or upload the file first for plain HTTP URLs", + } +} + fn uploaded_file_id(document: OcrDocument) -> Result { if !document.source().starts_with(REDUCTO_ID_PREFIX) { - return Err(Error::ReductoSource); + return Err(unsupported_source()); } if document.source()[REDUCTO_ID_PREFIX.len()..] .trim() @@ -310,6 +316,7 @@ fn uploaded_file_id(document: OcrDocument) -> Result { { return Err(Error::RequestField { path: "document file id".into(), + cause: None, }); } Ok(ReductoFileId(document.source().into())) @@ -393,6 +400,7 @@ fn build_pages_from_reducto(chunks: Vec) -> Result, E Some(Value::String(content)) => Ok(Some(content.as_str())), Some(_) => Err(Error::ResponseField { path: "result.chunks.blocks.content".into(), + cause: None, }), }) .collect::, _>>()?; @@ -437,9 +445,7 @@ fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result( response, false, @@ -547,6 +555,7 @@ async fn upload_bytes_async( let Some(file_id) = file_id else { return Err(Error::ResponseField { path: "file_id".into(), + cause: None, }); }; Ok(ReductoFileId(file_id.to_string())) @@ -554,6 +563,8 @@ async fn upload_bytes_async( #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; #[tokio::test] @@ -566,10 +577,7 @@ mod tests { let params = ReductoParseV3Config .map_ocr_params(&overrides, "parse-v3") .unwrap(); - let client = OcrClient::for_test( - litellm_http::Client::plain_for_test(), - litellm_http::Client::no_redirect_for_test(), - ); + let client = OcrClient::for_test(); let connection = OcrConnection::default(); let document = serde_json::from_value( json!({"type":"document_url","document_url":"reducto://ready.pdf"}), @@ -629,15 +637,16 @@ mod tests { ); } + #[rstest] + #[case::fractional_pages(json!({"num_pages":1.5}))] + #[case::non_numeric_pages(json!({"num_pages":[]}))] + #[case::non_numeric_credits(json!({"credits":{}}))] + fn invalid_usage_is_rejected(#[case] usage: Value) { + assert!(serde_json::from_value::(json!({"usage":usage})).is_err()); + } + #[test] - fn usage_uses_shared_validation_while_block_page_numbers_are_best_effort() { - for usage in [ - json!({"num_pages":1.5}), - json!({"num_pages":[]}), - json!({"credits":{}}), - ] { - assert!(serde_json::from_value::(json!({"usage":usage})).is_err()); - } + fn block_page_numbers_are_best_effort() { let response = serde_json::from_value(json!({"result":{"chunks":[{"blocks":[ {"content":"ignored", "bbox":{"page":"invalid"}}, {"content":"kept", "bbox":{"page":2.5}, "extra":null} diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs index 810462c57f8..d760da9073e 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs @@ -16,7 +16,9 @@ pub(super) fn vertex_config(request: &PreparedOcrRequest) -> Result Result<(), Error> { if connection.api_base.is_some() && connection.api_base_source == InputSource::Request { return Err(litellm_auth::Error::InvalidConfiguration( - "credentials cannot be sent to a request-controlled Vertex AI endpoint".into(), + litellm_auth::ErrorDetail::RequestControlledEndpoint { + provider: "Vertex AI", + }, ) .into()); } diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs index 67f9122ed67..5d680588612 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs @@ -235,8 +235,9 @@ pub fn normalize_response( &mut output, PythonJsonFormatter, )) - .map_err(|_| response_field("content"))?; - String::from_utf8(output).map_err(|_| response_field("content"))? + .map_err(|error| Error::response_field(content_path("content"), error))?; + String::from_utf8(output) + .map_err(|error| Error::response_field(content_path("content"), error))? }; (data, fallback) } @@ -261,7 +262,12 @@ pub fn normalize_response( }) }) .collect::, Error>>()?, - Some(_) => return Err(response_field("pages")), + Some(_) => { + return Err(Error::ResponseField { + path: content_path("pages"), + cause: None, + }); + } None => Vec::new(), }; let usage = ocr_data @@ -273,7 +279,12 @@ pub fn normalize_response( .transpose()?; let model = match ocr_data.get("model") { Some(Value::String(model)) => model.clone(), - Some(_) => return Err(response_field("model")), + Some(_) => { + return Err(Error::ResponseField { + path: content_path("model"), + cause: None, + }); + } None => model.to_string(), }; Ok(LiteLLMOcrResponse { @@ -369,10 +380,8 @@ impl serde_json::ser::Formatter for PythonJsonFormatter { } } -fn response_field(field: &str) -> Error { - Error::ResponseField { - path: format!("choices[0].message.content.{field}"), - } +fn content_path(field: &str) -> String { + format!("choices[0].message.content.{field}") } pub fn provider_model(model: &str) -> Result { @@ -380,6 +389,7 @@ pub fn provider_model(model: &str) -> Result { if local_model.is_empty() { return Err(Error::RequestField { path: "model".into(), + cause: None, }); } Ok(format!("{MODEL_PREFIX}{local_model}")) @@ -411,9 +421,7 @@ impl VertexAIDeepSeekOCRConfig { ]) }) .map(|url| url.into_string()) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } } @@ -625,20 +633,18 @@ mod tests { assert_eq!(result["future"], "kept"); } - #[test] - fn response_transform_rejects_missing_empty_and_malformed_content() { - for value in [ - json!({"choices":[]}), - json!({"choices":[{"message":{"content":{}}}]}), - json!({"choices":[{"message":{"content":""}}]}), - json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), - json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), - ] { - let result = serde_json::from_value::(value) - .map_err(|_| ()) - .and_then(|response| normalize_response("model", response).map_err(|_| ())); - assert!(result.is_err()); - } + #[rstest::rstest] + #[case::no_choices(json!({"choices":[]}))] + #[case::empty_object(json!({"choices":[{"message":{"content":{}}}]}))] + #[case::empty_string(json!({"choices":[{"message":{"content":""}}]}))] + #[case::malformed_string(json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}))] + #[case::malformed_object(json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}))] + fn response_transform_rejects_missing_empty_and_malformed_content(#[case] value: Value) { + let rejected = match serde_json::from_value::(value) { + Ok(response) => normalize_response("model", response).is_err(), + Err(_) => true, + }; + assert!(rejected); } #[test] diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index ab28ee8bff4..77881cd2e78 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -174,9 +174,7 @@ impl VertexAiOcrConfig { ]) }) .map(|url| url.into_string()) - .map_err(|_| Error::RequestField { - path: "api_base".into(), - }) + .map_err(|error| Error::request_field("api_base", error)) } } @@ -198,6 +196,7 @@ fn validate_location(location: &str) -> Result<(), Error> { } Err(Error::RequestField { path: "vertex_location".into(), + cause: None, }) } diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index 570e68ca4fd..2c2866d7090 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -20,4 +20,5 @@ time.workspace = true [dev-dependencies] jsonschema = { version = "0.55.1", default-features = false } +litellm-error.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index 0b113436a5b..37412ef1e14 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -20,11 +20,15 @@ pub struct IntegrityLimits { pub min_reference_ratio: f64, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum AliasIssue { + #[error("model {model:?} has aliases that are not an array")] InvalidList { model: String }, + #[error("model {model:?} has an alias that is not a string")] InvalidName { model: String }, + #[error("alias {alias:?} of model {model:?} is already a model name")] CanonicalCollision { model: String, alias: String }, + #[error("alias {alias:?} of model {model:?} is already an alias of another model")] AliasCollision { model: String, alias: String }, } diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index edb6ba1eb11..00902858be4 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -2,26 +2,20 @@ use thiserror::Error; #[derive(Debug, Error)] pub enum Error { - /// The body is not valid JSON, or a model entry fails typed deserialization. - #[error("invalid JSON: {0}")] + #[error("invalid catalog JSON")] Json(#[from] serde_json::Error), - /// The catalog has no entries at all. #[error("catalog is empty")] Empty, - /// A non-reserved top level value is not a JSON object. #[error("model {model:?} must be an object")] EntryNotObject { model: String }, - /// Canonical entry count is under the configured minimum. #[error("catalog has {actual} models, below minimum {minimum}")] BelowMinimum { actual: usize, minimum: usize }, - /// Canonical entry count is under the configured reference ratio. #[error("catalog has {actual} models, below {ratio} of reference count {reference}")] Shrunk { actual: usize, reference: usize, ratio: f64, }, - /// The configured minimum reference ratio is not finite or outside `[0, 1]`. #[error("minimum reference ratio must be finite and between zero and one")] InvalidRatio, } diff --git a/litellm-rust/crates/model-catalog/src/validation.rs b/litellm-rust/crates/model-catalog/src/validation.rs index 73f08a865c1..6dd4174148c 100644 --- a/litellm-rust/crates/model-catalog/src/validation.rs +++ b/litellm-rust/crates/model-catalog/src/validation.rs @@ -5,24 +5,57 @@ use thiserror::Error; use crate::{AliasIssue, Catalog, ModelInfo, UtcHours, Weekday}; -/// A registry entry violates the checked-in catalog contract. #[derive(Debug, Error)] pub enum RegistryValidationError { - #[error("{reason}")] - Entry { model: String, reason: String }, - #[error("alias issue: {0:?}")] + #[error("invalid registry entry {model:?}")] + Entry { + model: String, + #[source] + reason: EntryViolation, + }, + #[error(transparent)] Alias(AliasIssue), } -/// Validate one registry entry without restricting the tolerant catalog reader. +#[derive(Debug, Error)] +pub enum EntryViolation { + #[error("entry must be an object")] + NotObject, + #[error("aliases must be an array")] + AliasesNotArray, + #[error("aliases must contain strings")] + AliasNotString, + #[error("entry does not match the model schema")] + Schema(#[source] serde_json::Error), + #[error("litellm_provider is required")] + MissingProvider, + #[error("invalid deprecation_date")] + DeprecationDate(#[source] time::error::Parse), + #[error("off_peak_pricing needs hours or windows")] + OffPeakWithoutHours, + #[error("off_peak_pricing.hours_utc is invalid")] + InvalidHours, + #[error("off_peak_pricing.windows is empty")] + EmptyWindows, + #[error("off_peak_pricing.weekdays is invalid")] + InvalidWeekdays, + #[error("typed entry could not be serialized")] + Serialize(#[source] serde_json::Error), + #[error("entry has an unknown field, null, or changed value: {}", fields.join(", "))] + ChangedFields { fields: Vec }, + #[error("{field} must be nonnegative")] + Negative { field: String }, + #[error("{field} must be at least one")] + UpliftBelowOne { field: String }, +} + pub fn validate_model_entry(model: &str, value: &Value) -> Result<(), RegistryValidationError> { - validate_entry_inner(model, value).map_err(|reason| RegistryValidationError::Entry { + validate_entry_inner(value).map_err(|reason| RegistryValidationError::Entry { model: model.to_owned(), reason, }) } -/// Check every model and alias in a parsed catalog against registry rules. pub fn validate_registry(catalog: &Catalog) -> Result<(), RegistryValidationError> { if let Some(issue) = catalog.alias_issues().first() { return Err(RegistryValidationError::Alias(issue.clone())); @@ -53,70 +86,61 @@ fn keys(value: &Map) -> BTreeSet { value.keys().cloned().collect() } -fn symmetric_difference(left: &BTreeSet, right: &BTreeSet) -> BTreeSet { +fn symmetric_difference(left: &BTreeSet, right: &BTreeSet) -> Vec { left.symmetric_difference(right).cloned().collect() } -fn validate_entry_inner(model_name: &str, value: &Value) -> Result<(), String> { - let object = value - .as_object() - .ok_or_else(|| format!("{model_name} must be an object"))?; +fn validate_entry_inner(value: &Value) -> Result<(), EntryViolation> { + let object = value.as_object().ok_or(EntryViolation::NotObject)?; if let Some(aliases) = object.get("aliases") { - let names = aliases - .as_array() - .ok_or_else(|| format!("{model_name}.aliases must be an array"))?; + let names = aliases.as_array().ok_or(EntryViolation::AliasesNotArray)?; if names.iter().any(|name| !name.is_string()) { - return Err(format!("{model_name}.aliases must contain strings")); + return Err(EntryViolation::AliasNotString); } } - let info: ModelInfo = - serde_json::from_value(value.clone()).map_err(|error| format!("{model_name}: {error}"))?; + let info: ModelInfo = serde_json::from_value(value.clone()).map_err(EntryViolation::Schema)?; if info.litellm_provider.is_none() { - return Err(format!("{model_name}.litellm_provider is required")); + return Err(EntryViolation::MissingProvider); } - validate_dates_and_windows(model_name, &info)?; - let serialized = serde_json::to_value(info).map_err(|error| error.to_string())?; + validate_dates_and_windows(&info)?; + let serialized = serde_json::to_value(info).map_err(EntryViolation::Serialize)?; let mut expected = object.clone(); expected.remove("aliases"); if !json_eq(&Value::Object(expected.clone()), &serialized) { let actual = serialized .as_object() .expect("ModelInfo serializes as an object"); - return Err(format!( - "{model_name} has an unknown field, null, or changed value: {:?}", - symmetric_difference(&keys(&expected), &keys(actual)) - )); + return Err(EntryViolation::ChangedFields { + fields: symmetric_difference(&keys(&expected), &keys(actual)), + }); } - check_prices(model_name, value) + check_prices("", value) } -fn validate_dates_and_windows(model_name: &str, info: &ModelInfo) -> Result<(), String> { +fn validate_dates_and_windows(info: &ModelInfo) -> Result<(), EntryViolation> { if let Some(date) = &info.deprecation_date { let format = time::format_description::parse_borrowed::<2>("[year]-[month]-[day]").unwrap(); - time::Date::parse(date, &format) - .map_err(|error| format!("{model_name}.deprecation_date: {error}"))?; + time::Date::parse(date, &format).map_err(EntryViolation::DeprecationDate)?; } let Some(pricing) = &info.off_peak_pricing else { return Ok(()); }; if pricing.hours_utc.is_none() && pricing.windows.is_none() { - return Err(format!( - "{model_name}.off_peak_pricing needs hours or windows" - )); + return Err(EntryViolation::OffPeakWithoutHours); } if let Some(hours) = &pricing.hours_utc { validate_hours(hours)?; } if let Some(windows) = &pricing.windows { if windows.is_empty() { - return Err(format!("{model_name}.off_peak_pricing.windows is empty")); + return Err(EntryViolation::EmptyWindows); } windows.iter().try_for_each(|window| { validate_hours(&window.hours_utc)?; if let Some(days) = &window.weekdays && (days.is_empty() || days.iter().any(|day| !valid_weekday(day))) { - return Err(format!("{model_name}.off_peak_pricing.weekdays is invalid")); + return Err(EntryViolation::InvalidWeekdays); } Ok(()) })?; @@ -124,13 +148,13 @@ fn validate_dates_and_windows(model_name: &str, info: &ModelInfo) -> Result<(), Ok(()) } -fn validate_hours(hours: &UtcHours) -> Result<(), String> { +fn validate_hours(hours: &UtcHours) -> Result<(), EntryViolation> { let values = match hours { UtcHours::Single(value) => std::slice::from_ref(value), UtcHours::Multiple(values) => values.as_slice(), }; if values.is_empty() || values.iter().any(|value| !valid_utc_window(value)) { - return Err("off_peak_pricing.hours_utc is invalid".into()); + return Err(EntryViolation::InvalidHours); } Ok(()) } @@ -176,7 +200,7 @@ fn valid_weekday(day: &Weekday) -> bool { } } -fn check_prices(path: &str, value: &Value) -> Result<(), String> { +fn check_prices(path: &str, value: &Value) -> Result<(), EntryViolation> { let Some(object) = value.as_object() else { return Ok(()); }; @@ -188,13 +212,17 @@ fn check_prices(path: &str, value: &Value) -> Result<(), String> { && let Some(number) = field.as_f64() && number < 0.0 { - return Err(format!("{field_path} must be nonnegative")); + return Err(EntryViolation::Negative { + field: relative(&field_path), + }); } if key.contains("uplift_multiplier") && let Some(number) = field.as_f64() && number < 1.0 { - return Err(format!("{field_path} must be at least one")); + return Err(EntryViolation::UpliftBelowOne { + field: relative(&field_path), + }); } if matches!(key.as_str(), "range" | "max_results_range") && field.as_array().is_some_and(|values| { @@ -203,7 +231,9 @@ fn check_prices(path: &str, value: &Value) -> Result<(), String> { .any(|value| value.as_f64().is_some_and(|n| n < 0.0)) }) { - return Err(format!("{field_path} must be nonnegative")); + return Err(EntryViolation::Negative { + field: relative(&field_path), + }); } if matches!(key.as_str(), "metadata" | "provider_specific_entry") { return Ok(()); @@ -216,3 +246,7 @@ fn check_prices(path: &str, value: &Value) -> Result<(), String> { } }) } + +fn relative(field_path: &str) -> String { + field_path.trim_start_matches('.').to_owned() +} diff --git a/litellm-rust/crates/model-catalog/tests/registry_validation.rs b/litellm-rust/crates/model-catalog/tests/registry_validation.rs index 67589010aba..3b65ef13d78 100644 --- a/litellm-rust/crates/model-catalog/tests/registry_validation.rs +++ b/litellm-rust/crates/model-catalog/tests/registry_validation.rs @@ -52,9 +52,9 @@ fn fallback_generalizations_are_typed(repo_root: PathBuf) { #[case::null_aliases(serde_json::json!({"litellm_provider": "test", "aliases": null}), "aliases must be an array")] #[ignore] fn registry_validation_rejects_malformed_entries(#[case] entry: Value, #[case] expected: &str) { + let error = validate_model_entry("test", &entry).unwrap_err(); assert!( - validate_model_entry("test", &entry) - .unwrap_err() + litellm_error::Report::new(&error) .to_string() .contains(expected) ); diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index c435ba37e78..d90fe073b20 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -111,8 +111,11 @@ GIL handling to `litellm-host-python`. - Avoid copying large payloads more than needed. The current JSON round-trip is acceptable for the first scaffold, but future performance work should evaluate direct PyO3 conversion before expanding Rust coverage to image-heavy paths. -- Do not expose raw Rust errors that include document contents or upstream - bodies. +- Render a Rust error for Python only in its mapping function, with + `litellm_error::Report`, never `Debug` or a document's contents. An upstream + body reaches Python intact as the `RustUpstreamError` args for provider error + mapping. The 256-character cap applies only to HTTP error diagnostics, not + response bodies. Do not log those args. ## Tests diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 8afde8362f6..9d69cd8c14b 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -27,6 +27,7 @@ litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true +litellm-error.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true litellm-cache-response.workspace = true @@ -70,6 +71,7 @@ criterion.workspace = true futures-util.workspace = true rstest.workspace = true sqlx = { workspace = true, features = ["migrate"] } +tempfile.workspace = true tokio-tungstenite.workspace = true wiremock.workspace = true aws-sdk-secretsmanager = "1.117.0" diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index e190b12e665..87dba3efcb1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -12,10 +12,93 @@ use pyo3::{ prelude::*, }; -fn cache_error(error: Error) -> PyErr { - match error { - Error::InvalidEntry => PyValueError::new_err(error.to_string()), - Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), - _ => PyRuntimeError::new_err(error.to_string()), +/// A cache call the bridge rejects itself, before or after the cache ran. +#[derive(Debug, thiserror::Error)] +pub(crate) enum HandleError { + #[error("recreate the v2 cache after fork")] + Forked, + #[error("cache limits must be positive")] + Limits, + #[error("namespace and a positive cache limit are required")] + RedisLimits, + #[error("cache durations must be finite and positive")] + Duration(#[source] Option), + #[error("unsupported v2 cache control: {0}")] + Control(String), + #[error("native inference needs a v2 cache, the catalog keeps Python caches on Python")] + NotNative, + #[error("cache operation without configured cache")] + Unconfigured, + #[error("cache arguments unavailable")] + MissingArguments, + #[error("cache reply without pending operation")] + NothingPending, +} + +impl From for PyErr { + fn from(error: HandleError) -> Self { + let message = litellm_error::Report::new(&error).to_string(); + match error { + HandleError::Limits + | HandleError::RedisLimits + | HandleError::Duration(_) + | HandleError::Control(_) => PyValueError::new_err(message), + HandleError::Forked + | HandleError::NotNative + | HandleError::Unconfigured + | HandleError::MissingArguments + | HandleError::NothingPending => PyRuntimeError::new_err(message), + } + } +} + +fn cache_error(error: Error) -> PyErr { + let message = litellm_error::Report::new(&error).to_string(); + match error { + Error::InvalidEntry(_) => PyValueError::new_err(message), + Error::UnsupportedOperation => PyNotImplementedError::new_err(message), + Error::Unavailable(_) | Error::UnscopedFlush | Error::MissingPrompt => { + PyRuntimeError::new_err(message) + } + } +} + +#[cfg(test)] +mod tests { + use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + }; + use rstest::rstest; + + use super::{Error, cache_error}; + + fn refused() -> std::io::Error { + std::io::Error::other("connection refused") + } + + #[rstest] + #[case::unavailable( + Error::unavailable(refused()), + false, + "cache is unavailable: connection refused" + )] + #[case::invalid_entry( + Error::invalid_entry(refused()), + true, + "invalid cache entry: connection refused" + )] + fn python_sees_the_class_and_the_whole_chain( + #[case] error: Error, + #[case] caller_fault: bool, + #[case] message: &str, + ) { + Python::initialize(); + Python::attach(|py| { + let failure = cache_error(error); + assert_eq!(failure.is_instance_of::(py), caller_fault); + assert_eq!(failure.is_instance_of::(py), !caller_fault); + assert_eq!(failure.value(py).to_string(), message); + }); } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs index eb09b8131c3..c54cb641843 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs @@ -1,4 +1,4 @@ -use crate::cache::cache_error; +use crate::cache::{HandleError, cache_error}; use std::{ sync::Arc, time::{Duration, SystemTime, UNIX_EPOCH}, @@ -14,7 +14,7 @@ use litellm_cache_response::{ CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, }; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; #[pyclass( frozen, @@ -37,9 +37,7 @@ enum Storage { impl NativeCacheHandle { fn check_process(&self) -> PyResult<()> { if self.pid != std::process::id() { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "recreate the v2 cache after fork", - )); + return Err(HandleError::Forked.into()); } Ok(()) } @@ -61,7 +59,7 @@ impl NativeCacheHandle { fn memory(ttl: f64, capacity: usize, max_entry_bytes: usize) -> PyResult { let ttl = duration(ttl)?; if capacity == 0 || max_entry_bytes == 0 { - return Err(PyValueError::new_err("cache limits must be positive")); + return Err(HandleError::Limits.into()); } let storage = Arc::new(InMemoryCache::new(Some(capacity), Some(ttl))); let backend = Arc::new(ResponseCache::new(storage.clone()).with_config( @@ -89,9 +87,7 @@ impl NativeCacheHandle { ) -> PyResult { let ttl = duration(ttl)?; if namespace.is_empty() || max_entry_bytes == 0 { - return Err(PyValueError::new_err( - "namespace and a positive cache limit are required", - )); + return Err(HandleError::RedisLimits.into()); } let storage = Arc::new( release_gil(py, || { @@ -252,11 +248,12 @@ impl NativeCacheHandle { } } -fn duration(seconds: f64) -> PyResult { - Duration::try_from_secs_f64(seconds) - .ok() - .filter(|value| !value.is_zero()) - .ok_or_else(|| PyValueError::new_err("cache durations must be finite and positive")) +fn duration(seconds: f64) -> Result { + match Duration::try_from_secs_f64(seconds) { + Ok(duration) if !duration.is_zero() => Ok(duration), + Ok(_) => Err(HandleError::Duration(None)), + Err(error) => Err(HandleError::Duration(Some(error))), + } } pub(in crate::cache) fn native_handle<'py>( @@ -277,11 +274,7 @@ pub(in crate::cache) fn configured( Option>, litellm_cache_response::CacheOptions, )> { - let handle = native_handle(configured)?.ok_or_else(|| { - pyo3::exceptions::PyRuntimeError::new_err( - "native inference needs a v2 cache, the catalog keeps Python caches on Python", - ) - })?; + let handle = native_handle(configured)?.ok_or(HandleError::NotNative)?; let cache = handle.extract::>()?; cache.check_process()?; let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none()); @@ -296,9 +289,7 @@ pub(in crate::cache) fn configured( name.as_str(), "no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache" ) { - return Err(PyValueError::new_err(format!( - "unsupported v2 cache control: {name}" - ))); + return Err(HandleError::Control(name).into()); } } } @@ -316,7 +307,7 @@ pub(in crate::cache) fn configured( .map(|values| values.get_item(name)) .transpose()? .flatten() - .map(|value| duration(value.extract()?)) + .map(|value| Ok(duration(value.extract()?)?)) .transpose() }; Ok(( diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs index d8f8ac7c181..82b7d881df6 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/host.rs @@ -9,6 +9,7 @@ use pyo3::{ use serde_json::Value; use super::service::CacheCall; +use crate::cache::HandleError; enum Pending { Lookup(Reply, Error>>), @@ -43,16 +44,12 @@ impl PythonCache { pub fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult>> { let Some(cache) = self.cache.as_ref() else { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "cache operation without configured cache", - )); + return Err(HandleError::Unconfigured.into()); }; let arguments = self .arguments .as_ref() - .ok_or_else(|| { - pyo3::exceptions::PyRuntimeError::new_err("cache arguments unavailable") - })? + .ok_or(HandleError::MissingArguments)? .bind(py) .copy()?; let (method, result) = match call { @@ -97,33 +94,30 @@ impl PythonCache { py: Python<'_>, result: PyResult>, ) -> PyResult>> { - if let Err(error) = &result - && !error.is_instance_of::(py) - { - self.pending = None; - return Err(result.err().unwrap()); - } + let result = match result { + Err(error) if !error.is_instance_of::(py) => { + self.pending = None; + return Err(error); + } + result => result, + }; match self.pending.take() { Some(Pending::Lookup(reply)) => { - let value = result.map_err(|_| Error::Unavailable).and_then(|value| { + let value = result.map_err(Error::unavailable).and_then(|value| { if value.bind(py).is_none() { Ok(None) } else { from_py(value.bind(py)) .map(Some) - .map_err(|_| Error::InvalidEntry) + .map_err(Error::invalid_entry) } }); reply.send(value); } Some(Pending::Store(reply)) => { - reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable)); - } - None => { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "cache reply without pending operation", - )); + reply.send(result.map(|_| ()).map_err(Error::unavailable)); } + None => return Err(HandleError::NothingPending.into()), } Ok(None) } diff --git a/litellm-rust/crates/python-bridge/src/cache/python/service.rs b/litellm-rust/crates/python-bridge/src/cache/python/service.rs index f215e5e97f5..38872b85d30 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/service.rs @@ -17,18 +17,21 @@ fn from_python(value: Value) -> Result { let output = match value.get(STREAM_EVENTS_KEY) { Some(events) => { let events: Vec = - serde_json::from_value(events.clone()).map_err(|_| Error::InvalidEntry)?; + serde_json::from_value(events.clone()).map_err(Error::invalid_entry)?; CachedOutput::Stream(events.concat()) } None => CachedOutput::Response(value), }; - serde_json::to_value(ResponseEnvelope::new("messages", output)).map_err(|_| Error::InvalidEntry) + serde_json::to_value(ResponseEnvelope::new("messages", output)).map_err(Error::invalid_entry) } fn to_python(value: Value) -> Result { let envelope: ResponseEnvelope> = - serde_json::from_value(value).map_err(|_| Error::InvalidEntry)?; - match envelope.decode("messages").ok_or(Error::InvalidEntry)? { + serde_json::from_value(value).map_err(Error::invalid_entry)?; + match envelope + .decode("messages") + .ok_or(Error::InvalidEntry(None))? + { CachedOutput::Response(response) => Ok(response), CachedOutput::Stream(text) => Ok(serde_json::json!({ STREAM_EVENTS_KEY: text.split_inclusive("\n\n").collect::>() @@ -56,7 +59,7 @@ pub(in crate::cache) fn service>( namespace: String, ) -> std::sync::Arc where - P::Error: From, + P::Error: From + std::error::Error, { std::sync::Arc::new(PythonCacheService { services, @@ -69,7 +72,7 @@ where impl> ResponseCacheService for PythonCacheService

where - P::Error: From, + P::Error: From + std::error::Error, { fn config(&self) -> &ResponseCacheConfig { &self.config @@ -84,7 +87,7 @@ where self.services .call(|reply| CacheCall::Lookup { reply }) .await - .map_err(|_| Error::Unavailable)?? + .map_err(Error::unavailable)?? .map(from_python) .transpose() }) @@ -101,7 +104,7 @@ where self.services .call(|reply| CacheCall::Store { value, reply }) .await - .map_err(|_| Error::Unavailable)? + .map_err(Error::unavailable)? }) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index cd19533fc73..00fc7cd822f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -39,7 +39,7 @@ impl Selection { services: HostServices

, ) -> (Option, CacheOptions) where - P::Error: From, + P::Error: From + std::error::Error, { let service = match self.backend { Backend::Disabled => None, diff --git a/litellm-rust/crates/python-bridge/src/callable.rs b/litellm-rust/crates/python-bridge/src/callable.rs index d0ed76a3e8b..605b649e1f3 100644 --- a/litellm-rust/crates/python-bridge/src/callable.rs +++ b/litellm-rust/crates/python-bridge/src/callable.rs @@ -48,14 +48,8 @@ mod tests { Err(PyErr::from_value(error.clone())) } - #[rstest::rstest] - fn only_ordinary_exceptions_are_reported_under_the_template() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - pyo3::ffi::c_str!( - r#" + const FAILURES: &std::ffi::CStr = pyo3::ffi::c_str!( + r#" class CallerError(Exception): def __format__(self, specification): return 'unavailable' @@ -63,12 +57,19 @@ ordinary = CallerError('must use __format__') type_error = TypeError('signature') abort = KeyboardInterrupt('cancelled') "# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); + ); + fn failures(py: Python<'_>) -> Bound<'_, PyDict> { + let locals = PyDict::new(py); + py.run(FAILURES, Some(&locals), Some(&locals)).unwrap(); + locals + } + + #[test] + fn ordinary_exceptions_are_reported_under_the_template() { + Python::initialize(); + Python::attach(|py| { + let locals = failures(py); let original = raised(&locals, "ordinary"); let wrapped = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err(); assert!(wrapped.is_instance_of::(py)); @@ -78,12 +79,19 @@ abort = KeyboardInterrupt('cancelled') wrapped.value(py).str().unwrap().to_str().unwrap(), "Failed to reach the caller: unavailable" ); + }); + } - for name in ["type_error", "abort"] { - let original = raised(&locals, name); - let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err(); - assert!(error.value(py).is(&original)); - } + #[rstest::rstest] + #[case::type_error("type_error")] + #[case::keyboard_interrupt("abort")] + fn other_exceptions_pass_through_unwrapped(#[case] name: &str) { + Python::initialize(); + Python::attach(|py| { + let locals = failures(py); + let original = raised(&locals, name); + let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err(); + assert!(error.value(py).is(&original)); }); } diff --git a/litellm-rust/crates/python-bridge/src/coercion.rs b/litellm-rust/crates/python-bridge/src/coercion.rs index 77c8cabc22b..8cdf1e6c5bc 100644 --- a/litellm-rust/crates/python-bridge/src/coercion.rs +++ b/litellm-rust/crates/python-bridge/src/coercion.rs @@ -1,4 +1,5 @@ use litellm_core_utils::serde_compat::parse_str_bool; +use litellm_error::Report; use pyo3::{ exceptions::{PyAttributeError, PyRuntimeError, PyValueError}, prelude::*, @@ -9,15 +10,20 @@ use pyo3::{ pub(crate) enum ProjectionError { Python(PyErr), InvalidConfiguration(ProjectionDetail), - UnsupportedLiveObject(String), + UnsupportedLiveObject(ProjectionDetail), InternalSchemaFailure(ProjectionDetail), } #[derive(Debug, thiserror::Error)] pub(crate) enum ProjectionDetail { - #[error("{0}")] - Message(String), - #[error("{field}: {source}")] + #[error("{group}.{name}: expected {expected}, got {actual}")] + Expected { + group: &'static str, + name: &'static str, + expected: &'static str, + actual: String, + }, + #[error("{field}")] InvalidField { field: &'static str, #[source] @@ -30,12 +36,6 @@ pub(crate) enum ProjectionDetail { }, } -impl From for ProjectionDetail { - fn from(message: String) -> Self { - Self::Message(message) - } -} - impl From for ProjectionError { fn from(error: PyErr) -> Self { Self::Python(error) @@ -46,12 +46,12 @@ impl From for PyErr { fn from(error: ProjectionError) -> Self { match error { ProjectionError::Python(error) => error, - ProjectionError::InvalidConfiguration(detail) => { - PyValueError::new_err(detail.to_string()) + ProjectionError::InvalidConfiguration(detail) + | ProjectionError::UnsupportedLiveObject(detail) => { + PyValueError::new_err(Report::new(&detail).to_string()) } - ProjectionError::UnsupportedLiveObject(message) => PyValueError::new_err(message), ProjectionError::InternalSchemaFailure(detail) => { - PyRuntimeError::new_err(detail.to_string()) + PyRuntimeError::new_err(Report::new(&detail).to_string()) } } } @@ -124,26 +124,26 @@ impl<'py> Field<'py> { && getter.is(object.getattr("__getattribute__")?)) } - pub(crate) fn path(&self) -> String { - format!("{}.{}", self.group, self.name) - } - /// A member of this field's collection, reported under the same path. pub(crate) fn member(&self, value: Bound<'py, PyAny>) -> Self { Self::new(self.group, self.name, value) } - pub(crate) fn expected(&self, expected: &str) -> Result { - Ok(format!( - "{}: expected {expected}, got {}", - self.path(), - self.value.get_type().name()? - )) + pub(crate) fn expected( + &self, + expected: &'static str, + ) -> Result { + Ok(ProjectionDetail::Expected { + group: self.group, + name: self.name, + expected, + actual: self.value.get_type().name()?.to_string(), + }) } - pub(crate) fn invalid(&self, expected: &str) -> ProjectionError { + pub(crate) fn invalid(&self, expected: &'static str) -> ProjectionError { match self.expected(expected) { - Ok(message) => ProjectionError::InvalidConfiguration(message.into()), + Ok(detail) => ProjectionError::InvalidConfiguration(detail), Err(error) => error, } } @@ -163,24 +163,23 @@ impl<'py> Field<'py> { pub(crate) fn schema_bool(&self) -> Result { if !self.value.is_instance_of::() { return Err(ProjectionError::InternalSchemaFailure( - self.expected("a Boolean")?.into(), + self.expected("a Boolean")?, )); } Ok(self.exact_true()) } pub(crate) fn strict_string(&self) -> Result { - let value = self - .value - .cast::() - .map_err(|_| self.invalid("a string"))?; + let Ok(value) = self.value.cast::() else { + return Err(self.invalid("a string")); + }; Ok(value.to_str()?.to_owned()) } pub(crate) fn schema_string(&self) -> Result { if !self.value.is_instance_of::() { return Err(ProjectionError::InternalSchemaFailure( - self.expected("a string")?.into(), + self.expected("a string")?, )); } self.strict_string() @@ -474,12 +473,7 @@ text = Text(' False ') }); } - #[test] - fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() { - with_python(|py| { - let locals = PyDict::new(py); - py.run( - c" + const RAISING_ATTRIBUTES: &std::ffi::CStr = c" failure = AttributeError('descriptor failed') class Snapshot: @property @@ -491,46 +485,56 @@ class Intercepted: def __getattribute__(self, name): raise failure dynamic = Dynamic() intercepted = Intercepted() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let snapshot = locals.get_item("snapshot").unwrap().unwrap(); - let descriptor = PyErr::from(Field::read(&snapshot, "test", "flag").err().unwrap()); +"; + + #[rstest] + #[case::descriptor("snapshot")] + #[case::dynamic("dynamic")] + #[case::intercepted("intercepted")] + fn attribute_errors_raised_by_a_snapshot_surface_as_raised(#[case] name: &str) { + with_python(|py| { + let locals = PyDict::new(py); + py.run(RAISING_ATTRIBUTES, Some(&locals), Some(&locals)) + .unwrap(); + let value = locals.get_item(name).unwrap().unwrap(); + let error = PyErr::from(Field::read(&value, "test", "flag").err().unwrap()); assert!( - descriptor + error .value(py) .is(locals.get_item("failure").unwrap().unwrap()) ); - for name in ["dynamic", "intercepted"] { - let value = locals.get_item(name).unwrap().unwrap(); - let error = PyErr::from(Field::read(&value, "test", "flag").err().unwrap()); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - } + }); + } + + #[test] + fn missing_snapshot_fields_are_runtime_errors_naming_the_field() { + with_python(|py| { + let locals = PyDict::new(py); + py.run(RAISING_ATTRIBUTES, Some(&locals), Some(&locals)) + .unwrap(); + let snapshot = locals.get_item("snapshot").unwrap().unwrap(); let missing = PyErr::from(Field::read(&snapshot, "test", "missing").err().unwrap()); assert!(missing.is_instance_of::(py)); assert!(missing.to_string().contains("test.missing")); }); } - #[test] - fn configuration_errors_name_fields_without_exposing_values() { + #[rstest] + #[case::mapping("{'secret': 'do-not-print'}")] + #[case::nested("['host.test', {'secret': 'do-not-print'}]")] + fn configuration_errors_name_fields_without_exposing_values(#[case] source: &str) { + with_python(|py| { + let field = Field::new("test", "setting", evaluate(py, source)); + let error = PyErr::from(field.falsy_optional_string().err().unwrap()); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("test.setting")); + assert!(!error.to_string().contains("do-not-print")); + }); + } + + #[test] + fn invalid_collections_and_flags_are_configuration_errors() { with_python(|py| { - for source in [ - "{'secret': 'do-not-print'}", - "['host.test', {'secret': 'do-not-print'}]", - ] { - let field = Field::new("test", "setting", evaluate(py, source)); - let error = PyErr::from(field.falsy_optional_string().err().unwrap()); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("test.setting")); - assert!(!error.to_string().contains("do-not-print")); - } let hosts = Field::new( "url_policy", "user_url_allowed_hosts", diff --git a/litellm-rust/crates/python-bridge/src/credentials.rs b/litellm-rust/crates/python-bridge/src/credentials.rs index 04d8e94c8fe..a4bb064aae0 100644 --- a/litellm-rust/crates/python-bridge/src/credentials.rs +++ b/litellm-rust/crates/python-bridge/src/credentials.rs @@ -10,10 +10,19 @@ use pyo3::{ types::{PyDict, PyString}, }; -const NOT_CALLABLE: &str = "Azure AD token provider must be callable"; const NOT_A_STRING: &str = "Azure AD token must be a string, got {}"; const FAILED: &str = "Failed to get Azure AD token: {}"; +#[derive(Debug, thiserror::Error)] +#[error("Azure AD token provider must be callable")] +struct NotCallable; + +impl From for PyErr { + fn from(error: NotCallable) -> Self { + PyTypeError::new_err(litellm_error::Report::new(&error).to_string()) + } +} + /// The `azure_ad_token_provider` keyword argument, kept alive for the rest of the call. pub(crate) struct CallerTokenProvider { provider: Py, @@ -36,7 +45,7 @@ impl CallerTokenProvider { pub(crate) fn acquire(&self, py: Python<'_>) -> PyResult { let provider = self.provider.bind(py); if !provider.is_callable() { - return Err(PyTypeError::new_err(NOT_CALLABLE)); + return Err(NotCallable.into()); } let token = wrap_failure( py, diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs index 687a090e768..5368da4f930 100644 --- a/litellm-rust/crates/python-bridge/src/diagnostics.rs +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -1,5 +1,5 @@ use litellm_host_python::{release_count, runtime_started}; -use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; #[pyfunction] pub(crate) fn gil_stats(py: Python<'_>) -> PyResult> { @@ -18,8 +18,7 @@ pub(crate) fn process_state_started() -> bool { /// so the runtime can never start. Raises if it already has. Forked workers are unaffected. #[pyfunction] pub(crate) fn reserve_process_for_forking() -> PyResult<()> { - litellm_host_python::reserve_process_for_forking() - .map_err(|_| PyRuntimeError::new_err("the native runtime already started in this process")) + Ok(litellm_host_python::reserve_process_for_forking()?) } #[cfg(feature = "panic-test")] diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 4dab0297e73..b42c08b7e78 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -1,9 +1,12 @@ -use litellm_http::transport::Error as TransportError; +use std::error::Error as StdError; + +use litellm_error::Report; use litellm_inference::Error; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, }; +use reqwest::header::HeaderMap; pyo3::create_exception!( _native, @@ -16,20 +19,34 @@ pyo3::create_exception!( _native, RustUpstreamError, pyo3::exceptions::PyException, - "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." + "The provider call was already issued and the provider answered with a failure. Args are (status, body); the headers attribute holds the provider's (name, value) response headers." ); -pub(crate) fn route_error_to_pyerr(error: Error) -> PyErr { - match error { - Error::Transport(TransportError::Http { status, body }) => { - RustUpstreamError::new_err((status, body)) - } - other => by_fault(other.is_request(), other.to_string()), - } +const MESSAGES_REQUEST_ERROR: &str = "messages_request_error"; + +pub(crate) fn upstream_error( + py: Python<'_>, + status: u16, + body: &str, + headers: &HeaderMap, +) -> PyResult { + let error = RustUpstreamError::new_err((status, body.to_owned())); + let pairs: Vec<(&str, String)> = headers + .iter() + .map(|(name, value)| { + ( + name.as_str(), + String::from_utf8_lossy(value.as_bytes()).into_owned(), + ) + }) + .collect(); + error.value(py).setattr("headers", pairs)?; + Ok(error) } /// A request the caller got wrong is a `ValueError`; anything else is a `RuntimeError`. -pub(crate) fn by_fault(is_request: bool, message: String) -> PyErr { +pub(crate) fn by_fault(is_request: bool, error: &(dyn StdError + 'static)) -> PyErr { + let message = Report::new(error).to_string(); if is_request { PyValueError::new_err(message) } else { @@ -37,17 +54,90 @@ pub(crate) fn by_fault(is_request: bool, message: String) -> PyErr { } } +/// How a route flags a rejected request for its Python route module, which raises the public +/// exception for it instead of mapping the native one. +#[derive(Clone, Copy)] +pub(crate) enum RequestMarker { + Unmarked, + Messages, +} + +/// Maps a route failure once. A core decline and a Python exception raised inside the call are +/// returned as they are, everything else is mapped and handed to `public`, the route's public +/// exception mapping. +pub(crate) fn route_failure( + py: Python<'_>, + error: Error, + marker: RequestMarker, + public: impl FnOnce(PyErr) -> PyResult, +) -> PyResult { + if let Some(decline) = error.declined() { + return Ok(RustBridgeDeclined::new_err(decline.reason)); + } + if let Some(raised) = crate::secrets::python_error(py, &error) { + return Ok(raised); + } + if let Some((status, body)) = error.upstream_response() { + return public(upstream_error(py, status, body, &HeaderMap::new())?); + } + let native = match (marker, error.rejected_request()) { + (RequestMarker::Messages, Some(rejection)) => { + let native = PyValueError::new_err(Report::new(rejection).to_string()); + native.value(py).setattr(MESSAGES_REQUEST_ERROR, true)?; + native + } + (RequestMarker::Messages, None) | (RequestMarker::Unmarked, _) => { + by_fault(error.is_request(), &error) + } + }; + public(native) +} + +pub(crate) fn route_error_to_pyerr(error: Error) -> PyErr { + Python::attach(|py| { + route_failure(py, error, RequestMarker::Unmarked, Ok).unwrap_or_else(std::convert::identity) + }) +} + #[cfg(test)] mod tests { + use litellm_core_utils::decline::Decline; + use reqwest::header::{HeaderValue, RETRY_AFTER}; + use super::*; + fn upstream_payload(py: Python<'_>, failure: &PyErr) -> ((u16, String), Vec<(String, String)>) { + let value = failure.value(py); + ( + value.getattr("args").unwrap().extract().unwrap(), + value.getattr("headers").unwrap().extract().unwrap(), + ) + } + #[rstest::rstest] - #[case::unsupported(Error::Unsupported("test capability"), true)] #[case::invalid_provider(Error::InvalidProvider("unknown".into()), true)] - #[case::invalid_request(Error::InvalidRequest("empty messages".into()), true)] - #[case::connection(TransportError::Connect("unreachable".into()).into(), false)] - #[case::network(TransportError::Network("timed out".into()).into(), false)] - #[case::invalid_response(Error::InvalidResponse("missing usage".into()), false)] + #[case::invalid_request( + Error::invalid_request(litellm_llms::ErrorDetail::Expected { + subject: "messages", + expected: "non-empty", + }), + true + )] + #[case::lost_connection( + litellm_http::Error::from(tokio_tungstenite::tungstenite::Error::ConnectionClosed).into(), + false + )] + #[case::invalid_response( + Error::invalid_response(litellm_llms::ErrorDetail::MissingField { + subject: "response", + field: "usage", + }), + false + )] + #[case::post_call_decline( + Error::post_call(Error::Declined(Decline::new("result rejected"))), + false + )] fn route_failures_are_terminal(#[case] error: Error, #[case] is_request: bool) { Python::initialize(); Python::attach(|py| { @@ -59,29 +149,97 @@ mod tests { } #[rstest::rstest] - fn transport_status_survives_python_mapping() { + #[case::unmarked(RequestMarker::Unmarked)] + #[case::messages(RequestMarker::Messages)] + fn a_core_decline_hands_the_call_back_to_python_before_the_public_mapping( + #[case] marker: RequestMarker, + ) { Python::initialize(); Python::attach(|py| { - let upstream = route_error_to_pyerr( - TransportError::Http { - status: 429, - body: "slow down".into(), - } - .into(), - ); - assert!(upstream.is_instance_of::(py)); + let error = Error::from(litellm_llms::Error::Declined(Decline::new( + "unrecognized request parameter", + ))); + let declined = route_failure(py, error, marker, |_| { + panic!("a decline must not reach the public mapping") + }) + .unwrap(); + assert!(declined.is_instance_of::(py)); assert_eq!( - upstream + declined .value(py) .getattr("args") .unwrap() - .extract::<(u16, String)>() + .extract::<(String,)>() .unwrap(), - (429, "slow down".into()) + ("unrecognized request parameter".into(),) ); }); } + #[rstest::rstest] + #[case::unmarked(RequestMarker::Unmarked)] + #[case::messages(RequestMarker::Messages)] + fn the_provider_answer_passes_through_with_a_headers_attribute(#[case] marker: RequestMarker) { + Python::initialize(); + Python::attach(|py| { + let error = litellm_http::Error::upstream(429, "slow down".into()).into(); + let upstream = route_failure(py, error, marker, Ok).unwrap(); + assert!(upstream.is_instance_of::(py)); + assert_eq!( + upstream_payload(py, &upstream), + ((429, "slow down".into()), Vec::new()) + ); + assert!( + upstream + .value(py) + .getattr_opt(MESSAGES_REQUEST_ERROR) + .unwrap() + .is_none() + ); + }); + } + + #[test] + fn upstream_error_keeps_the_provider_headers() { + Python::initialize(); + Python::attach(|py| { + let headers = HeaderMap::from_iter([(RETRY_AFTER, HeaderValue::from_static("17"))]); + let upstream = upstream_error(py, 503, "busy", &headers).unwrap(); + assert_eq!( + upstream_payload(py, &upstream), + ( + (503, "busy".into()), + vec![("retry-after".into(), "17".into())] + ) + ); + }); + } + + #[rstest::rstest] + #[case::ascii("busy ".repeat(80))] + #[case::unicode("é".repeat(300))] + fn long_provider_bodies_reach_python_intact( + #[case] message: String, + #[values(RequestMarker::Unmarked, RequestMarker::Messages)] marker: RequestMarker, + ) { + let envelope = serde_json::json!({ + "type": "error", + "error": {"type": "overloaded_error", "message": message}, + "request_id": "provider-request", + }); + let body = envelope.to_string(); + let error = litellm_http::Error::upstream(529, body.clone()).into(); + Python::initialize(); + Python::attach(|py| { + let upstream = route_failure(py, error, marker, Ok).unwrap(); + assert!(upstream.is_instance_of::(py)); + let ((status, returned_body), headers) = upstream_payload(py, &upstream); + assert_eq!(status, 529); + assert_eq!(returned_body, body); + assert!(headers.is_empty()); + }); + } + #[test] fn missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors() { Python::initialize(); @@ -95,4 +253,83 @@ mod tests { assert!(invalid.is_instance_of::(py)); }); } + + #[test] + fn a_wrapped_failure_renders_every_layer() { + Python::initialize(); + Python::attach(|py| { + let error = Error::post_call( + litellm_http::Error::upstream(500, "hook backend down".into()).into(), + ); + let report = Report::new(&error).to_string(); + let failure = route_error_to_pyerr(error); + assert_eq!(failure.value(py).to_string(), report); + assert!(report.contains("hook backend down"), "{report}"); + }); + } + + #[rstest::rstest] + #[case::rejected_request( + Error::invalid_request(litellm_llms::ErrorDetail::UnsupportedParameter { + model: "claude".into(), + param: "top_k".into(), + value: "5".into(), + hint: String::new(), + }), + Some("claude does not support top_k=5. To drop unsupported params, set `litellm.drop_params = True`.") + )] + #[case::missing_field( + litellm_llms::Error::MissingField("max_tokens").into(), + Some("missing required field: max_tokens") + )] + #[case::unresolvable_provider(Error::InvalidProvider("openai".into()), None)] + #[case::invalid_response( + Error::invalid_response(litellm_llms::ErrorDetail::MissingField { + subject: "response", + field: "usage", + }), + None + )] + fn only_messages_request_rejections_carry_the_marker_and_python_text( + #[case] error: Error, + #[case] rejection: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let unmarked = route_failure(py, error.clone(), RequestMarker::Unmarked, Ok).unwrap(); + assert!( + unmarked + .value(py) + .getattr_opt(MESSAGES_REQUEST_ERROR) + .unwrap() + .is_none() + ); + let native = route_failure(py, error, RequestMarker::Messages, Ok).unwrap(); + let marked = native + .value(py) + .getattr_opt(MESSAGES_REQUEST_ERROR) + .unwrap() + .map(|value| value.extract::().unwrap()); + assert_eq!(marked.unwrap_or(false), rejection.is_some()); + if let Some(rejection) = rejection { + assert!(native.is_instance_of::(py)); + assert_eq!(native.value(py).to_string(), rejection); + } + }); + } + + #[test] + fn a_python_secret_manager_exception_skips_the_public_mapping() { + Python::initialize(); + Python::attach(|py| { + let raised = pyo3::exceptions::PyKeyError::new_err("vault down"); + let identity = raised.value(py).clone().unbind(); + let error = Error::from(crate::secrets::read_error(py, raised)); + let surfaced = route_failure(py, error, RequestMarker::Unmarked, |_| { + panic!("the raised exception must not be remapped") + }) + .unwrap(); + assert!(surfaced.value(py).is(identity.bind(py))); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 57bf8aafc4e..b4e20a8f750 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -132,26 +132,14 @@ pub(crate) fn host_client(py: Python<'_>, variant: ClientVariant) -> PyResult PyErr { - match error { - litellm_http::Error::Read { - tls_source: TlsSource::ClientIdentity, - .. - } - | litellm_http::Error::InvalidPem { - tls_source: TlsSource::ClientIdentity, - .. - } => PyValueError::new_err( + match error.tls_source() { + Some(TlsSource::ClientIdentity) => PyValueError::new_err( "http_settings.ssl_certificate: expected a readable PEM certificate and private key", ), - litellm_http::Error::Read { - tls_source: TlsSource::CaBundle, - .. + Some(TlsSource::CaBundle) => { + PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle") } - | litellm_http::Error::InvalidPem { - tls_source: TlsSource::CaBundle, - .. - } => PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle"), - _ => PyValueError::new_err("http_settings: native HTTP client configuration is invalid"), + None => PyValueError::new_err(litellm_error::Report::new(&error).to_string()), } } @@ -274,21 +262,28 @@ settings = types.SimpleNamespace(**defaults) fn client_error_uses_tls_source_when_paths_match() { Python::initialize(); Python::attach(|py| { - let path = PathBuf::from("/shared.pem"); - let ca_error = client_error(litellm_http::Error::InvalidPem { - path: path.clone(), - message: "invalid".into(), - tls_source: TlsSource::CaBundle, - }); + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(file.path(), b"not a certificate").unwrap(); + let path = file.path().to_path_buf(); + let defaults = Resolution::from(&HttpSettings::default()).config; + let build = |config: HttpClientConfig| { + HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client(&config, ClientVariant::Provider) + .map(drop) + .unwrap_err() + }; + let ca_error = client_error(build(HttpClientConfig { + verify: Verify::CaBundle(path.clone()), + ..defaults.clone() + })); assert_eq!( ca_error.to_string(), "ValueError: http_settings.ssl_verify: expected a readable PEM CA bundle" ); - let client_error = client_error(litellm_http::Error::InvalidPem { - path, - message: "invalid".into(), - tls_source: TlsSource::ClientIdentity, - }); + let client_error = client_error(build(HttpClientConfig { + client_certificate: Some(litellm_http::ClientIdentity::Pem(path)), + ..defaults + })); assert!(client_error.is_instance_of::(py)); assert_eq!( client_error.to_string(), @@ -297,6 +292,34 @@ settings = types.SimpleNamespace(**defaults) }); } + #[test] + fn client_error_without_a_tls_source_reports_the_build_failure() { + Python::initialize(); + Python::attach(|py| { + let error = HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client( + &HttpClientConfig { + user_agent: Some("bad\nagent".into()), + ..Resolution::from(&HttpSettings::default()).config + }, + ClientVariant::Provider, + ) + .map(drop) + .unwrap_err(); + assert_eq!(error.tls_source(), None); + let report = litellm_error::Report::new(&error).to_string(); + + let failure = client_error(error); + + assert!(failure.is_instance_of::(py)); + assert_eq!(failure.value(py).str().unwrap().to_string(), report); + assert!( + report.contains("could not build the HTTP client: "), + "{report}" + ); + }); + } + #[test] fn python_settings_flow_into_the_configured_layer() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/logger/mod.rs b/litellm-rust/crates/python-bridge/src/logger/mod.rs index bf5c735360b..8f40d94f458 100644 --- a/litellm-rust/crates/python-bridge/src/logger/mod.rs +++ b/litellm-rust/crates/python-bridge/src/logger/mod.rs @@ -81,8 +81,12 @@ impl NativeDiagnosticProcessor { } } -fn processing_error(_: fancy_regex::Error) -> PyErr { - PyRuntimeError::new_err("diagnostic processing failed") +#[derive(Debug, thiserror::Error)] +#[error("diagnostic processing failed")] +struct ProcessingError(#[source] fancy_regex::Error); + +fn processing_error(error: fancy_regex::Error) -> PyErr { + PyRuntimeError::new_err(litellm_error::Report::new(&ProcessingError(error)).to_string()) } struct PythonSink { diff --git a/litellm-rust/crates/python-bridge/src/logger/tests.rs b/litellm-rust/crates/python-bridge/src/logger/tests.rs index 1fca3be720e..9a07428808a 100644 --- a/litellm-rust/crates/python-bridge/src/logger/tests.rs +++ b/litellm-rust/crates/python-bridge/src/logger/tests.rs @@ -187,6 +187,7 @@ fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication locals .set_item("http_warning", wrap_pyfunction!(http_warning, py).unwrap()) .unwrap(); + locals.set_item("native", crate::native_module(py)).unwrap(); let importable = py .eval( c"__import__('importlib.util', fromlist=['util']).find_spec('dotenv') is not None", @@ -205,6 +206,7 @@ import asyncio import logging import sys sys.path.insert(0, repo_root) +sys.modules['litellm.rust_bridge._native'] = native import litellm from litellm._logging import verbose_logger, session_id_var, trace_id_var diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 038fa8c9630..dec8ff0b0ce 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -21,10 +21,27 @@ pub(crate) struct RouteOptions { pub(crate) timeout: Option, } -fn required_object(name: &'static str, value: Value) -> PyResult> { +/// A caller argument the bridge cannot use, raised as a `ValueError`. +#[derive(Debug, thiserror::Error)] +pub(crate) enum ArgumentError { + #[error("{0} is required")] + Required(&'static str), + #[error("{0} must be a dict")] + NotADict(&'static str), + #[error("header values must be strings")] + HeaderValue, +} + +impl From for PyErr { + fn from(error: ArgumentError) -> Self { + PyValueError::new_err(litellm_error::Report::new(&error).to_string()) + } +} + +fn required_object(name: &'static str, value: Value) -> Result, ArgumentError> { match value { Value::Object(values) => Ok(values), - _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), + _ => Err(ArgumentError::NotADict(name)), } } @@ -35,7 +52,7 @@ fn optional_object( if value.is_none() { return Ok(None); } - required_object(name, from_py_argument(value)?).map(Some) + Ok(Some(required_object(name, from_py_argument(value)?)?)) } pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option { @@ -57,11 +74,11 @@ pub(crate) fn python_timeout_seconds(py: Python<'_>, timeout: Py) -> PyRe pub(crate) fn required_field<'py>( fields: &Bound<'py, PyDict>, - name: &str, + name: &'static str, ) -> PyResult> { - fields + Ok(fields .get_item(name)? - .ok_or_else(|| PyValueError::new_err(format!("{name} is required"))) + .ok_or(ArgumentError::Required(name))?) } pub(crate) fn optional_field( @@ -171,21 +188,18 @@ pub(crate) fn request_input_sources<'a>( .collect()) } -pub(crate) fn marshal_headers(headers: Option) -> PyResult> { - let value = match headers { - Some(headers) => headers, - None => Value::Object(Map::new()), - }; - let Value::Object(headers) = value else { - return Err(PyValueError::new_err("headers must be a dict")); +pub(crate) fn marshal_headers( + headers: Option, +) -> Result, ArgumentError> { + let headers = match headers { + Some(headers) => required_object("headers", headers)?, + None => Map::new(), }; headers .into_iter() - .map(|(name, value)| { - value - .as_str() - .map(|value| (name, value.to_string())) - .ok_or_else(|| PyValueError::new_err("header values must be strings")) + .map(|(name, value)| match value { + Value::String(value) => Ok((name, value)), + _ => Err(ArgumentError::HeaderValue), }) .collect() } @@ -395,48 +409,52 @@ module.__getattr__ = fail float: Option, } - #[test] - fn numeric_adapters_agree_across_json_and_python_boundaries() { + #[rstest::rstest] + #[case::empty(json!({}))] + #[case::explicit_nulls(json!({"integers": null, "float": null}))] + #[case::lax_values(json!({ + "integers": [i64::MIN, i64::MAX, "9007199254740993.0", " +1_000.00 ", true, 3.0], + "float": " 1.25 ", + }))] + #[case::integer_overflow(json!({"integers": [u64::MAX]}))] + #[case::imprecise_integer_string(json!({"integers": ["1.0000000000000001"]}))] + #[case::fractional_integer(json!({"integers": [2.5]}))] + #[case::nan_string(json!({"float": "NaN"}))] + #[case::infinity_string(json!({"float": "inf"}))] + #[case::overflowing_float_string(json!({"float": "1e999"}))] + #[case::boolean_float(json!({"float": true}))] + #[case::large_integer_float(json!({"float": u64::MAX}))] + fn numeric_adapters_agree_across_json_and_python_boundaries(#[case] input: Value) { Python::initialize(); Python::attach(|py| { - for input in [ - json!({}), - json!({"integers": null, "float": null}), - json!({"integers": [i64::MIN, i64::MAX, "9007199254740993.0", " +1_000.00 ", true, 3.0], "float": " 1.25 "}), - json!({"integers": [u64::MAX]}), - json!({"integers": ["1.0000000000000001"]}), - json!({"integers": [2.5]}), - json!({"float": "NaN"}), - json!({"float": "inf"}), - json!({"float": "1e999"}), - json!({"float": true}), - json!({"float": u64::MAX}), - ] { - let expected = serde_json::from_value::(input.clone()); - let python = litellm_host_python::to_py(py, &input).unwrap(); - let actual = from_py::(python.bind(py)); - match (expected, actual) { - (Ok(expected), Ok(actual)) => { - assert_eq!(actual, expected); - let serialized = litellm_host_python::to_py(py, &actual).unwrap(); - assert_eq!( - from_py::(serialized.bind(py)).unwrap(), - serde_json::to_value(expected).unwrap() - ); - } - (Err(_), Err(_)) => {} - mismatch => panic!("boundary mismatch for {input}: {mismatch:?}"), + let expected = serde_json::from_value::(input.clone()); + let python = litellm_host_python::to_py(py, &input).unwrap(); + let actual = from_py::(python.bind(py)); + match (expected, actual) { + (Ok(expected), Ok(actual)) => { + assert_eq!(actual, expected); + let serialized = litellm_host_python::to_py(py, &actual).unwrap(); + assert_eq!( + from_py::(serialized.bind(py)).unwrap(), + serde_json::to_value(expected).unwrap() + ); } + (Err(_), Err(_)) => {} + mismatch => panic!("boundary mismatch for {input}: {mismatch:?}"), } - for source in [ - c"{'float': float('nan')}", - c"{'float': float('inf')}", - c"{'integers': [float('inf')]}", - c"{'integers': [2 ** 100]}", - ] { - let value = py.eval(source, None, None).unwrap(); - assert!(from_py::(&value).is_err()); - } + }); + } + + #[rstest::rstest] + #[case::nan(c"{'float': float('nan')}")] + #[case::infinity(c"{'float': float('inf')}")] + #[case::infinite_integer(c"{'integers': [float('inf')]}")] + #[case::integer_overflow(c"{'integers': [2 ** 100]}")] + fn non_finite_and_overflowing_python_numbers_are_rejected(#[case] source: &std::ffi::CStr) { + Python::initialize(); + Python::attach(|py| { + let value = py.eval(source, None, None).unwrap(); + assert!(from_py::(&value).is_err()); }); } diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index e62673d4d4f..21a57123921 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -13,7 +13,7 @@ use crate::{ }; async fn execute( - http: Result, + http: litellm_http::Client, secrets: std::sync::Arc, audio: Value, optional_params: Map, @@ -27,7 +27,7 @@ async fn execute( extra_headers, timeout, } = options; - AudioTranscriptionRoute::new(http?, crate::http::resources().auth.clone(), secrets) + AudioTranscriptionRoute::new(http, crate::http::resources().auth.clone(), secrets) .execute(AudioTranscriptionRequest { model: &model, audio, @@ -48,7 +48,8 @@ pub(crate) fn transcription(py: Python<'_>, call: NativeCall<'_>) -> PyResult( litellm_host_python::from_py_argument(&required_field(&arguments, "audio")?)?; let options = value_route_options(&arguments)?; let optional_params = optional_object_field(&arguments, "optional_params")?.unwrap_or_default(); - let http = crate::http::provider_client(py, &call.kwargs, true)?; + let http = + crate::http::provider_client(py, &call.kwargs, true)?.map_err(crate::http::client_error)?; let secrets = crate::secrets::source(py)?; run_async( py, diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 5ad9f8df8b8..4db9b37cbb8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use std::convert::Infallible; use crate::routes::inference::InferenceHost; @@ -44,9 +45,9 @@ impl PythonBinding for ChatCompletionsPythonHost { .get("stream") .is_some_and(|value| value == &serde_json::Value::Bool(true)) { - return Err(InvokeError::Native(Error::Unsupported( + return Err(InvokeError::Native(Error::Declined(Decline::new( "native Python chat_completions streaming", - ))); + )))); } Ok(call) } @@ -79,7 +80,7 @@ impl PythonBinding for ChatCompletionsPythonHost { self.0.error(py, error) } fn host_error(error: &PyErr) -> Error { - Error::InvalidRequest(error.to_string().into()) + crate::routes::inference::host_call_error(error) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/clickhouse_spend.rs b/litellm-rust/crates/python-bridge/src/routes/clickhouse_spend.rs index f8df7a9de23..80d7bd5360e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/clickhouse_spend.rs +++ b/litellm-rust/crates/python-bridge/src/routes/clickhouse_spend.rs @@ -8,24 +8,13 @@ use pyo3::{ }; fn map_error(error: Error) -> PyErr { - use litellm_storage_clickhouse::Error as StorageError; - match error { - Error::InsertTooLarge | Error::Storage(StorageError::InsertTooLarge) => { - PyOverflowError::new_err(error.to_string()) - } - Error::InvalidRow - | Error::InvalidLimit(_) - | Error::InvalidSchema - | Error::Storage( - StorageError::InvalidRow - | StorageError::InvalidLimit(_) - | StorageError::InvalidTable - | StorageError::InvalidSchema - | StorageError::EmptySql - | StorageError::InvalidParameters - | StorageError::InvalidQuery, - ) => PyValueError::new_err(error.to_string()), - Error::Migration(_) | Error::Storage(_) => PyRuntimeError::new_err(error.to_string()), + let message = litellm_error::Report::new(&error).to_string(); + if error.is_insert_too_large() { + PyOverflowError::new_err(message) + } else if error.is_request() { + PyValueError::new_err(message) + } else { + PyRuntimeError::new_err(message) } } @@ -99,26 +88,22 @@ mod tests { use super::*; use rstest::rstest; + use litellm_storage_clickhouse::Error as StorageError; + #[rstest] - #[case::row(Error::InvalidRow, "ValueError")] - #[case::insert_limit(Error::InvalidLimit("CLICKHOUSE_TRACE_MAX_INSERT_BYTES"), "ValueError")] - #[case::insert_timeout( - Error::Storage(litellm_storage_clickhouse::Error::InvalidLimit( - "CLICKHOUSE_INSERT_TIMEOUT_SECONDS" - )), + #[case::timestamp(Error::NonIntegerTimestamp { field: "start_time" }, "ValueError")] + #[case::insert_limit( + StorageError::InvalidLimit("CLICKHOUSE_TRACE_MAX_INSERT_BYTES").into(), "ValueError" )] - #[case::insert_budget(Error::InsertTooLarge, "OverflowError")] - #[case::schema( - Error::Storage(litellm_storage_clickhouse::Error::SchemaFailed(503)), - "RuntimeError" - )] + #[case::insert_budget(StorageError::InsertTooLarge.into(), "OverflowError")] + #[case::schema(StorageError::SchemaFailed(503).into(), "RuntimeError")] #[case::migration( Error::Migration(sqlx::migrate::MigrateError::VersionMismatch(6)), "RuntimeError" )] #[case::storage( - Error::Storage(litellm_storage_clickhouse::Error::InvalidUrl), + Config::new("spend".into(), "ftp://localhost", 1).err().unwrap(), "RuntimeError" )] fn spend_failures_preserve_public_exception_types( @@ -127,7 +112,7 @@ mod tests { ) { Python::initialize(); Python::attach(|py| { - let message = error.to_string(); + let message = litellm_error::Report::new(&error).to_string(); let exception = map_error(error); assert_eq!(exception.get_type(py).name().unwrap(), exception_name); assert_eq!( @@ -136,4 +121,15 @@ mod tests { ); }); } + + #[test] + fn spend_failures_render_their_cause() { + let error = Config::new("spend".into(), "not a url", 1).err().unwrap(); + let cause = std::error::Error::source(&error).unwrap().to_string(); + Python::initialize(); + Python::attach(|py| { + let message = map_error(error).value(py).str().unwrap().to_string(); + assert!(message.ends_with(&cause), "{message}"); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/routes/embeddings.rs b/litellm-rust/crates/python-bridge/src/routes/embeddings.rs index 0f73eefedd5..0935f42f85a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/embeddings.rs +++ b/litellm-rust/crates/python-bridge/src/routes/embeddings.rs @@ -1,14 +1,23 @@ +use litellm_error::Report; use pyo3::prelude::*; use super::NativeCall; use crate::errors::RustBridgeDeclined; +#[derive(Debug, thiserror::Error)] +#[error("native embeddings route is not implemented")] +struct NotImplemented; + +impl From for PyErr { + fn from(error: NotImplemented) -> Self { + RustBridgeDeclined::new_err(Report::new(&error).to_string()) + } +} + #[pyfunction] pub(crate) fn embedding(call: NativeCall<'_>) -> PyResult> { drop(call); - Err(RustBridgeDeclined::new_err( - "native embeddings route is not implemented", - )) + Err(NotImplemented.into()) } #[pyfunction] diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index c9eb60e47ed..ca30c94b02f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -6,11 +6,9 @@ use litellm_callbacks_legacy_python::LoggingOperation; use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{call::HostedMachine, protocol::Protocol}; use litellm_host_python::{PythonBinding, PythonHostCalls, from_py}; -use litellm_http::transport::Error as TransportError; use litellm_inference::Error; use litellm_secrets::source::SecretSource; use pyo3::{ - exceptions::PyValueError, gc::{PyTraverseError, PyVisit}, prelude::*, types::PyDict, @@ -20,8 +18,10 @@ use serde_json::{Map, Value}; use super::{NativeCall, RequestView}; use crate::{ - errors::{RustUpstreamError, route_error_to_pyerr}, - marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds}, + errors::{RequestMarker, route_failure}, + marshal::{ + ArgumentError, RouteOptions, optional_timeout, public_response, python_timeout_seconds, + }, }; pub(super) struct InferenceHost { @@ -44,7 +44,7 @@ impl InferenceHost { &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - input: &str, + input: &'static str, ) -> PyResult { self.view.resolve(py, arguments)?; let argument = |name: &str| self.argument(py, arguments, name); @@ -57,7 +57,7 @@ impl InferenceHost { .map(|value| python_timeout_seconds(py, value.unbind())) .transpose()? .flatten(); - let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?; + let model = string("model")?.ok_or(ArgumentError::Required("model"))?; let custom_llm_provider = string("custom_llm_provider")?; let provider = get_custom_llm_provider(&model, custom_llm_provider.as_deref()) .map_or("", |resolved| resolved.custom_llm_provider); @@ -82,10 +82,7 @@ impl InferenceHost { .transpose()?, timeout: optional_timeout(timeout), }, - input: from_py( - &argument(input)? - .ok_or_else(|| PyValueError::new_err(format!("{input} is required")))?, - )?, + input: from_py(&argument(input)?.ok_or(ArgumentError::Required(input))?)?, params, }) } @@ -120,29 +117,20 @@ impl InferenceHost { } pub fn error(&self, py: Python<'_>, error: Error) -> PyResult { - if let Error::Secret(source) = &error - && let Some(original) = crate::secrets::python_error(py, source.source_error()) - { - return Ok(original); - } - let native = match error { - Error::Transport(TransportError::Http { status, body }) => { - let error = RustUpstreamError::new_err((status, body)); - error - .value(py) - .setattr("headers", Vec::<(String, String)>::new())?; - error - } - other => route_error_to_pyerr(other), - }; - let mapped = py - .import(self.module)? - .getattr("map_failure")? - .call1((native.value(py), self.view.request(py)?))?; - Ok(PyErr::from_value(mapped)) + route_failure(py, error, RequestMarker::Unmarked, |native| { + let mapped = py + .import(self.module)? + .getattr("map_failure")? + .call1((native.value(py), self.view.request(py)?))?; + Ok(PyErr::from_value(mapped)) + }) } } +pub(super) fn host_call_error(_: &PyErr) -> Error { + Error::invalid_request(litellm_llms::ErrorDetail::HostCall) +} + pub(super) trait InferenceRoute: Sized + 'static { type Protocol: Protocol; const OPERATION: LoggingOperation; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 5e445fcaa7d..27ce8c943f0 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -3,7 +3,6 @@ use litellm_host_python::{PythonHostCalls, PythonOwned}; use bytes::Bytes; use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py}; -use litellm_http::transport::Error as TransportError; use litellm_inference_messages::{ Error, MessagesCall, MessagesSettings, MessagesShaping, litellm_params, messages_body, route::{Messages, MessagesStreamHead}, @@ -12,7 +11,7 @@ use litellm_llms::base_llm::messages::context::MessagesModelCapabilities; use litellm_llms_types::headers::ProviderSpecificHeaders; use litellm_router_types::LitellmParams; use pyo3::{ - exceptions::{PyException, PyValueError}, + exceptions::PyException, gc::{PyTraverseError, PyVisit}, prelude::*, types::{PyBytes, PyDict}, @@ -20,14 +19,16 @@ use pyo3::{ use serde_json::{Map, Value}; use crate::{ - errors::{RustUpstreamError, route_error_to_pyerr}, - marshal::{optional_timeout, project_optional_fields, public_response, python_timeout_seconds}, + errors::{RequestMarker, route_failure}, + marshal::{ + ArgumentError, optional_timeout, project_optional_fields, public_response, + python_timeout_seconds, + }, python_settings::missing_module, routes::RequestView, }; const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; -const REQUEST_ERROR_MARKER: &str = "messages_request_error"; fn merge_headers( forwarded: Option>, @@ -68,29 +69,6 @@ fn project_litellm_params<'py>( Ok(litellm_params(fields.into_iter().chain(folded).collect())) } -fn native_error(py: Python<'_>, error: Error) -> PyResult { - match error { - Error::Transport(TransportError::Http { status, body }) => { - let error = RustUpstreamError::new_err((status, body)); - error - .value(py) - .setattr("headers", Vec::<(String, String)>::new())?; - Ok(error) - } - Error::InvalidRequest(message) => { - let error = PyValueError::new_err(message.to_string()); - error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; - Ok(error) - } - Error::MissingField(field) => { - let error = PyValueError::new_err(format!("missing required field: {field}")); - error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; - Ok(error) - } - other => Ok(route_error_to_pyerr(other)), - } -} - /// The Python side of the Messages route: projects the prepared arguments and builds the /// public response, chunks and exceptions. pub(super) struct MessagesPythonHost { @@ -115,9 +93,8 @@ impl MessagesPythonHost { let string = |name: &str| -> PyResult> { argument(name)?.map(|value| value.extract()).transpose() }; - let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?; - let messages = - argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?; + let model = string("model")?.ok_or(ArgumentError::Required("model"))?; + let messages = argument("messages")?.ok_or(ArgumentError::Required("messages"))?; let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?; let body = [ ("model".to_string(), Value::String(model.clone())), @@ -276,16 +253,13 @@ impl PythonBinding for MessagesPythonHost { } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { - if let Error::Secret(source) = &error - && let Some(original) = crate::secrets::python_error(py, source.source_error()) - { - return Ok(original); - } - Ok(self.map_failure(py, native_error(py, error)?)) + route_failure(py, error, RequestMarker::Messages, |native| { + Ok(self.map_failure(py, native)) + }) } fn host_error(error: &PyErr) -> Error { - Error::InvalidRequest(error.to_string().into()) + crate::routes::inference::host_call_error(error) } } @@ -459,33 +433,9 @@ mod tests { projected, litellm_params(serde_json::from_value(fields).unwrap()).unwrap() ), - (Err(error), Err(())) => assert!(matches!(error, Error::InvalidRequest(_))), + (Err(error), Err(())) => assert!(error.is_request(), "{error:?}"), (projected, expected) => panic!("got {projected:?}, expected {expected:?}"), } }); } - - #[rstest] - #[case::rejected_request(Error::InvalidRequest("does not support top_k=5".into()), true)] - #[case::missing_field(Error::MissingField("max_tokens"), true)] - #[case::unresolvable_provider(Error::InvalidProvider("openai".into()), false)] - #[case::upstream_failure( - Error::Transport(TransportError::Http { status: 400, body: "bad".into() }), - false, - )] - fn only_request_rejections_carry_the_request_error_marker( - #[case] error: Error, - #[case] marked: bool, - ) { - Python::initialize(); - Python::attach(|py| { - let native = native_error(py, error).unwrap(); - let marker = native - .value(py) - .getattr_opt(REQUEST_ERROR_MARKER) - .unwrap() - .map(|value| value.extract::().unwrap()); - assert_eq!(marker.unwrap_or(false), marked); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs index f85d7d4c7cb..71eb743138d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -9,6 +9,29 @@ use pyo3::{ types::{PyBytes, PyString, PyType}, }; +/// A `type='file'` document the bridge cannot read, with the text Python's OCR path raises. +#[derive(Debug, thiserror::Error)] +enum DocumentError { + #[error( + "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes" + )] + MissingFile, + #[error( + "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object." + )] + BareString, + #[error( + "Unsupported file input type: {0}. Expected pathlib.Path, bytes, or a file-like object." + )] + UnsupportedType(String), +} + +impl From for PyErr { + fn from(error: DocumentError) -> Self { + PyValueError::new_err(litellm_error::Report::new(&error).to_string()) + } +} + /// A `type='file'` document as projected: paths and bytes are typed inputs already; a /// file-like object is a reader the projection consumes once every other field is read. pub(super) enum FileDocumentInput { @@ -45,25 +68,18 @@ impl FromPyObject<'_, '_> for FileDocumentInput { Err(error) if error.is_instance_of::(py) => None, Err(error) => return Err(error), }; - let missing = || { - PyValueError::new_err( - "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes", - ) - }; - let file = document.get_item("file").map_err(|error| { - if error.is_instance_of::(py) { - missing() - } else { - error + let file = match document.get_item("file") { + Ok(file) => file, + Err(error) if error.is_instance_of::(py) => { + return Err(DocumentError::MissingFile.into()); } - })?; + Err(error) => return Err(error), + }; if file.is_none() { - return Err(missing()); + return Err(DocumentError::MissingFile.into()); } if file.is_instance_of::() { - return Err(PyValueError::new_err( - "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", - )); + return Err(DocumentError::BareString.into()); } static PATH_LIKE: PyOnceLock> = PyOnceLock::new(); if file.is_instance(PATH_LIKE.import(py, "os", "PathLike")?)? { @@ -81,10 +97,7 @@ impl FromPyObject<'_, '_> for FileDocumentInput { } match PythonFileReader::from_file_like(&file)? { Some(reader) => Ok(Self::Deferred { reader, mime_type }), - None => Err(PyValueError::new_err(format!( - "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", - file.get_type(), - ))), + None => Err(DocumentError::UnsupportedType(file.get_type().to_string()).into()), } } } @@ -92,6 +105,7 @@ impl FromPyObject<'_, '_> for FileDocumentInput { #[cfg(test)] mod tests { use pyo3::{exceptions::PyTypeError, types::PyDict}; + use rstest::rstest; use super::*; @@ -108,40 +122,47 @@ mod tests { } } - #[test] - fn extraction_validates_required_file_and_optional_mime_type() { + enum Rejection { + Value(&'static str), + Type, + } + + #[rstest] + #[case::missing_file(c"{}", Rejection::Value("must include a 'file' field"))] + #[case::none_file(c"{'file': None}", Rejection::Value("must include a 'file' field"))] + #[case::none_mime_type(c"{'file': b'abc', 'mime_type': None}", Rejection::Type)] + #[case::numeric_mime_type(c"{'file': b'abc', 'mime_type': 7}", Rejection::Type)] + #[case::bare_str(c"{'file': 'scan.pdf'}", Rejection::Value("bare str"))] + #[case::unsupported_type( + c"{'file': object()}", + Rejection::Value("Unsupported file input type") + )] + fn extraction_rejects_invalid_documents( + #[case] expression: &std::ffi::CStr, + #[case] rejection: Rejection, + ) { Python::initialize(); Python::attach(|py| { - for expression in [c"{}", c"{'file': None}"] { - let document = py.eval(expression, None, None).unwrap(); - let error = document.extract::().err().unwrap(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("must include a 'file' field")); - } - for expression in [ - c"{'file': b'abc', 'mime_type': None}", - c"{'file': b'abc', 'mime_type': 7}", - ] { - let document = py.eval(expression, None, None).unwrap(); - let error = document.extract::().err().unwrap(); - assert!(error.is_instance_of::(py)); - } let error = py - .eval(c"{'file': 'scan.pdf'}", None, None) + .eval(expression, None, None) .unwrap() .extract::() .err() .unwrap(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("bare str")); - let error = py - .eval(c"{'file': object()}", None, None) - .unwrap() - .extract::() - .err() - .unwrap(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("Unsupported file input type")); + match rejection { + Rejection::Value(fragment) => { + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains(fragment), "{error}"); + } + Rejection::Type => assert!(error.is_instance_of::(py)), + } + }); + } + + #[test] + fn extraction_keeps_bytes_and_their_mime_type() { + Python::initialize(); + Python::attach(|py| { let document = py .eval(c"{'file': b'abc', 'mime_type': 'image/png'}", None, None) .unwrap(); diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index 2068b6a6e4b..31736c7b654 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -1,128 +1,203 @@ +use std::io::ErrorKind; + +use litellm_error::Report; use litellm_llms::base_llm::ocr::error::Error; use pyo3::{ exceptions::{PyFileNotFoundError, PyOSError}, prelude::*, }; -use crate::errors::{RustUpstreamError, by_fault}; +use crate::errors::{by_fault, upstream_error}; + +const REQUEST_FORMAT_ERROR: &str = "ocr_request_format_error"; pub(super) fn to_pyerr(error: Error) -> PyErr { - let status = error.http_status_code(); - let mapped = Python::attach(|py| -> PyResult { - Ok(match error { - Error::Provider { - status, - body, - headers, - } => upstream_error(py, status, body, headers)?, - Error::Transport(litellm_http::transport::Error::Http { status, body }) => { - upstream_error(py, status, body, Vec::new())? - } - Error::RequestFormat => { - let error = by_fault(true, Error::RequestFormat.to_string()); - error - .value(py) - .setattr("ocr_request_format_error", true) - .ok(); - error - } - Error::FileRead { path, source } if source.kind() == std::io::ErrorKind::NotFound => { - PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) - } - Error::FileRead { source, .. } => PyOSError::new_err(source.to_string()), - other => by_fault(is_request(&other), other.to_string()), - }) - }) - .unwrap_or_else(|error| error); - attach_status(mapped, status) + Python::attach(|py| native(py, error).unwrap_or_else(std::convert::identity)) } -fn is_request(error: &Error) -> bool { - error.is_request() - || matches!( - error, - Error::Auth(_) - | Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::MissingField(_) - | Error::MissingDocumentUrl - ) -} - -fn upstream_error( +/// Maps an OCR route failure once. A Python exception raised inside the call is returned as it +/// is, everything else is mapped and handed to `public`, the route's public exception mapping. +pub(super) fn route_failure( py: Python<'_>, - status: u16, - body: String, - headers: Vec<(String, String)>, + error: Error, + public: impl FnOnce(PyErr) -> PyResult, ) -> PyResult { - let error = RustUpstreamError::new_err((status, body)); - error.value(py).setattr("headers", headers)?; - Ok(error) + match crate::secrets::python_error(py, &error) { + Some(raised) => Ok(raised), + None => public(native(py, error)?), + } } -fn attach_status(error: PyErr, status: Option) -> PyErr { - if let Some(status) = status { - Python::attach(|py| { - let value = error.value(py); - value.setattr("status_code", status).ok(); - value.setattr("message", value.to_string()).ok(); - }); +fn native(py: Python<'_>, error: Error) -> PyResult { + let mapped = match (&error, error.upstream_response()) { + (_, Some((status, body, headers))) => upstream_error(py, status, body, headers)?, + (Error::RequestFormat, None) => { + let mapped = by_fault(true, &error); + mapped.value(py).setattr(REQUEST_FORMAT_ERROR, true)?; + mapped + } + (Error::FileRead { path, cause }, None) if cause.kind() == ErrorKind::NotFound => { + PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) + } + (Error::FileRead { cause, .. }, None) => { + PyOSError::new_err(Report::new(cause.as_ref()).to_string()) + } + (_, None) => by_fault(error.is_request(), &error), + }; + if let Some(status) = error.http_status_code() { + let value = mapped.value(py); + value.setattr("status_code", status)?; + value.setattr("message", value.str()?)?; } - error + Ok(mapped) } #[cfg(test)] mod tests { - use pyo3::exceptions::PyValueError; + use pyo3::exceptions::{PyRuntimeError, PyValueError}; + use rstest::rstest; use super::*; + use crate::errors::RustUpstreamError; - #[test] - fn preserves_python_validation_and_provider_details() { + fn status_code(py: Python<'_>, mapped: &PyErr) -> Option { + mapped + .value(py) + .getattr_opt("status_code") + .unwrap() + .map(|status| status.extract().unwrap()) + } + + fn file_read(kind: ErrorKind) -> Error { + Error::FileRead { + path: "/missing/scan.pdf".into(), + cause: std::sync::Arc::new(std::io::Error::new(kind, "disk said no")), + } + } + + #[rstest] + #[case::provider( + Error::Provider { + status: 429, + body: r#"{"message":"rate limited"}"#.into(), + headers: Box::new(reqwest::header::HeaderMap::from_iter([( + reqwest::header::RETRY_AFTER, + reqwest::header::HeaderValue::from_static("17"), + )])), + }, + vec![("retry-after".into(), "17".into())] + )] + #[case::transport(litellm_http::Error::upstream(429, r#"{"message":"rate limited"}"#.into()).into(), vec![])] + fn a_provider_answer_keeps_its_status_body_and_headers( + #[case] error: Error, + #[case] headers: Vec<(String, String)>, + ) { Python::initialize(); Python::attach(|py| { - let mapped = to_pyerr(Error::MissingDocumentUrl); - assert!(mapped.is_instance_of::(py)); - assert_eq!(mapped.value(py).to_string(), "Document URL is required"); - assert_eq!( - mapped - .value(py) - .getattr("status_code") - .unwrap() - .extract::() - .unwrap(), - 400 - ); - let mapped = to_pyerr(Error::Provider { - status: 429, - body: r#"{"message":"rate limited"}"#.to_string(), - headers: vec![("Retry-After".to_string(), "17".to_string())], - }); + let mapped = to_pyerr(error); assert!(mapped.is_instance_of::(py)); - let headers: Vec<(String, String)> = mapped - .value(py) - .getattr("headers") - .and_then(|headers| headers.extract()) - .expect("OCR failures retain provider headers"); - assert_eq!(headers, vec![("Retry-After".to_string(), "17".to_string())]); - let args: (u16, String) = mapped - .value(py) - .getattr("args") - .and_then(|args| args.extract()) - .expect("OCR failures retain status and unprefixed provider message"); - assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string())); + let value = mapped.value(py); + let args: (u16, String) = value.getattr("args").unwrap().extract().unwrap(); + assert_eq!(args, (429, r#"{"message":"rate limited"}"#.into())); + let kept: Vec<(String, String)> = value.getattr("headers").unwrap().extract().unwrap(); + assert_eq!(kept, headers); + assert_eq!(status_code(py, &mapped), Some(429)); + }); + } - let mapped = to_pyerr(Error::InvalidRequest("invalid format".into())); + #[rstest] + #[case::not_found(Some(404), 404, "OCR document download failed")] + #[case::rate_limited(Some(429), 429, "OCR document download failed")] + #[case::timeout(None, 408, "OCR document download timed out")] + #[tokio::test] + async fn document_download_failures_keep_the_upstream_exception( + #[case] response_status: Option, + #[case] expected_status: u16, + #[case] expected_body: &str, + ) { + use std::time::Duration; + + use litellm_http::{Client, media::MediaFetcher}; + use litellm_llms::base_llm::ocr::{ + document::inline_remote_document, transformation::OcrConnection, + }; + use litellm_llms_types::formats::ocr::OcrDocument; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 1]; + socket.read_exact(&mut request).await.unwrap(); + if let Some(status) = response_status { + socket.write_all(format!("HTTP/1.1 {status} Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").as_bytes()).await.unwrap(); + } else { + std::future::pending::<()>().await; + } + }); + let connection = OcrConnection { + timeout: Duration::from_millis(100), + ..OcrConnection::default() + }; + let error = inline_remote_document( + &MediaFetcher::for_test(Client::for_test()), + OcrDocument::DocumentUrl { + document_url: format!("http://{address}/document.pdf"), + extra_fields: Default::default(), + }, + &connection, + ) + .await + .unwrap_err(); + server.abort(); + if response_status.is_none() { + assert!(litellm_error::find::(&error).is_some()); + } + Python::initialize(); + Python::attach(|py| { + let mapped = to_pyerr(error); + assert!(mapped.is_instance_of::(py), "{mapped}"); + let args: (u16, String) = mapped.value(py).getattr("args").unwrap().extract().unwrap(); + assert_eq!(args, (expected_status, expected_body.into())); + assert_eq!(status_code(py, &mapped), Some(expected_status)); + }); + } + + #[rstest] + #[case::missing_document_url(Error::MissingDocumentUrl)] + #[case::invalid_features(Error::Features)] + #[case::invalid_provider(Error::InvalidProvider("unknown".into()))] + fn request_faults_are_value_errors_with_status_400(#[case] error: Error) { + Python::initialize(); + Python::attach(|py| { + let report = Report::new(&error).to_string(); + let mapped = to_pyerr(error); assert!(mapped.is_instance_of::(py)); - assert_eq!( - mapped - .value(py) - .getattr("status_code") - .unwrap() - .extract::() - .unwrap(), - 400 - ); + assert_eq!(mapped.value(py).to_string(), report); + assert_eq!(status_code(py, &mapped), Some(400)); + }); + } + + #[rstest] + #[case::missing_api_key(litellm_auth::Error::MissingApiKey { + provider: "Mistral", + environment_variable: "MISTRAL_API_KEY", + } + .into())] + #[case::missing_credential(Error::MissingCredential { + credential: "Azure AI credentials", + guidance: "set AZURE_AI_API_KEY", + })] + fn a_missing_credential_is_a_statusless_runtime_error(#[case] error: Error) { + Python::initialize(); + Python::attach(|py| { + let mapped = to_pyerr(error); + assert!(mapped.is_instance_of::(py)); + assert_eq!(status_code(py, &mapped), None); }); } @@ -135,19 +210,12 @@ mod tests { assert!(mapped.is_instance_of::(py)); assert!( value - .getattr("ocr_request_format_error") + .getattr(REQUEST_FORMAT_ERROR) .unwrap() .extract::() .unwrap() ); - assert_eq!( - value - .getattr("status_code") - .unwrap() - .extract::() - .unwrap(), - 400 - ); + assert_eq!(status_code(py, &mapped), Some(400)); assert_eq!( value .getattr("message") @@ -159,18 +227,11 @@ mod tests { }); } - fn file_read(kind: std::io::ErrorKind) -> Error { - Error::FileRead { - path: "/missing/scan.pdf".into(), - source: std::sync::Arc::new(std::io::Error::new(kind, "disk said no")), - } - } - #[test] fn missing_files_map_to_file_not_found_naming_the_path() { Python::initialize(); Python::attach(|py| { - let mapped = to_pyerr(file_read(std::io::ErrorKind::NotFound)); + let mapped = to_pyerr(file_read(ErrorKind::NotFound)); assert!(mapped.is_instance_of::(py)); assert_eq!( mapped.value(py).to_string(), @@ -183,28 +244,32 @@ mod tests { fn other_file_read_failures_map_to_os_error_with_the_io_message() { Python::initialize(); Python::attach(|py| { - let mapped = to_pyerr(file_read(std::io::ErrorKind::PermissionDenied)); + let mapped = to_pyerr(file_read(ErrorKind::PermissionDenied)); assert!(mapped.is_instance_of::(py)); assert!(!mapped.is_instance_of::(py)); assert_eq!(mapped.value(py).to_string(), "disk said no"); }); } - #[rstest::rstest] + #[rstest] #[case::oversized(Error::TooLarge { limit: 7 })] - #[case::malformed_field(Error::ResponseField { path: "pages[0].index".into() })] - fn response_failures_are_statusless_runtime_errors(#[case] error: Error) { + #[case::malformed_field(Error::response_field( + "pages[0].index", + "-1".parse::().unwrap_err() + ))] + fn response_failures_are_statusless_runtime_errors( + #[case] error: Error, + #[values("status_code", "ocr_request_format_error", "headers")] attribute: &str, + ) { Python::initialize(); Python::attach(|py| { - let message = error.to_string(); + let report = Report::new(&error).to_string(); let mapped = to_pyerr(error); let value = mapped.value(py); - assert!(mapped.is_instance_of::(py)); + assert!(mapped.is_instance_of::(py)); assert!(!mapped.is_instance_of::(py)); - assert_eq!(value.to_string(), message); - for attribute in ["status_code", "ocr_request_format_error", "headers"] { - assert!(!value.hasattr(attribute).unwrap(), "{attribute}"); - } + assert_eq!(value.to_string(), report); + assert!(!value.hasattr(attribute).unwrap()); }); } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index fc1b70392e2..a1a56baef8f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -12,7 +12,7 @@ use pyo3::{ }; use super::{ - errors::to_pyerr as ocr_error_to_pyerr, + errors::route_failure, project::{OcrHostHandles, SETTINGS_ERROR_MARKER, project_request}, }; use crate::{marshal::public_response, routes::RequestView}; @@ -135,16 +135,11 @@ impl PythonBinding for OcrPythonHost { } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { - if let Error::Secret(source) = &error - && let Some(original) = crate::secrets::python_error(py, source) - { - return Ok(original); - } - Ok(self.map_failure(py, ocr_error_to_pyerr(error))) + route_failure(py, error, |native| Ok(self.map_failure(py, native))) } - fn host_error(error: &PyErr) -> Error { - Error::InvalidRequest(error.to_string()) + fn host_error(_: &PyErr) -> Error { + Error::HostCall } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index fe2f83bff46..40287c669eb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -6,14 +6,16 @@ use litellm_inference_ocr::{ wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, }; use litellm_llms::base_llm::ocr::error::Error; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; use serde_json::{Map, Value}; use super::{document::FileDocumentInput, errors::to_pyerr as ocr_error_to_pyerr}; use crate::{ coercion::FieldSpec, credentials::{self, CallerTokenProvider}, - marshal::{project_optional_fields, python_timeout_seconds, request_input_sources}, + marshal::{ + ArgumentError, project_optional_fields, python_timeout_seconds, request_input_sources, + }, python_settings::{PythonSettings, Snapshot}, }; @@ -30,9 +32,9 @@ struct OcrArguments<'a, 'py> { } impl<'py> OcrArguments<'_, 'py> { - fn lookup(&self, name: &str) -> PyResult> { - litellm_host_python::lookup(self.kwargs, self.bound, name)? - .ok_or_else(|| PyValueError::new_err(format!("missing argument: {name}"))) + fn lookup(&self, name: &'static str) -> PyResult> { + Ok(litellm_host_python::lookup(self.kwargs, self.bound, name)? + .ok_or(ArgumentError::Required(name))?) } fn model(&self) -> PyResult { @@ -92,6 +94,7 @@ impl ProjectedDocument { { ocr_error_to_pyerr(Error::RequestField { path: "document.type".into(), + cause: None, }) } else { error @@ -218,8 +221,10 @@ fn fold_module_globals( pub(super) const SETTINGS_ERROR_MARKER: &str = "native_settings_error"; fn settings_error(py: Python<'_>, error: PyErr) -> PyErr { - error.value(py).setattr(SETTINGS_ERROR_MARKER, true).ok(); - error + match error.value(py).setattr(SETTINGS_ERROR_MARKER, true) { + Ok(()) => error, + Err(failure) => failure, + } } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index 9c430ba7868..b797e334c6a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::decline::Decline; use std::convert::Infallible; use crate::routes::inference::InferenceHost; @@ -118,9 +119,9 @@ impl PythonBinding for ResponsesPythonHost { .get("stream") .is_some_and(|value| value == &serde_json::Value::Bool(true)) { - return Err(InvokeError::Native(Error::Unsupported( + return Err(InvokeError::Native(Error::Declined(Decline::new( "native Python responses streaming", - ))); + )))); } Ok(call) } @@ -141,8 +142,10 @@ impl PythonBinding for ResponsesPythonHost { head: ::StreamHead, ) -> PyResult> { let _ = head; - Err(pyo3::exceptions::PyRuntimeError::new_err( - "native Python Responses streaming is not supported", + Err(crate::errors::route_error_to_pyerr( + Error::invalid_response(litellm_llms::ErrorDetail::Unhandled( + "native Python responses streaming", + )), )) } @@ -152,8 +155,10 @@ impl PythonBinding for ResponsesPythonHost { chunk: ::Chunk, ) -> PyResult> { let _ = chunk; - Err(pyo3::exceptions::PyRuntimeError::new_err( - "native Python Responses streaming is not supported", + Err(crate::errors::route_error_to_pyerr( + Error::invalid_response(litellm_llms::ErrorDetail::Unhandled( + "native Python responses streaming", + )), )) } @@ -161,7 +166,7 @@ impl PythonBinding for ResponsesPythonHost { self.host.error(py, error) } fn host_error(error: &PyErr) -> Error { - Error::InvalidRequest(error.to_string().into()) + crate::routes::inference::host_call_error(error) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs b/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs index 615a838a597..c89185dc21f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs @@ -1,7 +1,9 @@ -use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; use serde_json::Value; +use litellm_http::ClientVariant; +use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; + use crate::{ errors::route_error_to_pyerr, marshal::{marshal_headers, optional_timeout}, @@ -25,10 +27,12 @@ impl ResponsesWebSocketConnection { ) -> PyResult> { let headers = marshal_headers(headers)?; let timeout = optional_timeout(timeout_seconds); + let client = crate::http::host_client(py, ClientVariant::WebSocket)?; crate::execution::run_async_value(py, async move { - let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) - .await - .map_err(route_error_to_pyerr)?; + let inner = + RustResponsesWebSocketConnection::connect_url(&client, &url, &headers, timeout) + .await + .map_err(route_error_to_pyerr)?; Ok(ResponsesWebSocketConnection { inner }) }) } diff --git a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs index 2c26311231f..4b10eed087f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs @@ -2,6 +2,7 @@ use crate::execution::run_async; use std::sync::Arc; use std::{num::NonZero, thread::available_parallelism}; +use litellm_error::{Report, Source}; use litellm_host_python::enter_native; use litellm_token_counter::{ CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, @@ -50,10 +51,10 @@ impl TokenCounter { let _slot = encode_slots .acquire_owned() .await - .map_err(|error| Error::Task(error.to_string()))?; + .map_err(|error| Error::Task(Source::new(error)))?; tokio::task::spawn_blocking(move || count_body(&counter, &body)) .await - .map_err(|error| Error::Task(error.to_string()))? + .map_err(|error| Error::Task(Source::new(error)))? }, token_count_error_to_pyerr, ) @@ -70,19 +71,12 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result PyErr { - let message = error.to_string(); - match error { - Error::Load(_) - | Error::Ranks(_) - | Error::UnicodeClasses - | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message), - Error::RequestParse(_) - | Error::MissingInput - | Error::FloatText - | Error::ContentBlock - | Error::ArrayItems - | Error::JsonSerialization(_) - | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message), - Error::Encode(_) | Error::Decode(_) | Error::Task(_) => PyRuntimeError::new_err(message), + let message = Report::new(&error).to_string(); + if error.is_unsupported() { + RustBridgeDeclined::new_err(message) + } else if error.is_load() { + PyValueError::new_err(message) + } else { + PyRuntimeError::new_err(message) } } diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 71a06bfdc39..669c10e602a 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -192,7 +192,7 @@ handler.get_secret_from_manager = get_secret_from_manager KeyManagementSettings::default(), )), Arc::new(move |_: &str| fallback.map(str::to_owned)), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(FailurePolicy::EnvironmentFallback); (resolver, locals, handler) diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs index d1a1eafd468..7afd06e2373 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -128,9 +128,7 @@ pub(crate) fn read(py: Python<'_>) -> PyResult { .system .is_some_and(|system| system != backend.system()) { - return Err(pyo3::exceptions::PyValueError::new_err( - "native secret manager system does not match configuration", - )); + return Err(super::ManagerError::SystemMismatch.into()); } Ok(SecretManagerSnapshot { client: SecretManagerClient::Native(Box::new(backend)), diff --git a/litellm-rust/crates/python-bridge/src/secrets/error.rs b/litellm-rust/crates/python-bridge/src/secrets/error.rs index 5ccfd0a1554..ca51ee79921 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/error.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/error.rs @@ -1,7 +1,66 @@ use std::fmt; +use litellm_error::Report; use litellm_secrets::Error; -use pyo3::{exceptions::PyBaseException, prelude::*}; +use pyo3::{ + exceptions::{PyBaseException, PyRuntimeError, PyValueError}, + prelude::*, +}; + +/// A secret manager call the bridge rejects, with the text Python's own managers raise. +#[derive(Debug, thiserror::Error)] +pub(crate) enum ManagerError { + #[error("native secret manager must be recreated after fork")] + Forked, + #[error("unknown secret manager system")] + UnknownSystem(#[source] serde_json::Error), + #[error("invalid secret manager settings")] + Settings(#[source] litellm_error::Source), + #[error("native secret manager system does not match configuration")] + SystemMismatch, + #[error("secret manager does not support provider reads")] + ProviderReads, + #[error("invalid timeout")] + Timeout(#[source] std::time::TryFromFloatSecsError), + #[error("No secret found in CyberArk Secret Manager for {0}")] + CyberarkSecretMissing(String), + #[error("Current secret {0} not found")] + CurrentMissing(String), + #[error("Failed to verify new secret {0}")] + ReplacementMissing(String), + #[error("New secret value mismatch")] + ReplacementMismatch, + #[error("native secret manager mutation is unavailable")] + MutationUnsupported, + #[error(transparent)] + Secret(Error), +} + +impl ManagerError { + pub(crate) fn settings(error: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::Settings(litellm_error::Source::new(error)) + } +} + +impl From for PyErr { + fn from(error: ManagerError) -> Self { + let message = Report::new(&error).to_string(); + match error { + ManagerError::Forked => PyRuntimeError::new_err(message), + ManagerError::UnknownSystem(_) + | ManagerError::Settings(_) + | ManagerError::SystemMismatch + | ManagerError::ProviderReads + | ManagerError::Timeout(_) + | ManagerError::CyberarkSecretMissing(_) + | ManagerError::CurrentMissing(_) + | ManagerError::ReplacementMissing(_) + | ManagerError::ReplacementMismatch + | ManagerError::MutationUnsupported + | ManagerError::Secret(_) => PyValueError::new_err(message), + } + } +} #[derive(thiserror::Error)] #[error("Python secret manager failed")] @@ -14,18 +73,23 @@ impl fmt::Debug for PythonSecretError { } pub(super) fn external_error(py: Python<'_>, error: PyErr) -> Error { - Error::ExternalManager(Box::new(PythonSecretError(error.into_value(py)))) + Error::ExternalManager(litellm_error::Source::new(PythonSecretError( + error.into_value(py), + ))) } -pub(super) fn read_error(py: Python<'_>, error: PyErr) -> Error { - Error::ExternalRead(Box::new(PythonSecretError(error.into_value(py)))) +pub(crate) fn read_error(py: Python<'_>, error: PyErr) -> Error { + Error::ExternalRead(litellm_error::Source::new(PythonSecretError( + error.into_value(py), + ))) } -pub(crate) fn python_error(py: Python<'_>, error: &Error) -> Option { - let (Error::ExternalManager(source) | Error::ExternalRead(source)) = error else { - return None; - }; - source - .downcast_ref::() +/// The exception a Python secret manager raised anywhere in the chain, so a route surfaces it +/// unchanged instead of wrapping it. +pub(crate) fn python_error( + py: Python<'_>, + error: &(dyn std::error::Error + 'static), +) -> Option { + litellm_error::find::(error) .map(|error| PyErr::from_value(error.0.clone_ref(py).into_bound(py).into_any())) } diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs index cddd47d08db..57042e6a32e 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mod.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -11,7 +11,9 @@ mod vault; use std::sync::Arc; -pub(crate) use error::python_error; +#[cfg(test)] +pub(crate) use error::read_error; +pub(crate) use error::{ManagerError, python_error}; use litellm_secrets::source::{EnvironmentSecrets, SecretSource}; use pyo3::prelude::*; use python::PythonSecrets; diff --git a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs index eaab06a5b12..b04060651ac 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/mutation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/mutation.rs @@ -1,7 +1,8 @@ use super::operations::{PythonMutationError, PythonMutationResponse}; +use litellm_error::Report; use litellm_host_python::{json_loads, to_py}; use litellm_secrets::cyberark; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; pub(super) fn mutation_value( result: Result, @@ -23,18 +24,7 @@ pub(super) fn mutation_value( &serde_json::json!({"status": "error", "message": message}), ) } - Err(PythonMutationError::CurrentMissing(name)) => Err(PyValueError::new_err(format!( - "Current secret {name} not found" - ))), - Err(PythonMutationError::ReplacementMissing(name)) => Err(PyValueError::new_err(format!( - "Failed to verify new secret {name}" - ))), - Err(PythonMutationError::ReplacementMismatch) => { - Err(PyValueError::new_err("New secret value mismatch")) - } - Err(PythonMutationError::Unsupported) => Err(PyValueError::new_err( - "native secret manager mutation is unavailable", - )), + Err(PythonMutationError::Manager(error)) => Err(error.into()), }) } @@ -43,18 +33,35 @@ fn cyberark_failure( name: &str, failure: cyberark::WriteFailure, ) -> PyResult { - let message = match failure.source { + let message = if failure.source.is_unsafe_secret_name() { + format!("Invalid secret_name {}", name.into_pyobject(py)?.repr()?) + } else { + cyberark_message( + py, + failure.source, + failure.authentication, + failure.request_url.as_ref(), + )? + }; + Ok(if failure.authentication { + format!("Could not authenticate to CyberArk Conjur: {message}") + } else { + message + }) +} + +fn cyberark_message( + py: Python<'_>, + source: cyberark::Error, + authentication: bool, + request_url: Option<&reqwest::Url>, +) -> PyResult { + Ok(match source { cyberark::Error::Status(status) | cyberark::Error::AuthStatus(status) => { - let url = failure - .request_url - .as_ref() - .map_or("", reqwest::Url::as_str); + let url = request_url.map_or("", reqwest::Url::as_str); http_message(py, "POST", url, status)? } - cyberark::Error::Operation(litellm_secrets_types::Error::UnsafeSecretName) => { - format!("Invalid secret_name {}", name.into_pyobject(py)?.repr()?) - } - cyberark::Error::Http(source) if failure.authentication => match os_error_code(&source) { + cyberark::Error::Http(source) if authentication => match os_error_code(&source) { Some(code) => { let reason = py.import("os")?.getattr("strerror")?.call1((code,))?; py.import("builtins")? @@ -63,17 +70,12 @@ fn cyberark_failure( .str()? .extract()? } - None => cyberark::Error::Http(source).to_string(), + None => Report::new(&cyberark::Error::Http(source)).to_string(), }, cyberark::Error::Http(source) if source.is_connect() => { "All connection attempts failed".to_owned() } - source => source.to_string(), - }; - Ok(if failure.authentication { - format!("Could not authenticate to CyberArk Conjur: {message}") - } else { - message + source => Report::new(&source).to_string(), }) } diff --git a/litellm-rust/crates/python-bridge/src/secrets/operations.rs b/litellm-rust/crates/python-bridge/src/secrets/operations.rs index 8e338741aff..b458e3f34e8 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/operations.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/operations.rs @@ -4,6 +4,8 @@ use litellm_secrets::cyberark::AuthenticationRetry; use litellm_secrets::{Error, SecretManager}; use litellm_secrets_types::{PythonSecretRead, SecretOperationContext}; +use super::ManagerError; + pub(super) struct PythonReadRequest { pub secret_name: String, pub primary_secret_name: Option, @@ -62,15 +64,18 @@ pub(super) async fn read_python_provider( #[derive(Debug)] pub(super) enum PythonMutationError { - Unsupported, + Manager(ManagerError), Vault(Box), CyberarkWrite { name: String, failure: Box, }, - CurrentMissing(String), - ReplacementMissing(String), - ReplacementMismatch, +} + +impl From for PythonMutationError { + fn from(error: ManagerError) -> Self { + Self::Manager(error) + } } pub(super) async fn write_python_provider( @@ -89,7 +94,7 @@ pub(super) async fn write_python_provider( })?; Ok(write_success(name)) } - _ => Err(PythonMutationError::Unsupported), + _ => Err(ManagerError::MutationUnsupported.into()), } } @@ -115,7 +120,7 @@ pub(super) async fn delete_python_provider( "message": "CyberArk Conjur does not support direct secret deletion. Use policy updates to remove variables.", })) } - _ => Err(PythonMutationError::Unsupported), + _ => Err(ManagerError::MutationUnsupported.into()), } } @@ -138,7 +143,7 @@ pub(super) async fn rotate_python_provider( .flatten() .is_none() { - return Err(PythonMutationError::CurrentMissing(current_name.to_owned())); + return Err(ManagerError::CurrentMissing(current_name.to_owned()).into()); } client .write_with_retry( @@ -157,16 +162,16 @@ pub(super) async fn rotate_python_provider( .await .ok() .flatten() - .ok_or_else(|| PythonMutationError::ReplacementMissing(new_name.to_owned()))?; + .ok_or_else(|| ManagerError::ReplacementMissing(new_name.to_owned()))?; if actual != *value { - return Err(PythonMutationError::ReplacementMismatch); + return Err(ManagerError::ReplacementMismatch.into()); } if current_name != new_name { client.invalidate_cached_secret(current_name).await; } Ok(write_success(new_name)) } - _ => Err(PythonMutationError::Unsupported), + _ => Err(ManagerError::MutationUnsupported.into()), } } fn write_success(name: &str) -> serde_json::Value { diff --git a/litellm-rust/crates/python-bridge/src/secrets/provider.rs b/litellm-rust/crates/python-bridge/src/secrets/provider.rs index 5f4cd193fcf..ff10f596311 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/provider.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/provider.rs @@ -6,7 +6,9 @@ use litellm_secrets_types::{ AwsOperationContext, CyberarkOperationContext, GoogleOperationContext, HashicorpOperationContext, SecretOperationContext, }; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{prelude::*, types::PyDict}; + +use super::ManagerError; pub(super) fn read_request( system: KeyManagementSystem, @@ -39,9 +41,7 @@ pub(super) fn read_request( SecretOperationContext::Google(GoogleOperationContext::default()) } _ => { - return Err(PyValueError::new_err( - "secret manager does not support provider reads", - )); + return Err(ManagerError::ProviderReads.into()); } }; Ok(PythonReadRequest { @@ -89,9 +89,7 @@ fn aws_context( fn read_timeout(value: Option<&Bound<'_, PyAny>>) -> PyResult> { read_timeout_seconds(value)? - .map(|value| { - Duration::try_from_secs_f64(value).map_err(|_| PyValueError::new_err("invalid timeout")) - }) + .map(|value| Ok(Duration::try_from_secs_f64(value).map_err(ManagerError::Timeout)?)) .transpose() } diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index 5a606ab1039..deb07175e7a 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -84,7 +84,7 @@ mod tests { } async fn resolve(state: Arc, name: &'static str) -> Option { - ResolvedSecrets::from_state(state, litellm_http::Client::plain_for_test()) + ResolvedSecrets::from_state(state, litellm_http::Client::for_test()) .resolve(&[name]) .await .unwrap() @@ -182,7 +182,7 @@ mod tests { .await; let source = ResolvedSecrets::from_state( state(&server, KeyManagementSettings::default()), - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), ); let snapshot = source.resolve(&[declared]).await.unwrap(); assert_eq!(snapshot.get(undeclared), None); @@ -248,7 +248,7 @@ mod tests { async fn oidc_failures_are_not_converted_to_missing_secrets() { let result = ResolvedSecrets::from_state( Arc::new(SecretManagerState::default()), - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), ) .resolve(&["oidc/"]) .await; diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index 72ba91ded16..b956ec297b6 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -7,12 +7,9 @@ use litellm_secrets::{ read_secret_from_python_manager, }; use litellm_secrets_types::PythonSecretRead; -use pyo3::{ - exceptions::{PyAttributeError, PyRuntimeError, PyValueError}, - prelude::*, - types::PyDict, -}; +use pyo3::{exceptions::PyAttributeError, prelude::*, types::PyDict}; +use super::ManagerError; use crate::execution::{run_async_value, run_sync_value}; const SECRET_MANAGER_MODULE: &str = "litellm.rust_bridge.host.secret_manager"; @@ -35,9 +32,7 @@ pub(crate) struct NativeSecretManager { impl NativeSecretManager { pub(super) fn backend(&self) -> PyResult { if self.pid != std::process::id() { - return Err(PyRuntimeError::new_err( - "native secret manager must be recreated after fork", - )); + return Err(ManagerError::Forked.into()); } Ok(self.backend.clone()) } @@ -60,7 +55,7 @@ impl NativeSecretManager { enterprise_enabled, ) .await - .map_err(|error| PyValueError::new_err(error.to_string())) + .map_err(|error| ManagerError::Secret(error).into()) })?; Ok(Self { backend, @@ -82,7 +77,7 @@ impl NativeSecretManager { enterprise_enabled: bool, ) -> PyResult { let system = serde_json::from_value(serde_json::Value::String(system.to_owned())) - .map_err(|_| PyValueError::new_err("unknown secret manager system"))?; + .map_err(ManagerError::UnknownSystem)?; let settings = parse_settings(settings)?; Self::build( py, @@ -146,7 +141,7 @@ impl NativeSecretManager { .collect::>>()?; let settings = from_py::>(&config.getattr("settings")?) - .map_err(|_| PyValueError::new_err("invalid secret manager settings"))?; + .map_err(ManagerError::settings)?; let attributes = config .getattr("settings_attributes")? .extract::>()?; @@ -161,11 +156,11 @@ impl NativeSecretManager { system: serde_json::from_value(serde_json::Value::String( config.getattr("system")?.extract()?, )) - .map_err(|_| PyValueError::new_err("unknown secret manager system"))?, + .map_err(ManagerError::UnknownSystem)?, settings: serde_json::from_value(serde_json::Value::Object( settings.into_iter().chain(setting_overrides).collect(), )) - .map_err(|_| PyValueError::new_err("invalid secret manager settings"))?, + .map_err(ManagerError::settings)?, environment: captured .into_iter() .filter(|(key, _)| !environment_attributes.contains_key(key)) @@ -190,12 +185,8 @@ impl NativeSecretManager { } #[getter] - fn system(&self) -> String { - serde_json::to_value(self.configuration.system) - .expect("serializable system") - .as_str() - .expect("string system") - .to_owned() + fn system(&self) -> &'static str { + self.configuration.system.into() } #[pyo3(signature = (name, settings=None))] @@ -238,7 +229,7 @@ impl NativeSecretManager { run_sync_value(py, async move { super::operations::read_python_provider(&backend, &request, &ProcessEnvironment) .await - .map_err(|error| PyValueError::new_err(error.to_string())) + .map_err(|error| ManagerError::Secret(error).into()) .and_then(|value| python_secret_value(value, &request.secret_name)) }) } @@ -264,7 +255,7 @@ impl NativeSecretManager { run_async_value(py, async move { super::operations::read_python_provider(&backend, &request, &ProcessEnvironment) .await - .map_err(|error| PyValueError::new_err(error.to_string())) + .map_err(|error| ManagerError::Secret(error).into()) .and_then(|value| python_secret_value(value, &request.secret_name)) }) } @@ -438,8 +429,8 @@ fn cached( fn parse_settings(value: Option<&Bound<'_, PyAny>>) -> PyResult { value .map(|value| { - serde_json::from_value(from_py::(value)?) - .map_err(|_| PyValueError::new_err("invalid secret manager settings")) + Ok(serde_json::from_value(from_py::(value)?) + .map_err(ManagerError::settings)?) }) .transpose() .map(Option::unwrap_or_default) @@ -466,11 +457,10 @@ fn python_read_error( name: &str, error: litellm_secrets::Error, ) -> PyErr { - let message = match (system, error) { + match (system, error) { (KeyManagementSystem::Cyberark, litellm_secrets::Error::ManagedSecretMissing) => { - format!("No secret found in CyberArk Secret Manager for {name}") + ManagerError::CyberarkSecretMissing(name.to_owned()).into() } - (_, error) => error.to_string(), - }; - PyValueError::new_err(message) + (_, error) => ManagerError::Secret(error).into(), + } } diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault.rs b/litellm-rust/crates/python-bridge/src/secrets/vault.rs index 79281662848..30e3bd2c413 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault.rs @@ -2,8 +2,11 @@ mod operation; pub(super) use operation::{Failure, FailureKind, FailureStage, delete, rotate, write}; -use litellm_core_utils::exception_mapping_utils::{TimeoutPhase, timeout_message}; -use litellm_secrets::hashicorp::{Error, RawOperationError}; +use litellm_core_utils::exception_mapping_utils::{ + TimeoutPhase, timeout_message as transport_timeout_message, +}; +use litellm_error::Report; +use litellm_secrets::hashicorp::RawOperationError; use pyo3::prelude::*; use super::mutation::{error_value, http_message, json_value}; @@ -16,40 +19,39 @@ pub(super) fn failure_value( failure: Failure, context: &ErrorContext, ) -> PyResult> { - let message = match *failure.kind { - FailureKind::Native(RawOperationError::Http { - method, - url, - status, - body, - }) => match failure.stage { - FailureStage::Current(name) if status == 404 => { - format!("Current secret {name} not found") + let message = match (&*failure.kind, failure.stage) { + (FailureKind::Native(error), stage) if let Some(response) = error.http_failure() => { + match stage { + FailureStage::Current(name) if response.status == 404 => { + format!("Current secret {name} not found") + } + FailureStage::Replacement(name) if response.status == 404 => { + format!("Failed to verify new secret {name}") + } + FailureStage::Current(_) => format!( + "HTTP error occurred while checking current secret: {}", + response_text(py, &response.body)? + ), + FailureStage::Replacement(_) => format!( + "HTTP error occurred while verifying new secret: {}", + response_text(py, &response.body)? + ), + FailureStage::Mutation => { + http_message(py, &response.method, &response.url, response.status)? + } } - FailureStage::Replacement(name) if status == 404 => { - format!("Failed to verify new secret {name}") - } - FailureStage::Current(_) => format!( - "HTTP error occurred while checking current secret: {}", - response_text(py, &body)? - ), - FailureStage::Replacement(_) => format!( - "HTTP error occurred while verifying new secret: {}", - response_text(py, &body)? - ), - FailureStage::Mutation => http_message(py, &method, &url, status)?, - }, - FailureKind::ValueMismatch { expected, actual } => { - let actual = json_value(py, &actual)?; + } + (FailureKind::ValueMismatch { expected, actual }, _) => { + let actual = json_value(py, actual)?; format!( "New secret value mismatch. Expected: {}, Got: {}", expected.expose(), actual.bind(py).str()? ) } - kind => { + (kind, stage) => { let message = cause_message(py, kind, context)?; - match failure.stage { + match stage { FailureStage::Current(_) => { format!("Error checking current secret: {message}") } @@ -63,86 +65,91 @@ pub(super) fn failure_value( error_value(py, message) } -fn cause_message(py: Python<'_>, kind: FailureKind, context: &ErrorContext) -> PyResult { +fn cause_message(py: Python<'_>, kind: &FailureKind, context: &ErrorContext) -> PyResult { Ok(match kind { - FailureKind::Native(RawOperationError::Local(error)) => error.to_string(), + FailureKind::Native(error) => native_message(py, error, context)?, FailureKind::UnsafeName(name) => { format!("Invalid secret_name {}", name.into_pyobject(py)?.repr()?) } - FailureKind::Native(RawOperationError::Timeout { method, elapsed }) => { - if method == "POST" { - let kwargs = pyo3::types::PyDict::new(py); - kwargs.set_item( - "message", - timeout_message( - TimeoutPhase::Read, - context.timeout, - Some(elapsed.as_secs_f64()), - ), - )?; - kwargs.set_item("model", "default-model-name")?; - kwargs.set_item("llm_provider", "litellm-httpx-handler")?; - py.import("litellm.exceptions")? - .getattr("Timeout")? - .call((), Some(&kwargs))? - .str()? - .extract()? - } else if context.aiohttp { - "Timeout on reading data from socket".to_owned() - } else { - String::new() - } - } - FailureKind::Native(RawOperationError::Transport(source)) => { - if let Some(error) = request_error(&source) { - if error.is_timeout() { - String::new() - } else if error.is_connect() { - "All connection attempts failed".to_owned() - } else { - "HashiCorp Vault request failed".to_owned() - } - } else { - "HashiCorp Vault request failed".to_owned() - } - } FailureKind::MissingGet(value) => { - let value = json_value(py, &value)?; + let value = json_value(py, value)?; match value.bind(py).getattr("get") { Err(error) => error.value(py).str()?.extract()?, Ok(_) => "HashiCorp Vault response payload is malformed".to_owned(), } } - FailureKind::Json(body) => match json_value(py, &body) { + FailureKind::Json(body) => match json_value(py, body) { Err(error) => error.value(py).str()?.extract()?, Ok(_) => "HashiCorp Vault response payload is malformed".to_owned(), }, - FailureKind::Native(RawOperationError::Authentication { - source, - url, - certificate, - }) => { - let message = match source { - Error::LoginStatus { status } => http_message(py, "POST", &url, status)?, - error => error.to_string(), - }; - let mechanism = if certificate { "TLS cert" } else { "AppRole" }; - format!("Could not authenticate to Vault via {mechanism}: {message}") - } - FailureKind::Native(RawOperationError::Http { - method, - url, - status, - .. - }) => http_message(py, &method, &url, status)?, FailureKind::ValueMismatch { .. } => "New secret value mismatch".to_owned(), }) } -fn request_error<'a>(error: &'a (dyn std::error::Error + 'static)) -> Option<&'a reqwest::Error> { - error - .downcast_ref::() - .or_else(|| error.source().and_then(request_error)) +fn native_message( + py: Python<'_>, + error: &RawOperationError, + context: &ErrorContext, +) -> PyResult { + if let Some(error) = error.local() { + return Ok(Report::new(error).to_string()); + } + if let Some((method, elapsed)) = error.timeout() { + return timeout_message(py, method, elapsed, context); + } + if let Some(response) = error.http_failure() { + return http_message(py, &response.method, &response.url, response.status); + } + if let Some(login) = error.authentication() { + let message = match login.source.login_status() { + Some(status) => http_message(py, "POST", login.url, status)?, + None => Report::new(login.source).to_string(), + }; + let mechanism = if login.certificate { + "TLS cert" + } else { + "AppRole" + }; + return Ok(format!( + "Could not authenticate to Vault via {mechanism}: {message}" + )); + } + Ok(match error.transport() { + Some(error) if error.status() == Some(408) => String::new(), + Some(error) if error.is_connect() => "All connection attempts failed".to_owned(), + Some(_) | None => "HashiCorp Vault request failed".to_owned(), + }) +} + +fn timeout_message( + py: Python<'_>, + method: &str, + elapsed: std::time::Duration, + context: &ErrorContext, +) -> PyResult { + if method != "POST" { + return Ok(if context.aiohttp { + "Timeout on reading data from socket".to_owned() + } else { + String::new() + }); + } + let kwargs = pyo3::types::PyDict::new(py); + kwargs.set_item( + "message", + transport_timeout_message( + TimeoutPhase::Read, + context.timeout, + Some(elapsed.as_secs_f64()), + ), + )?; + kwargs.set_item("model", "default-model-name")?; + kwargs.set_item("llm_provider", "litellm-httpx-handler")?; + py.import("litellm.exceptions")? + .getattr("Timeout")? + .call((), Some(&kwargs))? + .str()? + .extract() } fn response_text(py: Python<'_>, body: &[u8]) -> PyResult { diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs index cf56beb5efe..41658f0baf8 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use litellm_secrets::{ SecretValue, - hashicorp::{Error, HashicorpVault, RawOperationError}, + hashicorp::{HashicorpVault, RawOperationError}, }; use litellm_secrets_types::{HashicorpOperationContext, SecretWriteContext}; use serde_json::value::RawValue; @@ -53,12 +53,10 @@ impl Failure { } fn native_failure(name: &str, error: RawOperationError) -> Failure { - match error { - RawOperationError::Local(Error::InvalidSecretName(_)) => { - FailureKind::UnsafeName(name.to_owned()).into() - } - error => FailureKind::Native(error).into(), + if error.is_unsafe_secret_name() { + return FailureKind::UnsafeName(name.to_owned()).into(); } + FailureKind::Native(error).into() } pub(crate) async fn write( diff --git a/litellm-rust/crates/python-bridge/src/tokenizer.rs b/litellm-rust/crates/python-bridge/src/tokenizer.rs index df219d55eb6..2026b2fdc42 100644 --- a/litellm-rust/crates/python-bridge/src/tokenizer.rs +++ b/litellm-rust/crates/python-bridge/src/tokenizer.rs @@ -8,16 +8,18 @@ use std::sync::Arc; #[cfg(feature = "fast")] use std::sync::OnceLock; +use litellm_error::Report; use litellm_host_python::{enter_native, release_gil}; #[cfg(feature = "fast")] use litellm_token_counter::fast::{FastCounter, FastTokenizer}; use litellm_token_counter::{Error, TextCodec}; use pyo3::{exceptions::PyUnicodeEncodeError, prelude::*, types::PyString}; -#[cfg(any(feature = "tiktoken", feature = "huggingface"))] -use pyo3::exceptions::PyValueError; #[cfg(feature = "huggingface")] -use pyo3::{exceptions::PyIOError, types::PyDict}; +use pyo3::{ + exceptions::{PyIOError, PyTypeError, PyValueError}, + types::PyDict, +}; #[cfg(feature = "tiktoken")] use pyo3::{ exceptions::{PyKeyError, PyRuntimeError}, @@ -51,6 +53,70 @@ pub(crate) fn load_tiktoken(py: Python<'_>, encoding: &str) -> PyResult for PyErr { + fn from(error: TokenizerError) -> Self { + let message = Report::new(&error).to_string(); + match error { + #[cfg(not(all(feature = "tiktoken", feature = "huggingface")))] + TokenizerError::MissingFeature(_) => RustBridgeDeclined::new_err(message), + #[cfg(all(feature = "tiktoken", feature = "huggingface"))] + TokenizerError::NotTiktoken | TokenizerError::NotHuggingFace => { + PyValueError::new_err(message) + } + #[cfg(feature = "tiktoken")] + TokenizerError::NoVocabulary | TokenizerError::EncodeSpecial(_) => { + PyRuntimeError::new_err(message) + } + #[cfg(feature = "tiktoken")] + TokenizerError::DecodeBytes(_) => PyKeyError::new_err(message), + #[cfg(feature = "huggingface")] + TokenizerError::Direction { .. } | TokenizerError::EncodingJson(_) => { + PyValueError::new_err(message) + } + #[cfg(feature = "huggingface")] + TokenizerError::PretokenizedMismatch => PyTypeError::new_err(message), + #[cfg(feature = "huggingface")] + TokenizerError::Download(_) => PyIOError::new_err(message), + } + } +} + pub(crate) enum Codec { #[cfg(feature = "tiktoken")] Tiktoken(TiktokenTokenizer), @@ -100,9 +166,7 @@ impl Tokenizer { #[cfg(not(feature = "tiktoken"))] { let _ = (py, encoding); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the tiktoken feature", - )) + Err(TokenizerError::MissingFeature("tiktoken").into()) } } @@ -118,9 +182,7 @@ impl Tokenizer { #[cfg(not(feature = "huggingface"))] { let _ = (py, tokenizer_json); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the huggingface feature", - )) + Err(TokenizerError::MissingFeature("huggingface").into()) } } @@ -144,16 +206,14 @@ impl Tokenizer { .getattr("hf_hub_download")? .call((), Some(&kwargs))? .extract()?; - let json = - release_gil(py, || std::fs::read_to_string(path)).map_err(PyIOError::new_err)?; + let json = release_gil(py, || std::fs::read_to_string(path)) + .map_err(TokenizerError::Download)?; Self::from_json(py, &json) } #[cfg(not(feature = "huggingface"))] { let _ = (py, identifier, revision, token); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the huggingface feature", - )) + Err(TokenizerError::MissingFeature("huggingface").into()) } } @@ -199,7 +259,7 @@ impl Tokenizer { let tokenizer = self.tiktoken()?; let text = self.text(text)?; release_gil(py, || tokenizer.encode_special(&text, &allowed)) - .map_err(PyRuntimeError::new_err) + .map_err(|error| TokenizerError::EncodeSpecial(error.into()).into()) } /// tiktoken's `encode_with_unstable`: `(stable_tokens, completions)`. @@ -258,8 +318,8 @@ impl Tokenizer { fn decode_bytes<'py>(&self, py: Python<'py>, ids: Vec) -> PyResult> { enter_native()?; let tokenizer = self.tiktoken()?; - let bytes = - release_gil(py, || tokenizer.decode_bytes(&ids)).map_err(PyKeyError::new_err)?; + let bytes = release_gil(py, || tokenizer.decode_bytes(&ids)) + .map_err(|error| TokenizerError::DecodeBytes(error.into()))?; Ok(PyBytes::new(py, &bytes)) } @@ -480,15 +540,16 @@ impl Tokenizer { match *self.inner { Codec::Tiktoken(ref tokenizer) => Ok(tokenizer), #[cfg(feature = "huggingface")] - Codec::HuggingFace(_) => Err(PyValueError::new_err("requires a tiktoken encoding")), + Codec::HuggingFace(_) => Err(TokenizerError::NotTiktoken.into()), } } #[cfg(feature = "tiktoken")] fn vocabulary(&self) -> PyResult<&Vocabulary> { - self.tiktoken()?.vocabulary().ok_or_else(|| { - PyRuntimeError::new_err("this encoding was built without its vocabulary") - }) + Ok(self + .tiktoken()? + .vocabulary() + .ok_or(TokenizerError::NoVocabulary)?) } #[cfg(feature = "huggingface")] @@ -496,7 +557,7 @@ impl Tokenizer { match *self.inner { Codec::HuggingFace(ref tokenizer) => Ok(tokenizer), #[cfg(feature = "tiktoken")] - Codec::Tiktoken(_) => Err(PyValueError::new_err("requires a Hugging Face tokenizer")), + Codec::Tiktoken(_) => Err(TokenizerError::NotHuggingFace.into()), } } } @@ -530,21 +591,21 @@ impl Sequence { match (self, is_pretokenized) { (Self::Text(text), false) => Ok(text.into()), (Self::Words(words), true) => Ok(words.into()), - _ => Err(pyo3::exceptions::PyTypeError::new_err( - "input must match is_pretokenized", - )), + _ => Err(TokenizerError::PretokenizedMismatch.into()), } } } #[cfg(feature = "huggingface")] -fn direction(value: &str, left: T, right: T, what: &str) -> PyResult { +fn direction(value: &str, left: T, right: T, what: &'static str) -> PyResult { match value { "left" => Ok(left), "right" => Ok(right), - other => Err(PyValueError::new_err(format!( - "invalid {what} direction {other:?}: expected 'left' or 'right'" - ))), + other => Err(TokenizerError::Direction { + what, + value: other.to_owned(), + } + .into()), } } @@ -564,7 +625,7 @@ impl HuggingFaceEncoding { fn new(json: Option<&str>) -> PyResult { let inner = match json { Some(json) => encoding_from_json(json) - .map_err(|error| PyValueError::new_err(error.to_string()))?, + .map_err(|error| TokenizerError::EncodingJson(error.into()))?, None => Encoding::default(), }; Ok(Self { inner }) @@ -586,7 +647,7 @@ impl HuggingFaceEncoding { py: Python<'py>, ) -> PyResult<(Bound<'py, pyo3::types::PyType>, (String,))> { let json = encoding_to_json(&self.inner) - .map_err(|error| PyValueError::new_err(error.to_string()))?; + .map_err(|error| TokenizerError::EncodingJson(error.into()))?; Ok((py.get_type::(), (json,))) } diff --git a/litellm-rust/crates/python-compat/src/error.rs b/litellm-rust/crates/python-compat/src/error.rs index 43990eb7602..6547907056e 100644 --- a/litellm-rust/crates/python-compat/src/error.rs +++ b/litellm-rust/crates/python-compat/src/error.rs @@ -10,15 +10,17 @@ pub enum Error { Unhashable(&'static str), #[error("value nests deeper than {MAX_DEPTH} levels")] TooDeep, - #[error("invalid pickle: {0}")] - InvalidPickle(String), + #[error("invalid pickle")] + InvalidPickle(#[source] serde_pickle::Error), + #[error("could not write pickle")] + PickleWrite(#[source] serde_pickle::Error), #[error("Object of type {0} is not JSON serializable")] NotJsonSerializable(&'static str), #[error("keys must be str, int, float, bool or None, not {0}")] InvalidJsonKey(&'static str), #[error("Out of range float values are not JSON compliant")] NonFiniteFloat, - #[error("integer does not fit in the JSON number range")] + #[error("integer does not fit in 64 bits")] IntegerOutOfRange, #[error("Object of type {0} cannot be pickled as plain data")] NotPicklable(&'static str), diff --git a/litellm-rust/crates/python-compat/src/literal.rs b/litellm-rust/crates/python-compat/src/literal.rs index 12e9036572f..336c6dfcb2e 100644 --- a/litellm-rust/crates/python-compat/src/literal.rs +++ b/litellm-rust/crates/python-compat/src/literal.rs @@ -58,13 +58,6 @@ impl Term { kind: Kind::Other, } } - - fn is_number(&self) -> bool { - matches!( - self.value, - Value::Int(_) | Value::Float(_) | Value::Complex { .. } - ) - } } struct Parser<'a> { @@ -204,18 +197,25 @@ impl Parser<'_> { self.enter()?; let operand = self.unary()?; self.depth -= 1; - if operand.kind != Kind::Constant || !operand.is_number() { + if operand.kind != Kind::Constant { return Err(Error::InvalidLiteral(self.offset + at)); } - let value = if !negative { - operand.value - } else { - match operand.value { - Value::Int(value) => Value::Int(-value), - Value::Float(value) => Value::Float(-value), - Value::Complex { re, im } => Value::Complex { re: -re, im: -im }, - _ => unreachable!("checked numeric above"), - } + let value = match (operand.value, negative) { + (value @ (Value::Int(_) | Value::Float(_) | Value::Complex { .. }), false) => value, + (Value::Int(value), true) => Value::Int(-value), + (Value::Float(value), true) => Value::Float(-value), + (Value::Complex { re, im }, true) => Value::Complex { re: -re, im: -im }, + ( + Value::None + | Value::Bool(_) + | Value::Str(_) + | Value::Bytes(_) + | Value::List(_) + | Value::Tuple(_) + | Value::Dict(_) + | Value::Set(_), + _, + ) => return Err(Error::InvalidLiteral(self.offset + at)), }; Ok(Term { value, diff --git a/litellm-rust/crates/python-compat/src/pickle.rs b/litellm-rust/crates/python-compat/src/pickle.rs index bff7442dc6b..dd7ac6a0e4f 100644 --- a/litellm-rust/crates/python-compat/src/pickle.rs +++ b/litellm-rust/crates/python-compat/src/pickle.rs @@ -22,10 +22,8 @@ pub fn loads(data: &[u8]) -> Result { let mut deserializer = serde_pickle::Deserializer::new(data, DeOptions::new()); let value = Seed { depth: 0 } .deserialize(&mut deserializer) - .map_err(|error| Error::InvalidPickle(error.to_string()))?; - deserializer - .end() - .map_err(|error| Error::InvalidPickle(error.to_string()))?; + .map_err(Error::InvalidPickle)?; + deserializer.end().map_err(Error::InvalidPickle)?; Ok(value) } @@ -33,8 +31,7 @@ pub fn loads(data: &[u8]) -> Result { /// default. Sets and complex numbers are rejected rather than silently changing type. pub fn dumps(value: &Value) -> Result, Error> { check_picklable(value, 0)?; - serde_pickle::to_vec(&Pickled(value), SerOptions::new()) - .map_err(|error| Error::InvalidPickle(error.to_string())) + serde_pickle::to_vec(&Pickled(value), SerOptions::new()).map_err(Error::PickleWrite) } fn check_picklable(value: &Value, depth: usize) -> Result<(), Error> { @@ -66,7 +63,7 @@ impl Serialize for Pickled<'_> { Value::Bool(value) => serializer.serialize_bool(*value), Value::Int(value) => { let value = i64::try_from(value) - .map_err(|_| serde::ser::Error::custom("integer out of i64 range"))?; + .map_err(|_| serde::ser::Error::custom(Error::IntegerOutOfRange))?; serializer.serialize_i64(value) } Value::Float(value) => serializer.serialize_f64(*value), @@ -94,7 +91,7 @@ impl Serialize for Pickled<'_> { map.end() } value @ (Value::Set(_) | Value::Complex { .. }) => Err(serde::ser::Error::custom( - format!("{} cannot be pickled as plain data", value.type_name()), + Error::NotPicklable(value.type_name()), )), } } diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs index 06019cc3c9d..66f4f51bda9 100644 --- a/litellm-rust/crates/secrets-aws/src/auth.rs +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -79,9 +79,7 @@ impl ProvideCredentials for Credentials { self.auth .resolve_credentials(self.config.clone(), &|name| self.environment.get(name)) .await - .map_err(|_| { - CredentialsError::provider_error("secret manager authentication failed") - }) + .map_err(CredentialsError::provider_error) }) } } diff --git a/litellm-rust/crates/secrets-aws/src/error.rs b/litellm-rust/crates/secrets-aws/src/error.rs index c8ca671e9b5..c537da204dd 100644 --- a/litellm-rust/crates/secrets-aws/src/error.rs +++ b/litellm-rust/crates/secrets-aws/src/error.rs @@ -11,7 +11,7 @@ pub enum Error { #[error("KMS response has no plaintext")] MissingPlaintext, #[error("AWS request timed out")] - Timeout, + Timeout(#[source] #[redact] Box>), #[error("AWS KMS decrypt failed")] Decrypt(#[from] #[redact] Box>), #[error("AWS Secrets Manager read failed")] @@ -33,7 +33,7 @@ pub enum Error { #[error("AWS Secrets Manager response has no string payload")] MissingString, #[error("primary secret is not a JSON object")] - PrimarySecret, + PrimarySecret(#[source] #[redact] Option), #[error(transparent)] Operation(#[from] litellm_secrets_types::Error), } diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager/read.rs b/litellm-rust/crates/secrets-aws/src/secret_manager/read.rs index b7583f9785f..d7f10a9847e 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager/read.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager/read.rs @@ -72,7 +72,7 @@ impl AwsSecretsManagerV2 { } let client = self.client_for_context(context)?; let value = match Self::read_with_client(&client, name, ReadPolicy::Python).await { - Err(Error::Read(_) | Error::MissingString | Error::Timeout) => None, + Err(Error::Read(_) | Error::MissingString | Error::Timeout(_)) => None, result => result?, }; Ok(PythonSecretRead::Value(value.map(Secret::String))) @@ -124,9 +124,10 @@ impl AwsSecretsManagerV2 { Self::read_with_client(&self.client, name, policy).await, policy, ) { - (Err(Error::Read(_) | Error::MissingString | Error::Timeout), ReadPolicy::Python) => { - Ok(None) - } + ( + Err(Error::Read(_) | Error::MissingString | Error::Timeout(_)), + ReadPolicy::Python, + ) => Ok(None), (result, _) => result, } } @@ -173,7 +174,7 @@ impl AwsSecretsManagerV2 { aws_sdk_secretsmanager::error::SdkError::TimeoutError(_) ) || matches!(&error, aws_sdk_secretsmanager::error::SdkError::DispatchFailure(failure) if failure.is_timeout()) => { - Err(Error::Timeout) + Err(Error::Timeout(Box::new(error))) } Err(error) if error @@ -216,9 +217,9 @@ fn resolve_payload(payload: PythonSecretRead, name: &str) -> Result Ok(value), PythonSecretRead::PrimaryJson(document) => { - let object: Value = - serde_json::from_str(document.expose()).map_err(|_| Error::PrimarySecret)?; - let object = object.as_object().ok_or(Error::PrimarySecret)?; + let object: Value = serde_json::from_str(document.expose()) + .map_err(|error| Error::PrimarySecret(Some(error)))?; + let object = object.as_object().ok_or(Error::PrimarySecret(None))?; Ok(object.get(name).cloned().map(Secret::from_json)) } } diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager/reads.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager/reads.rs index 3c39924f09f..30aec467543 100644 --- a/litellm-rust/crates/secrets-aws/tests/secret_manager/reads.rs +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager/reads.rs @@ -88,7 +88,7 @@ async fn failed_read_returns_none_but_invalid_primary_json_is_an_error( manager .read_secret_for_resolver("KEY", Some("invalid"), &|_: &str| None) .await, - Err(Error::PrimarySecret) + Err(Error::PrimarySecret(_)) )); assert!(matches!( manager.async_read_secret("no-string").await, @@ -143,7 +143,7 @@ async fn trait_read_applies_the_aws_operation_timeout() { }; assert!(matches!( BaseSecretManager::async_read_secret(&loaded_manager(&server), "key", &context).await, - Err(Error::Timeout) + Err(Error::Timeout(_)) )); } @@ -169,7 +169,7 @@ async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() { let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default()); assert!(matches!( manager.async_read_secret("key").await, - Err(Error::Timeout) + Err(Error::Timeout(_)) )); } diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml index 7ec03fb98da..2aff6129c37 100644 --- a/litellm-rust/crates/secrets-azure/Cargo.toml +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -15,6 +15,7 @@ litellm-core-utils.workspace = true reqwest.workspace = true serde.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true percent-encoding = "2.3" diff --git a/litellm-rust/crates/secrets-azure/src/error.rs b/litellm-rust/crates/secrets-azure/src/error.rs index 2862f08a56c..097f42cb89c 100644 --- a/litellm-rust/crates/secrets-azure/src/error.rs +++ b/litellm-rust/crates/secrets-azure/src/error.rs @@ -7,7 +7,7 @@ pub enum Error { #[error("{0} environment variable is missing")] MissingEnvironment(&'static str), #[error("AZURE_KEY_VAULT_URI is not a valid https vault URL")] - VaultUri, + VaultUri(#[source] Option), #[error("Azure Key Vault credentials are not configured")] MissingCredentials, #[error(transparent)] @@ -20,10 +20,26 @@ pub enum Error { Http( #[source] #[redact] - reqwest::Error, + litellm_http::Error, ), #[error("Azure Key Vault returned HTTP {0}")] Status(u16), #[error("Azure Key Vault response is missing the secret value")] MissingValue, } + +impl Error { + pub fn is_missing_value(&self) -> bool { + match self { + Self::MissingValue => true, + Self::Operation(_) + | Self::Timeout + | Self::MissingEnvironment(_) + | Self::VaultUri(_) + | Self::MissingCredentials + | Self::Auth(_) + | Self::Http(_) + | Self::Status(_) => false, + } + } +} diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs index 095c451927c..d97b7a22cfc 100644 --- a/litellm-rust/crates/secrets-azure/src/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -38,7 +38,7 @@ impl AzureKeyVault { environment: Arc, ) -> Result { if vault.host_str().is_none() { - return Err(Error::VaultUri); + return Err(Error::VaultUri(None)); } let inputs = AzureAuthInputs { azure_scope: ConfigValue::Value(Sourced::new( @@ -64,9 +64,9 @@ impl AzureKeyVault { let value = environment .get(AZURE_KEY_VAULT_URI) .ok_or(Error::MissingEnvironment(AZURE_KEY_VAULT_URI))?; - let vault = reqwest::Url::parse(&value).map_err(|_| Error::VaultUri)?; + let vault = reqwest::Url::parse(&value).map_err(|error| Error::VaultUri(Some(error)))?; if vault.scheme() != "https" || vault.host_str().is_none() { - return Err(Error::VaultUri); + return Err(Error::VaultUri(None)); } Self::with_client(client, vault, environment) } @@ -95,7 +95,7 @@ impl AzureKeyVault { let url = self .vault .join(&format!("secrets/{encoded_name}?api-version={API_VERSION}")) - .map_err(|_| Error::VaultUri)?; + .map_err(|error| Error::VaultUri(Some(error)))?; let response = self .client .get(url) @@ -103,14 +103,17 @@ impl AzureKeyVault { .header(reqwest::header::ACCEPT, "application/json") .send() .await - .map_err(Error::Http)?; + .map_err(|error| Error::Http(error.into()))?; if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(None); } if response.status() != reqwest::StatusCode::OK { return Err(Error::Status(response.status().as_u16())); } - let payload: SecretResponse = response.json().await.map_err(Error::Http)?; + let payload: SecretResponse = response + .json() + .await + .map_err(|error| Error::Http(error.into()))?; let value = payload.value.ok_or(Error::MissingValue)?; Ok(Some(SecretValue::new(value))) } diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs index fcbc46092e1..a0ba9e6ddf8 100644 --- a/litellm-rust/crates/secrets-azure/tests/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -11,7 +11,7 @@ use wiremock::{ fn manager(server: &MockServer) -> AzureKeyVault { AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), ) @@ -122,16 +122,17 @@ async fn missing_value_is_an_error() { } #[rstest] -#[case::missing(None, true)] -#[case::http(Some("http://vault.example"), false)] -#[case::relative(Some("vault.example"), false)] -#[case::malformed(Some("://"), false)] +#[case::missing(None, true, false)] +#[case::http(Some("http://vault.example"), false, false)] +#[case::relative(Some("vault.example"), false, true)] +#[case::malformed(Some("://"), false, true)] fn new_validates_vault_environment( #[case] uri: Option<&'static str>, #[case] missing_environment: bool, + #[case] unparsable: bool, ) { let result = AzureKeyVault::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), Arc::new(move |name: &str| { (name == "AZURE_KEY_VAULT_URI") .then(|| uri.map(str::to_owned)) @@ -145,7 +146,14 @@ fn new_validates_vault_environment( Err(Error::MissingEnvironment("AZURE_KEY_VAULT_URI")) )); } else { - assert!(matches!(result, Err(Error::VaultUri))); + let error = result.err().expect("invalid vault URI"); + assert!(matches!(error, Error::VaultUri(_))); + assert_eq!( + std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .is_some(), + unparsable + ); } } @@ -158,7 +166,7 @@ fn new_validates_vault_environment( #[case::local("http://localhost:8080", "https://localhost/.default")] fn derives_scope_from_vault_host(#[case] uri: &str, #[case] expected: &str) { let manager = AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), uri.parse().unwrap(), Arc::new(|_: &str| None), ) @@ -187,7 +195,7 @@ async fn missing_credentials_do_not_request_vault() { fn manager_without_credentials(server: &MockServer) -> AzureKeyVault { AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| { (name == "AZURE_CREDENTIAL").then(|| "ClientSecretCredential".to_owned()) diff --git a/litellm-rust/crates/secrets-azure/tests/live.rs b/litellm-rust/crates/secrets-azure/tests/live.rs index 429cd3013f7..d7713dc798d 100644 --- a/litellm-rust/crates/secrets-azure/tests/live.rs +++ b/litellm-rust/crates/secrets-azure/tests/live.rs @@ -10,7 +10,7 @@ use rstest::rstest; #[ignore] async fn reads_a_real_secret() { let environment = Arc::new(ProcessEnvironment); - let manager = AzureKeyVault::new(litellm_http::Client::plain_for_test(), environment).unwrap(); + let manager = AzureKeyVault::new(litellm_http::Client::for_test(), environment).unwrap(); let name = std::env::var("AZURE_KEY_VAULT_LIVE_SECRET_NAME").unwrap(); let secret = manager.get_secret(&name).await.unwrap().unwrap(); assert!(matches!(&secret, Secret::String(_))); diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml index f630d5857d8..2ae2c1ea5c3 100644 --- a/litellm-rust/crates/secrets-cyberark/Cargo.toml +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -14,6 +14,7 @@ moka.workspace = true reqwest.workspace = true serde_json.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true litellm-tracing.workspace = true percent-encoding = "2.3" diff --git a/litellm-rust/crates/secrets-cyberark/src/error.rs b/litellm-rust/crates/secrets-cyberark/src/error.rs index 4f21647225b..b5dbd7342fa 100644 --- a/litellm-rust/crates/secrets-cyberark/src/error.rs +++ b/litellm-rust/crates/secrets-cyberark/src/error.rs @@ -6,7 +6,7 @@ pub enum Error { Http( #[from] #[redact] - reqwest::Error, + litellm_http::Error, ), #[error("CyberArk Conjur authentication returned HTTP {0}")] AuthStatus(u16), @@ -17,15 +17,68 @@ pub enum Error { )] MissingCredentials, #[error("CyberArk client certificate could not be loaded")] - ClientCertificate, + ClientCertificate( + #[source] + #[redact] + litellm_http::Error, + ), #[error("CyberArk Conjur HTTP client could not be built")] - Client(#[redact] Box), + Client( + #[source] + #[redact] + litellm_http::Error, + ), #[error("invalid refresh interval")] - RefreshInterval, + RefreshInterval(#[source] std::num::ParseIntError), #[error("invalid CyberArk Conjur endpoint")] - Endpoint, + Endpoint(#[source] url::ParseError), #[error("CyberArk secret manager requires an enterprise license")] EnterpriseRequired, #[error(transparent)] Operation(#[from] litellm_secrets_types::Error), } + +impl Error { + pub fn is_unsafe_secret_name(&self) -> bool { + match self { + Self::Operation(error) => error.is_unsafe_secret_name(), + Self::Timeout + | Self::Http(_) + | Self::AuthStatus(_) + | Self::Status(_) + | Self::MissingCredentials + | Self::ClientCertificate(_) + | Self::Client(_) + | Self::RefreshInterval(_) + | Self::Endpoint(_) + | Self::EnterpriseRequired => false, + } + } +} + +#[derive(thiserror::Error, veil::Redact)] +#[error("CyberArk Conjur write failed")] +pub struct WriteFailure { + pub source: Error, + #[redact] + pub request_url: Option, + pub authentication: bool, +} + +impl WriteFailure { + pub(crate) fn local(source: Error) -> Self { + Self { + source, + request_url: None, + authentication: false, + } + } + + pub(crate) fn request(source: Error, url: reqwest::Url) -> Self { + Self { + source, + request_url: Some(url), + authentication: false, + } + } +} diff --git a/litellm-rust/crates/secrets-cyberark/src/lib.rs b/litellm-rust/crates/secrets-cyberark/src/lib.rs index 74a8b6febf2..d780fb43921 100644 --- a/litellm-rust/crates/secrets-cyberark/src/lib.rs +++ b/litellm-rust/crates/secrets-cyberark/src/lib.rs @@ -3,5 +3,5 @@ mod error; mod secret_manager; -pub use error::Error; -pub use secret_manager::{AuthenticationRetry, CyberArkSecretManager, DeleteOutcome, WriteFailure}; +pub use error::{Error, WriteFailure}; +pub use secret_manager::{AuthenticationRetry, CyberArkSecretManager, DeleteOutcome}; diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index 44473c9fe90..c5cb0308b13 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -17,7 +17,7 @@ use litellm_secrets_types::{ use moka::future::Cache; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode}; -use crate::Error; +use crate::{Error, WriteFailure}; const CYBERARK_API_BASE: &str = "CYBERARK_API_BASE"; const CYBERARK_ACCOUNT: &str = "CYBERARK_ACCOUNT"; @@ -62,38 +62,12 @@ pub enum AuthenticationRetry { Unauthorized, } -#[derive(veil::Redact)] -pub struct WriteFailure { - pub source: Error, - #[redact] - pub request_url: Option, - pub authentication: bool, -} - -impl WriteFailure { - fn local(source: Error) -> Self { - Self { - source, - request_url: None, - authentication: false, - } - } - - fn request(source: Error, url: reqwest::Url) -> Self { - Self { - source, - request_url: Some(url), - authentication: false, - } - } -} - impl CyberArkSecretManager { fn secret_url(&self, name: &str) -> Result { let encoded = utf8_percent_encode(name, SECRET_NAME_SAFE); self.endpoint .join(&format!("secrets/{}/variable/{}", self.account, encoded)) - .map_err(|_| Error::Endpoint) + .map_err(Error::Endpoint) } } diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs index ea2262392d5..1cda0ae5306 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/client.rs @@ -64,25 +64,18 @@ impl CyberArkSecretManager { }), ..config.clone() }; - let client = - pool.client(&config, ClientVariant::Provider) - .map_err(|error| match error { - litellm_http::Error::Read { - tls_source: TlsSource::ClientIdentity, - .. - } - | litellm_http::Error::InvalidPem { - tls_source: TlsSource::ClientIdentity, - .. - } => Error::ClientCertificate, - other => Error::Client(Box::new(other)), - })?; + let client = pool + .client(&config, ClientVariant::Provider) + .map_err(|error| match error.tls_source() { + Some(TlsSource::ClientIdentity) => Error::ClientCertificate(error), + Some(TlsSource::CaBundle) | None => Error::Client(error), + })?; let endpoint = reqwest::Url::parse( &environment .get(CYBERARK_API_BASE) .unwrap_or_else(|| DEFAULT_API_BASE.to_owned()), ) - .map_err(|_| Error::Endpoint)?; + .map_err(Error::Endpoint)?; let account = environment .get(CYBERARK_ACCOUNT) .unwrap_or_else(|| DEFAULT_ACCOUNT.to_owned()); @@ -95,7 +88,7 @@ impl CyberArkSecretManager { value .parse::() .map(Duration::from_secs) - .map_err(|_| Error::RefreshInterval) + .map_err(Error::RefreshInterval) }) .transpose()?; Ok(Self::with_client( @@ -115,7 +108,7 @@ impl CyberArkSecretManager { self.account, utf8_percent_encode(&self.username, SECRET_NAME_SAFE) )) - .map_err(|_| Error::Endpoint) + .map_err(Error::Endpoint) } pub(super) async fn authenticate( @@ -135,11 +128,14 @@ impl CyberArkSecretManager { context, ) .send() - .await?; + .await + .map_err(litellm_http::Error::from)?; if !response.status().is_success() { return Err(Error::AuthStatus(response.status().as_u16())); } - let token = SecretValue::new(STANDARD.encode(response.text().await?)); + let token = SecretValue::new( + STANDARD.encode(response.text().await.map_err(litellm_http::Error::from)?), + ); self.token.insert((), token.clone()).await; Ok(token) } diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/read.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/read.rs index daa9ed1114d..2862c6a956d 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/read.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/read.rs @@ -76,7 +76,8 @@ impl CyberArkSecretManager { context, ) .send() - .await?; + .await + .map_err(litellm_http::Error::from)?; let response = if matches!(retry, AuthenticationRetry::Unauthorized) && had_cached_token && response.status() == reqwest::StatusCode::UNAUTHORIZED @@ -89,7 +90,8 @@ impl CyberArkSecretManager { context, ) .send() - .await? + .await + .map_err(litellm_http::Error::from)? } else { response }; @@ -99,7 +101,7 @@ impl CyberArkSecretManager { if !response.status().is_success() { return Err(Error::Status(response.status().as_u16())); } - let value = SecretValue::new(response.text().await?); + let value = SecretValue::new(response.text().await.map_err(litellm_http::Error::from)?); Ok(Some(value)) } } diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs index 87f6cea3619..7a2c79f506a 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager/write.rs @@ -83,7 +83,9 @@ impl CyberArkSecretManager { ) .send() .await - .map_err(|source| WriteFailure::request(source.into(), url.clone())) + .map_err(|source| { + WriteFailure::request(litellm_http::Error::from(source).into(), url.clone()) + }) } pub(super) async fn ensure_variable_exists( @@ -141,8 +143,8 @@ impl CyberArkSecretManager { ); return; } - Err(error) => { - litellm_tracing::warn!("Error ensuring CyberArk variable exists: {error}"); + Err(_) => { + litellm_tracing::warn!("Error ensuring CyberArk variable exists"); return; } } diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs index 1bea77e49ae..8a0746f8002 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/configuration.rs @@ -77,7 +77,7 @@ async fn authentication_encodes_login(#[case] username: &str, #[case] expected_p .mount(&server) .await; let manager = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "acct".into(), username.into(), @@ -201,7 +201,7 @@ async fn trait_read_applies_cyberark_operation_timeout_to_authentication() { match result { Err(Error::Timeout) => {} - Err(Error::Http(error)) => assert!(error.is_timeout()), + Err(Error::Http(error)) => assert_eq!(error.status(), Some(408)), other => panic!("expected timeout, got {other:?}"), } } @@ -237,7 +237,7 @@ fn new_validates_credentials_before_license_and_configuration() { }), true ), - Err(Error::RefreshInterval) + Err(Error::RefreshInterval(_)) )); assert!(matches!( from_environment( @@ -248,7 +248,7 @@ fn new_validates_credentials_before_license_and_configuration() { }), true ), - Err(Error::Endpoint) + Err(Error::Endpoint(_)) )); } @@ -263,7 +263,7 @@ fn certificate_only_credentials_are_validated_as_a_client_identity() { true, ); - assert!(matches!(result, Err(Error::ClientCertificate))); + assert!(matches!(result, Err(Error::ClientCertificate(_)))); } #[rstest] @@ -347,7 +347,7 @@ fn invalid_client_identity_is_not_ignored( true, ); - assert!(matches!(result, Err(Error::ClientCertificate))); + assert!(matches!(result, Err(Error::ClientCertificate(_)))); } #[rstest] @@ -413,7 +413,7 @@ fn new_reports_missing_client_certificate_files() { }), true ), - Err(Error::ClientCertificate) + Err(Error::ClientCertificate(_)) )); } @@ -433,7 +433,7 @@ async fn trailing_slash_endpoint_preserves_base_path() { .await; let endpoint = format!("{}/prefix/", server.uri()).parse().unwrap(); let manager = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), endpoint, "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs index 6fa15ecba9b..d7bc71b6503 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/support.rs @@ -62,7 +62,7 @@ pub(super) fn from_environment( pub(super) fn manager(server: &MockServer, ttl: Duration) -> CyberArkSecretManager { CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs index 764a591bc71..3db8b6355e1 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager/writes.rs @@ -233,7 +233,7 @@ async fn writes_match_python_parity_fixture(parity_fixture: ParityFixture) { .mount(&server) .await; let manager = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), parity_fixture.account, parity_fixture.username, @@ -297,7 +297,7 @@ async fn live_conjur_round_trip() { .as_nanos() ); let manager = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), endpoint.clone(), account.clone(), username.clone(), @@ -312,7 +312,7 @@ async fn live_conjur_round_trip() { .await .unwrap(); let verifier = CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), endpoint.clone(), account.clone(), username.clone(), diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml index 3ce14fe7a12..06b3747def5 100644 --- a/litellm-rust/crates/secrets-google/Cargo.toml +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -17,6 +17,7 @@ base64.workspace = true crc32c = "0.6.8" serde_json.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true google-cloud-kms-v1 = "1.14.0" google-cloud-gax = { version = "1.14.0", default-features = false } diff --git a/litellm-rust/crates/secrets-google/src/error.rs b/litellm-rust/crates/secrets-google/src/error.rs index 159e4983ddc..486c7a3f35e 100644 --- a/litellm-rust/crates/secrets-google/src/error.rs +++ b/litellm-rust/crates/secrets-google/src/error.rs @@ -26,7 +26,7 @@ pub enum Error { Http( #[from] #[redact] - reqwest::Error, + litellm_http::Error, ), #[error("Google Secret Manager returned HTTP {0}")] Status(u16), @@ -35,15 +35,41 @@ pub enum Error { #[error("required environment variable is missing: {0}")] MissingEnvironment(&'static str), #[error("invalid refresh interval")] - RefreshInterval, + RefreshInterval(#[source] std::num::ParseIntError), #[error("payload is not valid base64")] - Base64(#[from] base64::DecodeError), + Base64( + #[from] + #[redact] + base64::DecodeError, + ), #[error("Google Secret Manager payload checksum mismatch")] - Checksum, + Checksum(#[source] Option), #[error("decrypted value is not UTF-8")] - Utf8, + Utf8(#[source] std::str::Utf8Error), #[error("invalid Google Secret Manager endpoint")] - Endpoint, + Endpoint(#[source] url::ParseError), #[error("Google Secret Manager requires an enterprise license")] EnterpriseRequired, } + +impl Error { + pub fn is_not_found(&self) -> bool { + match self { + Self::Status(status) => *status == 404, + Self::Operation(_) + | Self::Timeout + | Self::Client(_) + | Self::Auth(_) + | Self::Kms(_) + | Self::Http(_) + | Self::MissingPayload + | Self::MissingEnvironment(_) + | Self::RefreshInterval(_) + | Self::Base64(_) + | Self::Checksum(_) + | Self::Utf8(_) + | Self::Endpoint(_) + | Self::EnterpriseRequired => false, + } + } +} diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs index 08ba466b799..463d1f6fd73 100644 --- a/litellm-rust/crates/secrets-google/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -92,12 +92,12 @@ impl GoogleSecretManager { let ttl = environment .get(GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL) .filter(|v| !v.is_empty()) - .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .map(|v| v.parse::().map_err(Error::RefreshInterval)) .transpose()? .unwrap_or( environment .get(SECRET_MANAGER_REFRESH_INTERVAL) - .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .map(|v| v.parse::().map_err(Error::RefreshInterval)) .transpose()? .unwrap_or(DEFAULT_REFRESH_INTERVAL.as_secs() as i64), ); @@ -164,20 +164,21 @@ impl GoogleSecretManager { ), percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC) )) - .map_err(|_| Error::Endpoint)?; + .map_err(Error::Endpoint)?; let response = self .client .get(url) .headers(self.credentials.request_headers().await?) .send() - .await?; + .await + .map_err(litellm_http::Error::from)?; if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(None); } if response.status() != reqwest::StatusCode::OK { return Err(Error::Status(response.status().as_u16())); } - let response: Response = response.json().await?; + let response: Response = response.json().await.map_err(litellm_http::Error::from)?; let Some(payload) = response.payload else { return Err(Error::MissingPayload); }; @@ -186,12 +187,15 @@ impl GoogleSecretManager { }; let bytes = STANDARD.decode(data)?; if let Some(expected) = payload.data_crc32c { - let expected = expected.parse::().map_err(|_| Error::Checksum)?; + let expected = expected + .parse::() + .map_err(|error| Error::Checksum(Some(error)))?; if crc32c::crc32c(&bytes) != expected { - return Err(Error::Checksum); + return Err(Error::Checksum(None)); } } - let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?; + let plaintext = + String::from_utf8(bytes).map_err(|error| Error::Utf8(error.utf8_error()))?; let value = SecretValue::new(plaintext); Ok(Some(value)) } diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs index e9bee7633f5..615e331cb96 100644 --- a/litellm-rust/crates/secrets-google/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -11,7 +11,7 @@ use wiremock::{ fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager { GoogleSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "project".into(), Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())), @@ -152,8 +152,8 @@ async fn failed_or_missing_reads_are_not_cached( assert!(matches!(result, Err(Error::MissingPayload))); } ExpectedReadFailure::Base64 => assert!(matches!(result, Err(Error::Base64(_)))), - ExpectedReadFailure::Utf8 => assert!(matches!(result, Err(Error::Utf8))), - ExpectedReadFailure::Checksum => assert!(matches!(result, Err(Error::Checksum))), + ExpectedReadFailure::Utf8 => assert!(matches!(result, Err(Error::Utf8(_)))), + ExpectedReadFailure::Checksum => assert!(matches!(result, Err(Error::Checksum(None)))), } drop(failing); Mock::given(path( @@ -215,7 +215,7 @@ async fn always_read_and_expired_cache_fetch_again( fn google_manager_requires_host_license_and_project_configuration() { assert!(matches!( GoogleSecretManager::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), Arc::new(|_: &str| None), false ), @@ -223,7 +223,7 @@ fn google_manager_requires_host_license_and_project_configuration() { )); assert!(matches!( GoogleSecretManager::new( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), Arc::new(|_: &str| None), true ), @@ -244,8 +244,8 @@ fn google_manager_rejects_invalid_refresh_intervals(#[case] variable: &'static s }); assert!(matches!( - GoogleSecretManager::new(litellm_http::Client::plain_for_test(), environment, true), - Err(Error::RefreshInterval) + GoogleSecretManager::new(litellm_http::Client::for_test(), environment, true), + Err(Error::RefreshInterval(_)) )); } diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml index 7dd3d3c674f..8731a8deb7a 100644 --- a/litellm-rust/crates/secrets-hashicorp/Cargo.toml +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -6,8 +6,11 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-error.workspace = true litellm-core-utils.workspace = true +litellm-http.workspace = true litellm-secrets-types.workspace = true +reqwest.workspace = true rustify.workspace = true rustify_derive.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs index d32491199c4..4232667ae16 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/config.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -153,9 +153,11 @@ fn refresh_interval(environment: &dyn Lookup) -> Result { let Some(value) = value else { return Ok(DEFAULT_REFRESH_INTERVAL); }; - let seconds: i64 = value.parse().map_err(|_| Error::RefreshInterval)?; + let seconds: i64 = value + .parse() + .map_err(|error| Error::RefreshInterval(Some(error)))?; if seconds < 0 { - return Err(Error::RefreshInterval); + return Err(Error::RefreshInterval(None)); } Ok(Duration::from_secs(seconds as u64)) } diff --git a/litellm-rust/crates/secrets-hashicorp/src/error.rs b/litellm-rust/crates/secrets-hashicorp/src/error.rs index 1813cb486bc..73399a44eaa 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/error.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/error.rs @@ -1,42 +1,144 @@ +use rustify::errors::ClientError as RustifyClientError; +use vaultrs::error::ClientError; + #[derive(thiserror::Error, veil::Redact)] pub enum Error { #[error("HashiCorp Vault requires an enterprise license")] EnterpriseRequired, #[error("invalid secret name")] - InvalidSecretName(litellm_secrets_types::Error), + InvalidSecretName(#[source] litellm_secrets_types::Error), #[error(transparent)] Operation(#[from] litellm_secrets_types::Error), #[error("HashiCorp Vault client failed")] Client( - #[from] + #[source] #[redact] vaultrs::error::ClientError, ), - #[error("HashiCorp Vault client settings are invalid: {message}")] - ClientSettings { message: String }, - #[error("HashiCorp Vault TLS identity could not be configured for {path}: {message}")] + #[error("HashiCorp Vault request failed")] + Transport( + #[source] + #[redact] + litellm_http::Error, + ), + #[error("HashiCorp Vault client settings are invalid")] + ClientSettings(#[source] vaultrs::client::VaultClientSettingsBuilderError), + #[error("HashiCorp Vault TLS identity could not be configured for {path}")] TlsIdentity { path: std::path::PathBuf, - message: String, + source: litellm_error::Source, }, #[error("HashiCorp Vault login returned HTTP {status}")] LoginStatus { status: u16 }, #[error("HashiCorp Vault login response is malformed")] - MalformedLogin, + MalformedLogin( + #[source] + #[redact] + vaultrs::error::ClientError, + ), #[error("HashiCorp Vault authentication is not configured")] NoAuthConfigured, #[error("HashiCorp Vault returned HTTP {status}")] Status { status: u16 }, #[error("HashiCorp Vault response payload is malformed")] - MalformedPayload, + MalformedPayload( + #[source] + #[redact] + vaultrs::error::ClientError, + ), #[error("HashiCorp Vault secret value is not a string")] NonStringValue, #[error("HashiCorp Vault data key conflicts with description")] DataKeyConflictsWithDescription, #[error("HashiCorp Vault secret version exceeds CAS range")] - CasVersionOverflow, + CasVersionOverflow(#[source] std::num::TryFromIntError), #[error("HashiCorp Vault operation timed out")] Timeout, #[error("invalid HashiCorp Vault refresh interval")] - RefreshInterval, + RefreshInterval(#[source] Option), +} + +impl Error { + pub fn login_status(&self) -> Option { + match self { + Self::LoginStatus { status } => Some(*status), + Self::EnterpriseRequired + | Self::InvalidSecretName(_) + | Self::Operation(_) + | Self::Client(_) + | Self::Transport(_) + | Self::ClientSettings(_) + | Self::TlsIdentity { .. } + | Self::MalformedLogin(_) + | Self::NoAuthConfigured + | Self::Status { .. } + | Self::MalformedPayload(_) + | Self::NonStringValue + | Self::DataKeyConflictsWithDescription + | Self::CasVersionOverflow(_) + | Self::Timeout + | Self::RefreshInterval(_) => None, + } + } + + pub fn is_unsafe_secret_name(&self) -> bool { + match self { + Self::InvalidSecretName(error) | Self::Operation(error) => { + error.is_unsafe_secret_name() + } + Self::EnterpriseRequired + | Self::Client(_) + | Self::Transport(_) + | Self::ClientSettings(_) + | Self::TlsIdentity { .. } + | Self::LoginStatus { .. } + | Self::MalformedLogin(_) + | Self::NoAuthConfigured + | Self::Status { .. } + | Self::MalformedPayload(_) + | Self::NonStringValue + | Self::DataKeyConflictsWithDescription + | Self::CasVersionOverflow(_) + | Self::Timeout + | Self::RefreshInterval(_) => false, + } + } +} + +impl From for Error { + fn from(error: ClientError) -> Self { + match error { + ClientError::RestClientError { source } => match transport(source) { + Ok(error) => Self::Transport(error), + Err(source) => Self::Client(ClientError::RestClientError { source }), + }, + ClientError::RestClientBuildError { source } => Self::Transport(source.into()), + error => Self::Client(error), + } + } +} + +pub(crate) fn transport( + error: RustifyClientError, +) -> Result { + match error { + RustifyClientError::RequestError { + source, + url, + method, + } => source + .downcast::() + .map(litellm_http::Error::from) + .map_err(|source| RustifyClientError::RequestError { + source, + url, + method, + }), + RustifyClientError::ResponseError { source } => source + .downcast::() + .map(litellm_http::Error::from) + .map_err(|source| RustifyClientError::ResponseError { source }), + RustifyClientError::ReqwestBuildError { source } => Ok(source.into()), + error => Err(error), + } } diff --git a/litellm-rust/crates/secrets-hashicorp/src/lib.rs b/litellm-rust/crates/secrets-hashicorp/src/lib.rs index 7aab5a38b01..15c03724f72 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/lib.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/lib.rs @@ -7,4 +7,6 @@ pub mod secret_manager; pub use config::{AppRoleAuth, HashicorpVaultConfig, TlsCertAuth}; pub use error::Error; -pub use secret_manager::{HashicorpVault, RawOperationError, SecretLocation}; +pub use secret_manager::{ + HashicorpVault, HttpFailure, LoginFailure, RawOperationError, SecretLocation, +}; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index f0d4fe8c0b4..447beacda85 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -31,7 +31,7 @@ use vaultrs::{ use crate::{Error, HashicorpVaultConfig, TlsCertAuth, cert_login::CertLoginRequest}; -pub use raw::RawOperationError; +pub use raw::{HttpFailure, LoginFailure, RawOperationError}; const CACHE_CAPACITY: u64 = 200; @@ -178,23 +178,25 @@ fn map_api_error(error: ClientError, context: ErrorContext) -> Error { ErrorContext::Login => Error::LoginStatus { status: code }, ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code }, }, - ClientError::JsonParseError { source } => match context { - ErrorContext::Login => Error::MalformedLogin, - ErrorContext::Read => Error::MalformedPayload, - ErrorContext::Secret => Error::Client(ClientError::JsonParseError { source }), + error @ ClientError::JsonParseError { .. } => match context { + ErrorContext::Login => Error::MalformedLogin(error), + ErrorContext::Read => Error::MalformedPayload(error), + ErrorContext::Secret => Error::Client(error), }, - ClientError::ResponseEmptyError | ClientError::ResponseDataEmptyError => { - malformed_response(context) + error @ (ClientError::ResponseEmptyError | ClientError::ResponseDataEmptyError) => { + malformed_response(context, error) } ClientError::RestClientError { source } => match source { RustifyClientError::ServerResponseError { code, .. } => match context { ErrorContext::Login => Error::LoginStatus { status: code }, ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code }, }, - RustifyClientError::ResponseParseError { .. } => malformed_response(context), - source => Error::Client(ClientError::RestClientError { source }), + source @ RustifyClientError::ResponseParseError { .. } => { + malformed_response(context, ClientError::RestClientError { source }) + } + source => ClientError::RestClientError { source }.into(), }, - error => Error::Client(error), + error => error.into(), } } @@ -208,10 +210,10 @@ fn api_status(error: &ClientError) -> Option { } } -fn malformed_response(context: ErrorContext) -> Error { +fn malformed_response(context: ErrorContext, error: ClientError) -> Error { match context { - ErrorContext::Login => Error::MalformedLogin, - ErrorContext::Read => Error::MalformedPayload, - ErrorContext::Secret => Error::Client(ClientError::ResponseDataEmptyError), + ErrorContext::Login => Error::MalformedLogin(error), + ErrorContext::Read => Error::MalformedPayload(error), + ErrorContext::Secret => Error::Client(error), } } diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/client.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/client.rs index c1dc073dcf8..01d3807d290 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/client.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/client.rs @@ -83,10 +83,8 @@ impl HashicorpVault { .ca_certs(Vec::new()) .verify(true) .build() - .map_err(|message| Error::ClientSettings { - message: message.to_string(), - })?; - VaultClient::new(settings).map_err(Error::Client) + .map_err(Error::ClientSettings)?; + VaultClient::new(settings).map_err(Error::from) } } @@ -94,16 +92,16 @@ fn identity_for(tls: Option<&TlsCertAuth>) -> Result, Error> { tls.map(|tls| { let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { path: tls.cert_path.clone(), - message: source.to_string(), + source: litellm_error::Source::new(source), })?; let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { path: tls.key_path.clone(), - message: source.to_string(), + source: litellm_error::Source::new(source), })?; Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| { Error::TlsIdentity { path: tls.cert_path.clone(), - message: source.to_string(), + source: litellm_error::Source::new(source), } }) }) diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/raw.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/raw.rs index 0325561a7a0..9c722c0b6c8 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/raw.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/raw.rs @@ -4,26 +4,123 @@ use vaultrs::api::kv2::requests::{ DeleteLatestSecretVersionRequest, ReadSecretRequest, SetSecretRequest, }; -#[derive(veil::Redact)] +#[derive(thiserror::Error, veil::Redact)] pub enum RawOperationError { + #[error(transparent)] Local(Error), + #[error("HashiCorp Vault authentication failed")] Authentication { source: Error, url: String, certificate: bool, }, - Http { - method: String, - url: String, - status: u16, + #[error("HashiCorp Vault {} request returned HTTP {}", .0.method, .0.status)] + Http(HttpFailure), + #[error("HashiCorp Vault request failed")] + Transport( + #[source] #[redact] - body: Vec, - }, - Transport(#[redact] RustifyClientError), - Timeout { - method: String, - elapsed: Duration, - }, + litellm_http::Error, + ), + #[error("HashiCorp Vault client failed")] + Client( + #[source] + #[redact] + RustifyClientError, + ), + #[error("HashiCorp Vault {method} request timed out")] + Timeout { method: String, elapsed: Duration }, +} + +#[derive(veil::Redact)] +pub struct HttpFailure { + pub method: String, + pub url: String, + pub status: u16, + #[redact] + pub body: Vec, +} + +pub struct LoginFailure<'a> { + pub source: &'a Error, + pub url: &'a str, + pub certificate: bool, +} + +impl RawOperationError { + pub fn is_unsafe_secret_name(&self) -> bool { + match self { + Self::Local(error) => error.is_unsafe_secret_name(), + Self::Authentication { .. } + | Self::Http(_) + | Self::Transport(_) + | Self::Client(_) + | Self::Timeout { .. } => false, + } + } + + pub fn local(&self) -> Option<&Error> { + match self { + Self::Local(error) => Some(error), + Self::Authentication { .. } + | Self::Http(_) + | Self::Transport(_) + | Self::Client(_) + | Self::Timeout { .. } => None, + } + } + + pub fn authentication(&self) -> Option> { + match self { + Self::Authentication { + source, + url, + certificate, + } => Some(LoginFailure { + source, + url, + certificate: *certificate, + }), + Self::Local(_) + | Self::Http(_) + | Self::Transport(_) + | Self::Client(_) + | Self::Timeout { .. } => None, + } + } + + pub fn http_failure(&self) -> Option<&HttpFailure> { + match self { + Self::Http(failure) => Some(failure), + Self::Local(_) + | Self::Authentication { .. } + | Self::Transport(_) + | Self::Client(_) + | Self::Timeout { .. } => None, + } + } + + pub fn timeout(&self) -> Option<(&str, Duration)> { + match self { + Self::Timeout { method, elapsed } => Some((method, *elapsed)), + Self::Local(_) + | Self::Authentication { .. } + | Self::Http(_) + | Self::Transport(_) + | Self::Client(_) => None, + } + } + + pub fn transport(&self) -> Option<&litellm_http::Error> { + match self { + Self::Transport(error) => Some(error), + Self::Local(_) + | Self::Authentication { .. } + | Self::Http(_) + | Self::Client(_) + | Self::Timeout { .. } => None, + } + } } impl HashicorpVault { @@ -118,7 +215,7 @@ impl HashicorpVault { let request = endpoint .with_middleware(&client.middle) .request(client.http.base()) - .map_err(RawOperationError::Transport)?; + .map_err(RawOperationError::Client)?; let method = request.method().to_string(); let url = format!( "{}/v1/{}{}/data/{}", @@ -141,14 +238,17 @@ impl HashicorpVault { })?, None => client.http.send(request).await, } - .map_err(RawOperationError::Transport)?; + .map_err(|error| match crate::error::transport(error) { + Ok(error) => RawOperationError::Transport(error), + Err(error) => RawOperationError::Client(error), + })?; if !response.status().is_success() { - return Err(RawOperationError::Http { + return Err(RawOperationError::Http(HttpFailure { method, url, status: response.status().as_u16(), body: response.into_body(), - }); + })); } Ok(response.into_body()) } diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/write.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/write.rs index 0b1acb3a179..15a0e4c7dff 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager/write.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager/write.rs @@ -37,7 +37,7 @@ impl HashicorpVault { .await { Ok(metadata) => u32::try_from(metadata.current_version) - .map_err(|_| Error::CasVersionOverflow)?, + .map_err(Error::CasVersionOverflow)?, Err(error) if api_status(&error) == Some(404) => 0, Err(_) => return Err(map_api_error(error, ErrorContext::Secret)), }; diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/configuration.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/configuration.rs index be3059ce986..b45b843fde8 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/configuration.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/configuration.rs @@ -32,7 +32,7 @@ fn invalid_refresh_intervals_are_rejected(#[case] value: &str) { assert!(matches!( HashicorpVaultConfig::from_environment(environment.as_ref()), - Err(Error::RefreshInterval) + Err(Error::RefreshInterval(_)) )); } @@ -79,6 +79,39 @@ async fn approle_login_uses_namespace_and_reuses_the_token() { assert!(manager.async_read_secret("name-2").await.unwrap().is_none()); } +#[rstest] +#[tokio::test] +async fn an_unreadable_client_certificate_keeps_its_io_cause() { + let server: MockServer = MockServer::start().await; + let directory: tempfile::TempDir = tempfile::tempdir().unwrap(); + let cert_path = directory.path().join("missing.crt"); + let values: HashMap = HashMap::from([ + ("HCP_VAULT_ADDR".to_owned(), server.uri()), + ( + "HCP_VAULT_CLIENT_CERT".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_CLIENT_KEY".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ]); + let environment: Arc = + Arc::new(move |name: &str| values.get(name).cloned()); + let manager: HashicorpVault = HashicorpVault::new(environment, true).unwrap(); + + let error = manager.async_read_secret("name").await.unwrap_err(); + + let cause = std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .and_then(|cause| cause.downcast_ref::()) + .expect("io cause"); + assert_eq!(cause.kind(), std::io::ErrorKind::NotFound); + let report = litellm_error::Report::new(&error).to_string(); + assert_eq!(report.matches(cert_path.to_str().unwrap()).count(), 1); + assert!(report.ends_with(&cause.to_string())); +} + #[rstest] #[tokio::test] async fn tls_login_posts_the_role_and_uses_the_client_identity() { diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/reads.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/reads.rs index 0c662eb5b55..c31b29fb536 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/reads.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/reads.rs @@ -138,7 +138,7 @@ async fn read_responses_distinguish_absence_and_malformed_payloads( .await; match expected { ExpectedRead::Missing => assert!(result.unwrap().is_none()), - ExpectedRead::Malformed => assert!(matches!(result, Err(Error::MalformedPayload))), + ExpectedRead::Malformed => assert!(matches!(result, Err(Error::MalformedPayload(_)))), ExpectedRead::NonString => assert!(matches!(result, Err(Error::NonStringValue))), } } @@ -383,3 +383,52 @@ async fn concurrent_reads_share_a_load_but_verification_fetches_fresh( "new" ); } + +fn unreachable_manager() -> HashicorpVault { + let address = std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap(); + let environment: Arc = Arc::new(move |name: &str| match name { + "HCP_VAULT_ADDR" => Some(format!("http://{address}")), + "HCP_VAULT_TOKEN" => Some("token".to_owned()), + _ => None, + }); + HashicorpVault::new(environment, true).unwrap() +} + +#[tokio::test] +async fn unreachable_vault_reads_are_connect_failures_without_the_url() { + let error = unreachable_manager() + .async_read_secret("do-not-leak") + .await + .unwrap_err(); + + assert!(matches!(&error, Error::Transport(transport) if transport.is_connect())); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.starts_with("HashiCorp Vault request failed: "), + "{report}" + ); + assert!(!report.contains("do-not-leak"), "{report}"); +} + +#[tokio::test] +async fn unreachable_vault_raw_reads_are_connect_failures_without_the_url() { + let error = unreachable_manager() + .read_raw("do-not-leak", &HashicorpOperationContext::default()) + .await + .unwrap_err(); + + assert!( + error + .transport() + .is_some_and(litellm_http::Error::is_connect) + ); + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.starts_with("HashiCorp Vault request failed: "), + "{report}" + ); + assert!(!report.contains("do-not-leak"), "{report}"); +} diff --git a/litellm-rust/crates/secrets-types/src/error.rs b/litellm-rust/crates/secrets-types/src/error.rs index 027dd54d8cc..891d8a0fba9 100644 --- a/litellm-rust/crates/secrets-types/src/error.rs +++ b/litellm-rust/crates/secrets-types/src/error.rs @@ -11,3 +11,15 @@ pub enum Error { #[error("secret manager received an incompatible operation context")] InvalidOperationContext, } + +impl Error { + pub fn is_unsafe_secret_name(&self) -> bool { + match self { + Self::UnsafeSecretName => true, + Self::CurrentSecretMissing + | Self::NewSecretMissing + | Self::NewSecretMismatch + | Self::InvalidOperationContext => false, + } + } +} diff --git a/litellm-rust/crates/secrets-types/tests/cache.rs b/litellm-rust/crates/secrets-types/tests/cache.rs index cec4f5198e3..18efde4289b 100644 --- a/litellm-rust/crates/secrets-types/tests/cache.rs +++ b/litellm-rust/crates/secrets-types/tests/cache.rs @@ -110,7 +110,7 @@ async fn concurrent_misses_share_a_load_and_cancellation_allows_a_retry() { } #[tokio::test] -async fn refresh_bypasses_cached_values_and_errors_and_absence_are_retried() { +async fn refresh_bypasses_cached_values() { let cache = SecretCache::new(10, Duration::from_secs(60)); cache.insert("key", "old").await; assert_eq!( @@ -127,9 +127,21 @@ async fn refresh_bypasses_cached_values_and_errors_and_absence_are_retried() { .unwrap(), Some("fresh") ); - for result in [Err("failure"), Ok(None), Ok(Some("recovered"))] { - assert_eq!(cache.read("retry", async { result }).await, result); - } +} + +#[rstest] +#[case::failure(Err("failure"))] +#[case::absence(Ok(None))] +#[tokio::test] +async fn errors_and_absence_are_retried(#[case] first: Result, &'static str>) { + let cache = SecretCache::new(10, Duration::from_secs(60)); + assert_eq!(cache.read("retry", async { first }).await, first); + assert_eq!( + cache + .read("retry", async { Ok::<_, &str>(Some("recovered")) }) + .await, + Ok(Some("recovered")) + ); } #[tokio::test] diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 3655ce8bbc2..023d61136e4 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -23,6 +23,7 @@ litellm-secrets-hashicorp = { workspace = true, optional = true } litellm-secrets-azure = { workspace = true, optional = true } litellm-secrets-cyberark = { workspace = true, optional = true } litellm-core-utils.workspace = true +litellm-error.workspace = true litellm-http.workspace = true base64.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/secrets/src/compatibility.rs b/litellm-rust/crates/secrets/src/compatibility.rs index 621f7a3f432..368c9931b51 100644 --- a/litellm-rust/crates/secrets/src/compatibility.rs +++ b/litellm-rust/crates/secrets/src/compatibility.rs @@ -40,11 +40,9 @@ pub async fn get_secret_from_python_manager( }; match result { #[cfg(feature = "google")] - Err(Error::Google(litellm_secrets_google::Error::Status(404))) => { - Err(Error::ManagedSecretMissing) - } + Err(Error::Google(error)) if error.is_not_found() => Err(Error::ManagedSecretMissing), #[cfg(feature = "azure")] - Err(Error::Azure(litellm_secrets_azure::Error::MissingValue)) => Ok(None), + Err(Error::Azure(error)) if error.is_missing_value() => Ok(None), Ok(None) if matches!(manager, SecretManager::External(_)) && manager.system() == KeyManagementSystem::AzureKeyVault => diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 7fc756c5c7d..ade8555242d 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -7,9 +7,9 @@ pub enum Error { #[error("encrypted environment value is missing")] MissingCiphertext, #[error("ciphertext is not valid base64 for the configured manager")] - InvalidCiphertext, + InvalidCiphertext(#[source] Option), #[error("decrypted value is not UTF-8")] - Utf8, + Utf8(#[source] std::str::Utf8Error), #[error(transparent)] Client(#[from] litellm_http::Error), #[error("unsupported OIDC provider or missing build feature")] @@ -19,21 +19,21 @@ pub enum Error { #[error("OIDC environment variable is missing")] MissingEnvironment, #[error("OIDC request failed")] - OidcHttp, + OidcHttp(#[source] litellm_http::Error), #[error("OIDC provider returned HTTP {0}")] OidcStatus(u16), #[error("OIDC response is invalid")] - OidcResponse, + OidcResponse(#[source] litellm_http::Error), #[error("OIDC file path must be absolute and within the credential allowlist")] UnsafeOidcPath, #[error("OIDC file could not be read")] - OidcFile, + OidcFile(#[source] std::io::Error), #[error("secret cannot be converted to {expected}")] TypeMismatch { expected: &'static str }, #[error("external secret manager failed")] - ExternalManager(#[source] Box), + ExternalManager(#[source] litellm_error::Source), #[error("external secret manager read failed")] - ExternalRead(#[source] Box), + ExternalRead(#[source] litellm_error::Source), #[cfg(feature = "aws")] #[error(transparent)] Aws(#[from] litellm_secrets_aws::Error), diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index db37d793384..f42e09c6bdb 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -87,7 +87,8 @@ pub async fn get_secret_from_manager( let plaintext = client .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Permissive)?) .await?; - let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + let value = + String::from_utf8(plaintext).map_err(|error| Error::Utf8(error.utf8_error()))?; Ok(Some(Secret::String(SecretValue::new(value.trim())))) } #[cfg(feature = "google")] @@ -98,7 +99,8 @@ pub async fn get_secret_from_manager( let plaintext = client .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Canonical)?) .await?; - let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + let value = + String::from_utf8(plaintext).map_err(|error| Error::Utf8(error.utf8_error()))?; Ok(Some(Secret::String(SecretValue::new(value)))) } #[cfg(feature = "aws")] @@ -149,9 +151,9 @@ fn decode_ciphertext(value: &str, mode: Base64Mode) -> Result, Error> { }; let ciphertext = STANDARD .decode(&encoded) - .map_err(|_| Error::InvalidCiphertext)?; + .map_err(|error| Error::InvalidCiphertext(Some(error)))?; if canonical && STANDARD.encode(&ciphertext) != encoded { - return Err(Error::InvalidCiphertext); + return Err(Error::InvalidCiphertext(None)); } Ok(ciphertext) } diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index adeb6c3db00..24e6912c797 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -191,7 +191,7 @@ impl OidcResolver { .header("Accept", "application/json; api-version=2.0") .send() .await - .map_err(|_| Error::OidcHttp)?; + .map_err(|error| Error::OidcHttp(error.into()))?; if response.status() != reqwest::StatusCode::OK { return Err(Error::OidcStatus(response.status().as_u16())); } @@ -199,7 +199,10 @@ impl OidcResolver { struct Token { value: Option, } - let token: Token = response.json().await.map_err(|_| Error::OidcResponse)?; + let token: Token = response + .json() + .await + .map_err(|error| Error::OidcResponse(error.into()))?; if let Some(value) = &token.value { self.cache .insert( @@ -225,11 +228,14 @@ impl OidcResolver { .header("Metadata-Flavor", "Google") .send() .await - .map_err(|_| Error::OidcHttp)?; + .map_err(|error| Error::OidcHttp(error.into()))?; if response.status() != reqwest::StatusCode::OK { return Err(Error::OidcStatus(response.status().as_u16())); } - let token = response.text().await.map_err(|_| Error::OidcResponse)?; + let token = response + .text() + .await + .map_err(|error| Error::OidcResponse(error.into()))?; let now = (self.clock)(); let ttl = oidc_token_cache_ttl(&token, now, GOOGLE_TOKEN_MAX_TTL); let value = SecretValue::new(token); @@ -259,7 +265,7 @@ async fn read_file(path: &str) -> Result { tokio::fs::read_to_string(path) .await .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) - .map_err(|_| Error::OidcFile) + .map_err(Error::OidcFile) } async fn read_allowed_file( @@ -271,7 +277,7 @@ async fn read_allowed_file( } let resolved = tokio::fs::canonicalize(path) .await - .map_err(|_| Error::OidcFile)?; + .map_err(Error::OidcFile)?; let allowed = environment .get(OIDC_ALLOWED_CREDENTIAL_DIRS) .filter(|v| !v.is_empty()) @@ -283,7 +289,7 @@ async fn read_allowed_file( return tokio::fs::read_to_string(&resolved) .await .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) - .map_err(|_| Error::OidcFile); + .map_err(Error::OidcFile); } } Err(Error::UnsafeOidcPath) diff --git a/litellm-rust/crates/secrets/tests/aws.rs b/litellm-rust/crates/secrets/tests/aws.rs index 910b1855b77..abad04f31c5 100644 --- a/litellm-rust/crates/secrets/tests/aws.rs +++ b/litellm-rust/crates/secrets/tests/aws.rs @@ -53,7 +53,7 @@ async fn read_results_follow_the_selected_failure_policy( }, )), Arc::new(move |_: &str| environment.map(str::to_owned)), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(policy); let result = resolver @@ -97,7 +97,7 @@ async fn primary_secret_values_other_than_strings_resolve_to_none( let resolver = SecretResolver::new_python_compatible( Arc::new(state(&server, settings)), Arc::new(|_: &str| Some("fallback".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); let text = value.as_str(); assert_eq!( @@ -157,7 +157,7 @@ async fn gating_prediction_matches_actual_lookup( let resolver = SecretResolver::new_python_compatible( Arc::new(state), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver @@ -218,6 +218,6 @@ async fn aws_handler_reads_ciphertext_decodes_trims_and_redacts() { )); assert!(matches!( get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some("abc".into())).await, - Err(Error::InvalidCiphertext) + Err(Error::InvalidCiphertext(_)) )); } diff --git a/litellm-rust/crates/secrets/tests/azure.rs b/litellm-rust/crates/secrets/tests/azure.rs index 60f165add54..6e81017b0f6 100644 --- a/litellm-rust/crates/secrets/tests/azure.rs +++ b/litellm-rust/crates/secrets/tests/azure.rs @@ -22,7 +22,7 @@ async fn azure_handler_reads_missing_and_failed_secrets() { .await; let manager = SecretManager::AzureKeyVault( AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), std::sync::Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), ) @@ -81,7 +81,7 @@ async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_ .mount(&server) .await; let manager = AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "token".into())), ) @@ -92,7 +92,7 @@ async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_ Default::default(), )), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver diff --git a/litellm-rust/crates/secrets/tests/common_read_contract.rs b/litellm-rust/crates/secrets/tests/common_read_contract.rs index 2e8fbf46008..e6928038799 100644 --- a/litellm-rust/crates/secrets/tests/common_read_contract.rs +++ b/litellm-rust/crates/secrets/tests/common_read_contract.rs @@ -59,7 +59,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { ), Provider::Azure => SecretManager::AzureKeyVault( AzureKeyVault::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), environment, ) @@ -67,7 +67,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { ), Provider::Google => SecretManager::GoogleSecretManager( GoogleSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "project".into(), environment, @@ -84,7 +84,7 @@ fn manager(provider: Provider, server: &MockServer) -> SecretManager { .unwrap(), ), Provider::Cyberark => SecretManager::Cyberark(CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), @@ -240,7 +240,7 @@ async fn python_read_failures_preserve_provider_fallback_rules( KeyManagementSettings::default(), )), Arc::new(move |_: &str| environment_value.map(str::to_owned)), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); let expected = if matches!(provider, Provider::Aws) { None diff --git a/litellm-rust/crates/secrets/tests/cyberark.rs b/litellm-rust/crates/secrets/tests/cyberark.rs index b94bd9dc0ad..0a6c394937e 100644 --- a/litellm-rust/crates/secrets/tests/cyberark.rs +++ b/litellm-rust/crates/secrets/tests/cyberark.rs @@ -24,7 +24,7 @@ async fn cyberark_handler_reads_values_and_surfaces_errors() { .mount(&server) .await; let manager = SecretManager::Cyberark(CyberArkSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "acct".into(), "admin".into(), diff --git a/litellm-rust/crates/secrets/tests/google.rs b/litellm-rust/crates/secrets/tests/google.rs index 0b45a90c10b..5542c864dfd 100644 --- a/litellm-rust/crates/secrets/tests/google.rs +++ b/litellm-rust/crates/secrets/tests/google.rs @@ -25,7 +25,7 @@ async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) _ => None, }); let manager = GoogleSecretManager::with_client( - litellm_http::Client::plain_for_test(), + litellm_http::Client::for_test(), server.uri().parse().unwrap(), "project".into(), environment.clone(), @@ -40,7 +40,7 @@ async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) let resolver = SecretResolver::new_python_compatible( Arc::new(state), environment, - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(FailurePolicy::Propagate); let result = resolver.get_secret_str("KEY", None).await; @@ -108,7 +108,7 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites STANDARD.encode("encrypted") ))) .await, - Err(Error::InvalidCiphertext) + Err(Error::InvalidCiphertext(_)) )); assert!(matches!( get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await, diff --git a/litellm-rust/crates/secrets/tests/hashicorp.rs b/litellm-rust/crates/secrets/tests/hashicorp.rs index e10e903c25a..bfb289cdb97 100644 --- a/litellm-rust/crates/secrets/tests/hashicorp.rs +++ b/litellm-rust/crates/secrets/tests/hashicorp.rs @@ -56,7 +56,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { }, )), Arc::new(|_: &str| None), - litellm_secrets::OidcResolver::new(litellm_http::Client::plain_for_test()), + litellm_secrets::OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( found_resolver @@ -131,7 +131,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let failed_resolver = SecretResolver::new_python_compatible( Arc::new(failed_state), Arc::new(|_: &str| None), - litellm_secrets::OidcResolver::new(litellm_http::Client::plain_for_test()), + litellm_secrets::OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(FailurePolicy::Propagate); assert!(matches!( diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs index 6f72bd9d645..867a662cbbb 100644 --- a/litellm-rust/crates/secrets/tests/oidc.rs +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -30,7 +30,7 @@ async fn environment_sources_resolve_expected_value( ("CIRCLE_OIDC_TOKEN_V2", "circle-v2"), ]); assert_eq!( - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .resolve(reference, env.as_ref()) .await .unwrap() @@ -43,7 +43,7 @@ async fn environment_sources_resolve_expected_value( #[tokio::test] async fn environment_sources_bypass_boolean_conversion_and_defaults() { let env = environment(&[("TOKEN", "true")]); - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); + let oidc = OidcResolver::new(litellm_http::Client::for_test()); let resolver = SecretResolver::new(Arc::new(SecretManagerState::default()), env, oidc); assert_eq!( resolver @@ -94,7 +94,7 @@ async fn github_requests_are_authenticated_cached_and_revalidate_environment() { ), ("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "request-token"), ]); - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); + let oidc = OidcResolver::new(litellm_http::Client::for_test()); for _ in 0..2 { assert_eq!( oidc.resolve("oidc/github/https://service/oidc/path", env.as_ref()) @@ -131,7 +131,7 @@ async fn file_allowlist_resolves_symlinks_while_environment_paths_remain_explici ("PATH_TOKEN", private.to_str().unwrap()), ("AZURE_FEDERATED_TOKEN_FILE", token.to_str().unwrap()), ]); - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()); + let oidc = OidcResolver::new(litellm_http::Client::for_test()); assert_eq!( oidc.resolve(&format!("oidc/file/{}", token.display()), env.as_ref()) .await @@ -213,7 +213,7 @@ async fn google_expiry_caps_cache_and_preserves_audience( .expect(calls) .mount(&server) .await; - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()) + let oidc = OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()) .with_clock(now); for _ in 0..2 { @@ -235,7 +235,7 @@ async fn google_expiry_caps_cache_and_preserves_audience( #[tokio::test] async fn google_oidc_requires_its_build_feature() { assert!(matches!( - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .resolve("oidc/google/audience", environment(&[]).as_ref()) .await, Err(Error::UnsupportedOidc) @@ -246,7 +246,7 @@ async fn google_oidc_requires_its_build_feature() { #[tokio::test] async fn azure_oidc_without_a_token_file_requires_its_build_feature() { assert!(matches!( - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .resolve("oidc/azure/scope", environment(&[]).as_ref()) .await, Err(Error::UnsupportedOidc) @@ -262,7 +262,7 @@ async fn invalid_references_fail_before_environment_lookup( #[case] reference: &str, #[case] unsupported: bool, ) { - let error = OidcResolver::new(litellm_http::Client::plain_for_test()) + let error = OidcResolver::new(litellm_http::Client::for_test()) .resolve(reference, &|_: &str| { panic!("invalid reference reached environment lookup") }) @@ -284,7 +284,7 @@ async fn unreadable_expiry_keeps_python_cache_fallback(#[case] token: &str) { .expect(1) .mount(&server) .await; - let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + let resolver = OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()); for _ in 0..2 { assert_eq!( @@ -336,7 +336,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] }) } } - let oidc = OidcResolver::new(litellm_http::Client::plain_for_test()) + let oidc = OidcResolver::new(litellm_http::Client::for_test()) .with_azure_token_provider(Arc::new(Provider(failed))); let resolver = SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), @@ -364,7 +364,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] #[tokio::test] async fn missing_oidc_environment_is_an_error(#[case] reference: &str) { assert!(matches!( - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .resolve(reference, environment(&[]).as_ref()) .await, Err(Error::MissingEnvironment) @@ -383,7 +383,7 @@ async fn google_oidc_failures_are_not_cached_or_hidden_by_defaults() { let resolver = SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), environment(&[]), - OidcResolver::new(litellm_http::Client::plain_for_test()) + OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()), ); for _ in 0..2 { @@ -424,7 +424,7 @@ async fn google_tokens_expire_at_the_python_cache_deadline( .expect(2) .mount(&server) .await; - let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + let resolver = OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()) .with_clock(|| UNIX_EPOCH + Duration::from_secs(1000)); assert_eq!( @@ -477,7 +477,7 @@ async fn google_cache_uses_payload_expiry_without_requiring_a_jwt_header() { .expect(2) .mount(&server) .await; - let resolver = OidcResolver::new(litellm_http::Client::plain_for_test()) + let resolver = OidcResolver::new(litellm_http::Client::for_test()) .with_google_identity_endpoint(server.uri().parse().unwrap()) .with_clock(|| UNIX_EPOCH + Duration::from_secs(1000)); for _ in 0..2 { diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index 37557c36fb5..8435e051315 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -12,79 +12,86 @@ fn resolver(value: Option<&str>) -> SecretResolver { SecretResolver::new_python_compatible( Arc::new(SecretManagerState::default()), Arc::new(move |_: &str| value.clone()), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) } #[rstest::rstest] -#[case::environment(false)] -#[case::manager(true)] +#[case::capitalized_true("True", Some(true))] +#[case::padded_false(" FALSE ", Some(false))] +#[case::empty("", None)] +#[case::parenthesized("(True)", None)] +#[case::json("{\"key\":1}", None)] #[tokio::test] -async fn native_reads_preserve_strings_and_report_conversion_errors(#[case] managed: bool) { - for raw in ["True", " FALSE ", "", "(True)", "{\"key\":1}"] { - let state = if managed { - SecretManagerState::new( - SecretManager::External(Arc::new(FixedManager::custom(Ok(Some(Secret::String( - SecretValue::new(raw), - )))))), - KeyManagementSettings::default(), - ) - } else { - SecretManagerState::default() - }; - let resolver = SecretResolver::new( - Arc::new(state), - Arc::new(move |_: &str| Some(raw.to_owned())), - OidcResolver::new(litellm_http::Client::plain_for_test()), - ); - assert_eq!( - resolver - .get_secret_str("key", None) - .await - .unwrap() - .unwrap() - .expose(), - raw - ); - assert_eq!( - resolver.get_secret("key", None).await.unwrap(), - Some(Secret::String(SecretValue::new(raw))) - ); - if raw == "True" || raw == " FALSE " { - assert_eq!( - resolver.get_secret_bool("key", None).await.unwrap(), - Some(raw == "True") - ); - } else { - assert!(matches!( - resolver.get_secret_bool("key", None).await, - Err(Error::TypeMismatch { - expected: "boolean" - }) - )); - } +async fn native_reads_preserve_strings_and_report_conversion_errors( + #[case] raw: &'static str, + #[case] boolean: Option, + #[values(false, true)] managed: bool, +) { + let state = if managed { + SecretManagerState::new( + SecretManager::External(Arc::new(FixedManager::custom(Ok(Some(Secret::String( + SecretValue::new(raw), + )))))), + KeyManagementSettings::default(), + ) + } else { + SecretManagerState::default() + }; + let resolver = SecretResolver::new( + Arc::new(state), + Arc::new(move |_: &str| Some(raw.to_owned())), + OidcResolver::new(litellm_http::Client::for_test()), + ); + assert_eq!( + resolver + .get_secret_str("key", None) + .await + .unwrap() + .unwrap() + .expose(), + raw + ); + assert_eq!( + resolver.get_secret("key", None).await.unwrap(), + Some(Secret::String(SecretValue::new(raw))) + ); + match boolean { + Some(expected) => assert_eq!( + resolver.get_secret_bool("key", None).await.unwrap(), + Some(expected) + ), + None => assert!(matches!( + resolver.get_secret_bool("key", None).await, + Err(Error::TypeMismatch { + expected: "boolean" + }) + )), } } +#[rstest::rstest] +#[case::absent(Ok(None))] +#[case::failed(Err(()))] #[tokio::test] -async fn native_defaults_apply_to_absence_but_never_hide_provider_failures() { - for reply in [Ok(None), Err(())] { - let resolver = SecretResolver::new( - Arc::new(SecretManagerState::new( - SecretManager::External(Arc::new(FixedManager::custom(reply.clone()))), - KeyManagementSettings::default(), - )), - Arc::new(|_: &str| None), - OidcResolver::new(litellm_http::Client::plain_for_test()), - ); - let result = resolver - .get_secret_str("key", Some(SecretValue::new("default"))) - .await; - match reply { - Ok(None) => assert_eq!(result.unwrap().unwrap().expose(), "default"), - Err(()) => assert!(matches!(result, Err(Error::MissingCiphertext))), - Ok(Some(_)) => unreachable!(), - } +async fn native_defaults_apply_to_absence_but_never_hide_provider_failures( + #[case] reply: Result, ()>, +) { + let resolver = SecretResolver::new( + Arc::new(SecretManagerState::new( + SecretManager::External(Arc::new(FixedManager::custom(reply.clone()))), + KeyManagementSettings::default(), + )), + Arc::new(|_: &str| None), + OidcResolver::new(litellm_http::Client::for_test()), + ); + let result = resolver + .get_secret_str("key", Some(SecretValue::new("default"))) + .await; + match reply { + Ok(None) => assert_eq!(result.unwrap().unwrap().expose(), "default"), + Err(()) => assert!(matches!(result, Err(Error::MissingCiphertext))), + Ok(Some(_)) => unreachable!(), } } @@ -189,7 +196,7 @@ fn managed(reply: Result, ()>, environment: Option<&'static str>) KeyManagementSettings::default(), )), Arc::new(move |_: &str| environment.map(str::to_owned)), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ) .with_failure_policy(FailurePolicy::EnvironmentFallback) } @@ -281,7 +288,7 @@ async fn prefix_is_removed_once_and_resolved_from_environment() { let resolver = SecretResolver::new_python_compatible( Arc::new(state), Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver @@ -340,7 +347,7 @@ async fn excluded_hosted_keys_keep_the_python_manager_conversion_path( }, )), Arc::new(move |_: &str| Some(raw.to_owned())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver.get_secret("KEY", None).await.unwrap(), @@ -373,7 +380,7 @@ async fn azure_callback_absence_preserves_none_but_errors_fall_back( KeyManagementSettings::default(), )), Arc::new(|_: &str| Some("environment".into())), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); assert_eq!( resolver diff --git a/litellm-rust/crates/secrets/tests/source.rs b/litellm-rust/crates/secrets/tests/source.rs index 17210940a67..41f0ecea518 100644 --- a/litellm-rust/crates/secrets/tests/source.rs +++ b/litellm-rust/crates/secrets/tests/source.rs @@ -15,7 +15,7 @@ mod tests { #[case] expected: Option<&str>, ) { unsafe { std::env::set_var(name, value) }; - let secret = EnvironmentSecrets::python_compatible(litellm_http::Client::plain_for_test()) + let secret = EnvironmentSecrets::python_compatible(litellm_http::Client::for_test()) .resolve(&[name]) .await .unwrap() @@ -42,7 +42,7 @@ async fn dynamic_names_use_the_same_resolver_and_snapshots_never_do_fresh_lookup reads.fetch_add(1, Ordering::SeqCst); (name != "missing").then(|| name.to_owned()) }), - OidcResolver::new(litellm_http::Client::plain_for_test()), + OidcResolver::new(litellm_http::Client::for_test()), ); let snapshot = source.resolve(&["declared", "missing"]).await.unwrap(); let name = format!("runtime-{}", "key"); diff --git a/litellm-rust/crates/spend-clickhouse/Cargo.toml b/litellm-rust/crates/spend-clickhouse/Cargo.toml index 893c5daded3..ef7f4a0efdf 100644 --- a/litellm-rust/crates/spend-clickhouse/Cargo.toml +++ b/litellm-rust/crates/spend-clickhouse/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] flate2.workspace = true +litellm-error.workspace = true litellm-http.workspace = true litellm-storage-clickhouse.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/spend-clickhouse/src/error.rs b/litellm-rust/crates/spend-clickhouse/src/error.rs index 0d037ef06f8..c5a9f542ca6 100644 --- a/litellm-rust/crates/spend-clickhouse/src/error.rs +++ b/litellm-rust/crates/spend-clickhouse/src/error.rs @@ -1,15 +1,32 @@ +use litellm_error::Source; + #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("invalid ClickHouse insert row")] - InvalidRow, - #[error("{0} must be a positive integer")] - InvalidLimit(&'static str), - #[error("database must be a nonempty SQL identifier and retention must be positive")] - InvalidSchema, - #[error("ClickHouse insert exceeds the encoded size limit")] - InsertTooLarge, + #[error("{field} must be an integer millisecond timestamp")] + NonIntegerTimestamp { field: &'static str }, + #[error("{field} is outside the supported timestamp range")] + TimestampOutOfRange { field: &'static str, source: Source }, #[error(transparent)] Storage(#[from] litellm_storage_clickhouse::Error), #[error(transparent)] Migration(#[from] sqlx::migrate::MigrateError), } + +impl Error { + pub fn is_request(&self) -> bool { + match self { + Self::NonIntegerTimestamp { .. } | Self::TimestampOutOfRange { .. } => true, + Self::Storage(error) => error.is_request(), + Self::Migration(_) => false, + } + } + + pub fn is_insert_too_large(&self) -> bool { + match self { + Self::Storage(error) => error.is_insert_too_large(), + Self::NonIntegerTimestamp { .. } + | Self::TimestampOutOfRange { .. } + | Self::Migration(_) => false, + } + } +} diff --git a/litellm-rust/crates/spend-clickhouse/src/insert.rs b/litellm-rust/crates/spend-clickhouse/src/insert.rs index 495f0d8ecb2..66d0ea59536 100644 --- a/litellm-rust/crates/spend-clickhouse/src/insert.rs +++ b/litellm-rust/crates/spend-clickhouse/src/insert.rs @@ -13,7 +13,8 @@ use sha2::{Digest, Sha256}; use time::{OffsetDateTime, format_description::well_known::Rfc3339}; use crate::Error; -use litellm_storage_clickhouse::Connection; +use litellm_error::Source; +use litellm_storage_clickhouse::{Connection, Error as StorageError}; fn max_insert_bytes() -> Result { let name = "CLICKHOUSE_TRACE_MAX_INSERT_BYTES"; @@ -22,9 +23,9 @@ fn max_insert_bytes() -> Result { .parse::() .ok() .filter(|value| *value > 0) - .ok_or(Error::InvalidLimit(name)), + .ok_or(StorageError::InvalidLimit(name).into()), Err(std::env::VarError::NotPresent) => Ok(64 * 1024 * 1024), - Err(_) => Err(Error::InvalidLimit(name)), + Err(_) => Err(StorageError::InvalidLimit(name).into()), } } @@ -55,7 +56,11 @@ pub async fn insert_rows( pub fn encode_rows(rows: Vec) -> Result { let body = write_rows(&rows, None, Vec::new(), usize::MAX)?; - String::from_utf8(body).map_err(|_| Error::InvalidRow) + String::from_utf8(body).map_err(invalid_row) +} + +fn invalid_row(error: impl std::error::Error + Send + Sync + 'static) -> Error { + StorageError::InvalidRow(Source::new(error)).into() } fn prepare_insert( @@ -73,9 +78,9 @@ fn prepare_insert( )?; let body = encoder .into_inner() - .map_err(|_| Error::InvalidRow)? + .map_err(|error| invalid_row(error.into_error()))? .finish() - .map_err(|_| Error::InvalidRow)?; + .map_err(invalid_row)?; Ok((token, body)) } @@ -102,7 +107,7 @@ impl Write for LimitedWriter { fn write(&mut self, bytes: &[u8]) -> std::io::Result { if bytes.len() > self.remaining { self.exceeded = true; - return Err(std::io::Error::other(Error::InsertTooLarge)); + return Err(std::io::Error::other(StorageError::InsertTooLarge)); } let written = self.inner.write(bytes)?; self.remaining -= written; @@ -126,17 +131,18 @@ fn write_rows( exceeded: false, }; for (index, row) in rows.iter().enumerate() { + let row = EncodedRow::new(row, received_ms)?; let result = (|| { if index != 0 { writer.write_all(b"\n").map_err(serde_json::Error::io)?; } - serde_json::to_writer(&mut writer, &EncodedRow { row, received_ms }) + serde_json::to_writer(&mut writer, &row) })(); - if result.is_err() { + if let Err(error) = result { return Err(if writer.exceeded { - Error::InsertTooLarge + StorageError::InsertTooLarge.into() } else { - Error::InvalidRow + invalid_row(error) }); } } @@ -144,25 +150,37 @@ fn write_rows( } struct EncodedRow<'a> { - row: &'a InsertRow, + fields: Vec<(&'a str, Cow<'a, Value>)>, received_ms: Option, } +impl<'a> EncodedRow<'a> { + fn new(row: &'a InsertRow, received_ms: Option) -> Result { + let fields = row + .iter() + .map(|(name, value)| Ok((name.as_str(), insert_value(name, value)?))) + .collect::>()?; + Ok(Self { + fields, + received_ms, + }) + } +} + impl Serialize for EncodedRow<'_> { fn serialize(&self, serializer: S) -> Result { let mut map = serializer.serialize_map(None)?; let mut received_ms = self.received_ms; - for (name, value) in self.row { - if name.as_str() >= "EngineReceivedMs" + for (name, value) in &self.fields { + if *name >= "EngineReceivedMs" && let Some(timestamp) = received_ms.take() { map.serialize_entry("EngineReceivedMs", ×tamp)?; } - if name == "EngineReceivedMs" && self.received_ms.is_some() { + if *name == "EngineReceivedMs" && self.received_ms.is_some() { continue; } - let value = insert_value(name, value).map_err(serde::ser::Error::custom)?; - map.serialize_entry(name, &value)?; + map.serialize_entry(name, value)?; } if let Some(timestamp) = received_ms { map.serialize_entry("EngineReceivedMs", ×tamp)?; @@ -171,21 +189,23 @@ impl Serialize for EncodedRow<'_> { } } +const TIMESTAMP_FIELDS: [&str; 3] = ["start_time", "end_time", "completion_start_time"]; + fn insert_value<'a>(name: &str, value: &'a Value) -> Result, Error> { - let multiplier = match name { - "start_time" | "end_time" | "completion_start_time" => 1_000_000, - _ => return Ok(Cow::Borrowed(value)), + let Some(&field) = TIMESTAMP_FIELDS.iter().find(|field| **field == name) else { + return Ok(Cow::Borrowed(value)); }; - if name == "completion_start_time" && value.is_null() { + if field == "completion_start_time" && value.is_null() { return Ok(Cow::Borrowed(value)); } - let timestamp = value.as_i64().ok_or(Error::InvalidRow)?; - let datetime = OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp) * multiplier) - .map_err(|_| Error::InvalidRow)?; + let out_of_range = |source| Error::TimestampOutOfRange { field, source }; + let timestamp = value.as_i64().ok_or(Error::NonIntegerTimestamp { field })?; + let datetime = OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp) * 1_000_000) + .map_err(|error| out_of_range(Source::new(error)))?; datetime .format(&Rfc3339) .map(|value| Cow::Owned(Value::String(value))) - .map_err(|_| Error::InvalidRow) + .map_err(|error| out_of_range(Source::new(error))) } #[cfg(test)] @@ -204,10 +224,11 @@ mod tests { ]; let encoded = write_rows(&rows, None, Vec::new(), usize::MAX).unwrap(); assert!(write_rows(&rows, None, Vec::new(), encoded.len()).is_ok()); - assert!(matches!( - write_rows(&rows, None, Vec::new(), encoded.len() - 1), - Err(Error::InsertTooLarge) - )); + assert!( + write_rows(&rows, None, Vec::new(), encoded.len() - 1) + .unwrap_err() + .is_insert_too_large() + ); } #[rstest] @@ -264,10 +285,11 @@ mod tests { fn stamped_insert_enforces_the_encoded_limit() { let rows = vec![BTreeMap::new()]; assert!(prepare_insert(&rows, 1, 22).is_ok()); - assert!(matches!( - prepare_insert(&rows, 1, 21), - Err(Error::InsertTooLarge) - )); + assert!( + prepare_insert(&rows, 1, 21) + .unwrap_err() + .is_insert_too_large() + ); } #[rstest] @@ -282,15 +304,36 @@ mod tests { assert_eq!(insert_value(field, &input).unwrap().as_ref(), &expected); } + fn rejected_start_time(input: Value) -> Error { + prepare_insert(&[BTreeMap::from([("start_time".into(), input)])], 456, 4096).unwrap_err() + } + #[rstest] #[case::fractional(json!(1.5))] #[case::string(json!("1234"))] #[case::null(Value::Null)] - #[case::out_of_range(json!(i64::MAX))] - fn invalid_spend_dates_fail_before_transport(#[case] input: Value) { + fn non_integer_spend_dates_fail_before_transport(#[case] input: Value) { + let error = rejected_start_time(input); + assert!(error.is_request()); assert!(matches!( - prepare_insert(&[BTreeMap::from([("start_time".into(), input)])], 456, 4096), - Err(Error::InvalidRow) + error, + Error::NonIntegerTimestamp { + field: "start_time" + } )); } + + #[rstest] + fn out_of_range_spend_dates_keep_their_cause() { + let error = rejected_start_time(json!(i64::MAX)); + assert!(error.is_request()); + let cause = std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .expect("time cause"); + assert!( + cause + .downcast_ref::() + .is_some() + ); + } } diff --git a/litellm-rust/crates/spend-clickhouse/src/lib.rs b/litellm-rust/crates/spend-clickhouse/src/lib.rs index f4c108b1e7f..9da658aa3e8 100644 --- a/litellm-rust/crates/spend-clickhouse/src/lib.rs +++ b/litellm-rust/crates/spend-clickhouse/src/lib.rs @@ -4,7 +4,7 @@ mod schema; pub use error::Error; pub use insert::{encode_rows, insert_rows}; -use litellm_storage_clickhouse::Storage; +use litellm_storage_clickhouse::{Error as StorageError, Storage}; pub use schema::ensure_schema; #[derive(Clone)] @@ -16,7 +16,7 @@ pub struct Config { impl Config { pub fn new(database: String, url: &str, retention_days: u32) -> Result { if retention_days == 0 { - return Err(Error::InvalidSchema); + return Err(StorageError::InvalidSchema.into()); } Ok(Self { storage: Storage::new(database, url)?, diff --git a/litellm-rust/crates/spend-clickhouse/tests/storage.rs b/litellm-rust/crates/spend-clickhouse/tests/storage.rs index 2b3eb75c658..81b6527f9e1 100644 --- a/litellm-rust/crates/spend-clickhouse/tests/storage.rs +++ b/litellm-rust/crates/spend-clickhouse/tests/storage.rs @@ -77,7 +77,7 @@ async fn database() -> Database { 7, ) .unwrap(), - client: Client::no_redirect_for_test(), + client: Client::for_test(), } } diff --git a/litellm-rust/crates/spend-clickhouse/tests/transport.rs b/litellm-rust/crates/spend-clickhouse/tests/transport.rs index ffc3423f65e..49a4b365135 100644 --- a/litellm-rust/crates/spend-clickhouse/tests/transport.rs +++ b/litellm-rust/crates/spend-clickhouse/tests/transport.rs @@ -29,7 +29,7 @@ async fn spend_insert_controls_target_and_retry_settings() { .mount(&server) .await; let config = Config::new("spend_test".into(), &format!("{}?query=DROP&async_insert=0&wait_for_async_insert=0&input_format_skip_unknown_fields=1&database=wrong&readonly=1", server.uri()), 7).unwrap(); - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let row = BTreeMap::from([ ("request_id".into(), json!("request")), ("start_time".into(), json!(1234)), @@ -77,7 +77,7 @@ async fn spend_insert_preserves_storage_failure(#[case] status: u16) { .await; let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap(); let error = insert_rows( - &Client::no_redirect_for_test(), + &Client::for_test(), config.storage().writer(), "spend_test", vec![BTreeMap::from([("request_id".into(), json!("request"))])], @@ -95,7 +95,7 @@ async fn empty_spend_insert_does_not_contact_storage() { let server = MockServer::start().await; let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap(); insert_rows( - &Client::no_redirect_for_test(), + &Client::for_test(), config.storage().writer(), "spend_test", vec![], @@ -116,7 +116,7 @@ async fn failed_schema_setup_is_reported() { .await; let config = Config::new("spend_test".into(), &server.uri(), 7).unwrap(); assert!(matches!( - ensure_schema(&Client::no_redirect_for_test(), &config).await, + ensure_schema(&Client::for_test(), &config).await, Err(Error::Storage( litellm_storage_clickhouse::Error::SchemaFailed(403) )) diff --git a/litellm-rust/crates/storage-clickhouse/Cargo.toml b/litellm-rust/crates/storage-clickhouse/Cargo.toml index 39d04f2de5d..8b985fc40d8 100644 --- a/litellm-rust/crates/storage-clickhouse/Cargo.toml +++ b/litellm-rust/crates/storage-clickhouse/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] flate2.workspace = true +litellm-error.workspace = true litellm-http.workspace = true serde.workspace = true serde_json.workspace = true @@ -17,6 +18,7 @@ thiserror.workspace = true url.workspace = true [dev-dependencies] +litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } rstest.workspace = true testcontainers-modules = { workspace = true, features = ["clickhouse"] } diff --git a/litellm-rust/crates/storage-clickhouse/src/error.rs b/litellm-rust/crates/storage-clickhouse/src/error.rs index fa0c5eabd73..ef7b5046455 100644 --- a/litellm-rust/crates/storage-clickhouse/src/error.rs +++ b/litellm-rust/crates/storage-clickhouse/src/error.rs @@ -1,21 +1,23 @@ +use litellm_error::Source; + #[derive(Clone, Debug, thiserror::Error)] pub enum Error { #[error("invalid ClickHouse insert row")] - InvalidRow, + InvalidRow(#[source] Source), #[error("{0} must be a positive integer")] InvalidLimit(&'static str), #[error("invalid ClickHouse insert table")] InvalidTable, #[error("invalid ClickHouse HTTP URL")] - InvalidUrl, + InvalidUrl(#[source] url::ParseError), + #[error("ClickHouse URL must use http or https and name a host")] + UnsupportedUrl, #[error("database must be a nonempty SQL identifier and retention must be positive")] InvalidSchema, #[error("SQL query must not be empty")] EmptySql, #[error("invalid ClickHouse query parameters")] - InvalidParameters, - #[error("unknown ClickHouse read query")] - InvalidQuery, + InvalidParameters(#[source] Source), #[error("ClickHouse query failed with HTTP status {0}")] QueryFailed(u16), #[error("ClickHouse insert failed with HTTP status {0}")] @@ -24,10 +26,60 @@ pub enum Error { InsertTooLarge, #[error("ClickHouse schema setup failed with HTTP status {0}")] SchemaFailed(u16), + #[error("ClickHouse migrations are forward-only")] + Irreversible, #[error("ClickHouse query exceeded the response size limit")] ResponseTooLarge, - #[error("ClickHouse returned an invalid or failed JSON query response")] - InvalidResponse, - #[error("ClickHouse query transport failed")] - Transport, + #[error("ClickHouse returned an invalid response")] + InvalidResponse(#[source] Source), + #[error("ClickHouse returned a failed or unexpected response")] + UnexpectedResponse, + #[error("ClickHouse transport failed")] + Transport(#[from] litellm_http::Error), +} + +impl Error { + pub fn is_request(&self) -> bool { + match self { + Self::InvalidRow(_) + | Self::InvalidLimit(_) + | Self::InvalidTable + | Self::InvalidSchema + | Self::EmptySql + | Self::InvalidParameters(_) => true, + Self::InvalidUrl(_) + | Self::UnsupportedUrl + | Self::QueryFailed(_) + | Self::InsertFailed(_) + | Self::InsertTooLarge + | Self::SchemaFailed(_) + | Self::Irreversible + | Self::ResponseTooLarge + | Self::InvalidResponse(_) + | Self::UnexpectedResponse + | Self::Transport(_) => false, + } + } + + pub fn is_insert_too_large(&self) -> bool { + match self { + Self::InsertTooLarge => true, + Self::InvalidRow(_) + | Self::InvalidLimit(_) + | Self::InvalidTable + | Self::InvalidUrl(_) + | Self::UnsupportedUrl + | Self::InvalidSchema + | Self::EmptySql + | Self::InvalidParameters(_) + | Self::QueryFailed(_) + | Self::InsertFailed(_) + | Self::SchemaFailed(_) + | Self::Irreversible + | Self::ResponseTooLarge + | Self::InvalidResponse(_) + | Self::UnexpectedResponse + | Self::Transport(_) => false, + } + } } diff --git a/litellm-rust/crates/storage-clickhouse/src/insert.rs b/litellm-rust/crates/storage-clickhouse/src/insert.rs index 1cdad1897c1..29837359f0b 100644 --- a/litellm-rust/crates/storage-clickhouse/src/insert.rs +++ b/litellm-rust/crates/storage-clickhouse/src/insert.rs @@ -1,6 +1,7 @@ use std::{io::Write, time::Duration}; use flate2::{Compression, write::GzEncoder}; +use litellm_error::Source; use litellm_http::Client; use crate::{Connection, Error, valid_identifier}; @@ -34,13 +35,15 @@ pub async fn insert_encoded_rows( return Err(Error::InvalidTable); } let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); - encoder - .write_all(encoded.as_bytes()) - .map_err(|_| Error::InvalidRow)?; - let body = encoder.finish().map_err(|_| Error::InvalidRow)?; + encoder.write_all(encoded.as_bytes()).map_err(invalid_row)?; + let body = encoder.finish().map_err(invalid_row)?; insert_compressed_rows(client, connection, database, table, token, body).await } +fn invalid_row(error: std::io::Error) -> Error { + Error::InvalidRow(Source::new(error)) +} + pub async fn insert_compressed_rows( client: &Client, connection: &Connection, @@ -91,7 +94,7 @@ pub async fn insert_compressed_rows( .body(body) .send() .await - .map_err(|_| Error::Transport)?; + .map_err(litellm_http::Error::from)?; if !response.status().is_success() { return Err(Error::InsertFailed(response.status().as_u16())); } diff --git a/litellm-rust/crates/storage-clickhouse/src/lib.rs b/litellm-rust/crates/storage-clickhouse/src/lib.rs index e2ef62d1bc3..60a4e77a36c 100644 --- a/litellm-rust/crates/storage-clickhouse/src/lib.rs +++ b/litellm-rust/crates/storage-clickhouse/src/lib.rs @@ -16,9 +16,9 @@ pub struct Connection { impl Connection { pub fn parse(value: &str) -> Result { - let url = Url::parse(value).map_err(|_| Error::InvalidUrl)?; + let url = Url::parse(value).map_err(Error::InvalidUrl)?; if !matches!(url.scheme(), "http" | "https") || url.host().is_none() { - return Err(Error::InvalidUrl); + return Err(Error::UnsupportedUrl); } Ok(Self { url }) } @@ -33,11 +33,11 @@ impl Connection { connection .url .set_username(user) - .map_err(|_| Error::InvalidUrl)?; + .map_err(|()| Error::UnsupportedUrl)?; connection .url .set_password(Some(password)) - .map_err(|_| Error::InvalidUrl)?; + .map_err(|()| Error::UnsupportedUrl)?; let pairs: Vec<_> = connection .url .query_pairs() diff --git a/litellm-rust/crates/storage-clickhouse/src/migrate.rs b/litellm-rust/crates/storage-clickhouse/src/migrate.rs index c249020faaf..4da26baa681 100644 --- a/litellm-rust/crates/storage-clickhouse/src/migrate.rs +++ b/litellm-rust/crates/storage-clickhouse/src/migrate.rs @@ -12,7 +12,7 @@ use sqlx::{ migrate::{AppliedMigration, Migrate, MigrateError, Migration}, }; -use crate::{Connection, Error, READ_LIMITS, valid_identifier}; +use crate::{Connection, Error, READ_LIMITS, read::invalid_response, valid_identifier}; pub async fn execute_statement( client: &Client, @@ -22,7 +22,7 @@ pub async fn execute_statement( ) -> Result<(), Error> { let body = execute_sql(client, connection, sql, timeout).await?; if !body.trim().is_empty() { - return Err(Error::InvalidResponse); + return Err(Error::UnexpectedResponse); } Ok(()) } @@ -56,18 +56,18 @@ async fn execute_sql( .body(sql.to_owned()) .send() .await - .map_err(|_| Error::Transport)?; + .map_err(litellm_http::Error::from)?; if !response.status().is_success() { return Err(Error::SchemaFailed(response.status().as_u16())); } let mut body = Vec::new(); - while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? { + while let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? { if body.len() + chunk.len() > READ_LIMITS.response_bytes { return Err(Error::ResponseTooLarge); } body.extend_from_slice(&chunk); } - String::from_utf8(body).map_err(|_| Error::InvalidResponse) + String::from_utf8(body).map_err(invalid_response) } /// The startup runner records success after execution, never marks migrations dirty, and skips locks @@ -131,13 +131,13 @@ fn encode_hex(bytes: &[u8]) -> String { fn decode_hex(value: &str) -> Result, Error> { let (pairs, remainder) = value.as_bytes().as_chunks::<2>(); if !remainder.is_empty() { - return Err(Error::InvalidResponse); + return Err(Error::UnexpectedResponse); } pairs .iter() .map(|pair| { - let high = decode_hex_digit(pair[0]).ok_or(Error::InvalidResponse)?; - let low = decode_hex_digit(pair[1]).ok_or(Error::InvalidResponse)?; + let high = decode_hex_digit(pair[0]).ok_or(Error::UnexpectedResponse)?; + let low = decode_hex_digit(pair[1]).ok_or(Error::UnexpectedResponse)?; Ok((high << 4) | low) }) .collect() @@ -231,7 +231,7 @@ where .filter(|line| !line.trim().is_empty()) .map(|line| { let row = serde_json::from_str::(line) - .map_err(|_| migrate_error(Error::InvalidResponse))?; + .map_err(|error| migrate_error(invalid_response(error)))?; let checksum = decode_hex(&row.checksum).map_err(migrate_error)?; Ok(AppliedMigration { version: row.version, @@ -288,11 +288,7 @@ where _table_name: &'e str, _migration: &'e Migration, ) -> MigrateFuture<'e, Result> { - Box::pin(async { - Err(MigrateError::Execute(SqlxError::AnyDriverError(Box::new( - std::io::Error::other("ClickHouse migrations are forward-only"), - )))) - }) + Box::pin(async { Err(migrate_error(Error::Irreversible)) }) } } diff --git a/litellm-rust/crates/storage-clickhouse/src/read.rs b/litellm-rust/crates/storage-clickhouse/src/read.rs index 03a5522e595..40d4ad0a073 100644 --- a/litellm-rust/crates/storage-clickhouse/src/read.rs +++ b/litellm-rust/crates/storage-clickhouse/src/read.rs @@ -1,5 +1,6 @@ use std::{collections::BTreeMap, time::Duration}; +use litellm_error::Source; use litellm_http::Client; use serde::{Deserialize, Serialize, de::DeserializeOwned}; @@ -115,7 +116,7 @@ async fn execute_read_with_limits( .post(url) .timeout(Duration::from_secs(15)) .body(sql.to_owned()); - let mut response = request.send().await.map_err(|_| Error::Transport)?; + let mut response = request.send().await.map_err(litellm_http::Error::from)?; if !response.status().is_success() { if response .headers() @@ -128,20 +129,19 @@ async fn execute_read_with_limits( } let mut body = Vec::new(); - while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? { + while let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? { if body.len() + chunk.len() > limits.response_bytes { return Err(Error::ResponseTooLarge); } body.extend_from_slice(&chunk); } - let json: serde_json::Value = - serde_json::from_slice(&body).map_err(|_| Error::InvalidResponse)?; + let json: serde_json::Value = serde_json::from_slice(&body).map_err(invalid_response)?; if json.get("exception").is_some() || !json.get("data").is_some_and(serde_json::Value::is_array) { - return Err(Error::InvalidResponse); + return Err(Error::UnexpectedResponse); } - String::from_utf8(body).map_err(|_| Error::InvalidResponse) + String::from_utf8(body).map_err(invalid_response) } pub trait Query { @@ -158,8 +158,8 @@ struct Rows { } fn parameters(params: &T) -> Result, Error> { - let value = serde_json::to_value(params).map_err(|_| Error::InvalidParameters)?; - serde_json::from_value(value).map_err(|_| Error::InvalidParameters) + let value = serde_json::to_value(params).map_err(invalid_parameters)?; + serde_json::from_value(value).map_err(invalid_parameters) } pub async fn fetch( @@ -198,5 +198,13 @@ pub async fn fetch_json( fn decode_rows(body: &str) -> Result, Error> { serde_json::from_str::>(body) .map(|rows| rows.data) - .map_err(|_| Error::InvalidResponse) + .map_err(invalid_response) +} + +pub(crate) fn invalid_response(error: impl std::error::Error + Send + Sync + 'static) -> Error { + Error::InvalidResponse(Source::new(error)) +} + +fn invalid_parameters(error: serde_json::Error) -> Error { + Error::InvalidParameters(Source::new(error)) } diff --git a/litellm-rust/crates/storage-clickhouse/tests/migrations.rs b/litellm-rust/crates/storage-clickhouse/tests/migrations.rs index 29a85a8c4ec..09e56766891 100644 --- a/litellm-rust/crates/storage-clickhouse/tests/migrations.rs +++ b/litellm-rust/crates/storage-clickhouse/tests/migrations.rs @@ -48,7 +48,7 @@ async fn database() -> TestResult { Ok(ClickHouseDatabase { _container: container, url, - client: Client::no_redirect_for_test(), + client: Client::for_test(), }) } @@ -244,7 +244,7 @@ async fn failed_migration_is_not_recorded_and_retains_storage_error( #[rstest] fn invalid_database_identifier_is_rejected() { - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer("http://127.0.0.1:1").expect("valid URL"); assert!(matches!( ClickHouseMigrate::new( @@ -327,7 +327,7 @@ async fn duplicate_ledger_rows_are_tolerated( #[rstest] #[tokio::test] async fn revert_reports_forward_only_error() { - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer("http://127.0.0.1:1").expect("valid URL"); let migration = migration( 1, @@ -346,11 +346,7 @@ async fn revert_reports_forward_only_error() { .revert("_sqlx_migrations", &migration) .await .expect_err("ClickHouse migrations cannot be reverted"); - assert!( - error - .to_string() - .contains("ClickHouse migrations are forward-only") - ); + assert!(matches!(storage_error(&error), Some(Error::Irreversible))); } #[rstest] @@ -369,7 +365,7 @@ async fn applied_int64_versions_accept_numeric_and_quoted_json( ) .mount(&mock_server) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer(&mock_server.uri()).expect("valid URL"); let migrator = migrator(vec![]); let mut adapter = ClickHouseMigrate::new( @@ -391,6 +387,7 @@ async fn applied_int64_versions_accept_numeric_and_quoted_json( #[tokio::test] async fn schema_requests_override_unsafe_connection_settings( #[future(awt)] mock_server: MockServer, + #[values("wait_end_of_query", "send_progress_in_http_headers", "async_insert")] name: &str, ) { Mock::given(method("POST")) .and(query_param("wait_end_of_query", "1")) @@ -406,7 +403,7 @@ async fn schema_requests_override_unsafe_connection_settings( mock_server.uri() ); execute_statement( - &Client::no_redirect_for_test(), + &Client::for_test(), &Connection::writer(&url).expect("valid URL"), "CREATE DATABASE IF NOT EXISTS trace_test", REQUEST_TIMEOUT, @@ -417,20 +414,14 @@ async fn schema_requests_override_unsafe_connection_settings( .received_requests() .await .expect("requests recorded"); - for name in [ - "wait_end_of_query", - "send_progress_in_http_headers", - "async_insert", - ] { - assert_eq!( - requests[0] - .url - .query_pairs() - .filter(|(key, _)| key == name) - .count(), - 1 - ); - } + assert_eq!( + requests[0] + .url + .query_pairs() + .filter(|(key, _)| key == name) + .count(), + 1 + ); } #[rstest] @@ -445,7 +436,7 @@ async fn oversized_ledger_response_is_rejected_before_migrations( ) .mount(&mock_server) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer(&mock_server.uri()).expect("valid URL"); let migrator = migrator(vec![]); let mut adapter = ClickHouseMigrate::new( diff --git a/litellm-rust/crates/storage-clickhouse/tests/transport.rs b/litellm-rust/crates/storage-clickhouse/tests/transport.rs index f70fe52e1de..6b2c3fb8ef1 100644 --- a/litellm-rust/crates/storage-clickhouse/tests/transport.rs +++ b/litellm-rust/crates/storage-clickhouse/tests/transport.rs @@ -13,7 +13,7 @@ async fn insert_rejects_invalid_identifiers( #[case] table: &str, #[case] invalid_database: bool, ) { - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::writer("http://localhost:8123").expect("valid URL"); let result = insert_encoded_rows(&client, &connection, database, table, "token", "{}").await; @@ -24,7 +24,7 @@ async fn insert_rejects_invalid_identifiers( #[rstest] #[tokio::test] async fn read_rejects_empty_sql() { - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::reader("http://localhost:8123", "litellm").expect("valid URL"); assert!(matches!( @@ -55,18 +55,24 @@ impl litellm_storage_clickhouse::Query for TypedQuery { const SQL: &'static str = "SELECT typed_parameters"; } +enum Expected { + Rows, + Undecodable, + Failed, +} + #[rstest] #[case::valid( r#"{"meta":[],"data":[{"answer":"ok"}],"rows":1,"statistics":{"elapsed":0.1}}"#, - true + Expected::Rows )] -#[case::wrong_type(r#"{"data":[{"answer":1}]}"#, false)] -#[case::missing_column(r#"{"data":[{}]}"#, false)] -#[case::exception(r#"{"data":[],"exception":"failed"}"#, false)] +#[case::wrong_type(r#"{"data":[{"answer":1}]}"#, Expected::Undecodable)] +#[case::missing_column(r#"{"data":[{}]}"#, Expected::Undecodable)] +#[case::exception(r#"{"data":[],"exception":"failed"}"#, Expected::Failed)] #[tokio::test] async fn typed_fetch_encodes_parameters_and_validates_rows( #[case] body: &str, - #[case] valid: bool, + #[case] expected: Expected, ) { use litellm_storage_clickhouse::{fetch, fetch_json}; use wiremock::{ @@ -87,7 +93,7 @@ async fn typed_fetch_encodes_parameters_and_validates_rows( .expect(2) .mount(&server) .await; - let client = Client::no_redirect_for_test(); + let client = Client::for_test(); let connection = Connection::parse(&server.uri()).unwrap(); let params = QueryParams { signed: i64::MIN, @@ -98,17 +104,24 @@ async fn typed_fetch_encodes_parameters_and_validates_rows( }; let rows = fetch::(&client, &connection, ¶ms).await; let envelope = fetch_json::(&client, &connection, ¶ms).await; - if valid { - assert_eq!( - rows.unwrap(), - vec![QueryRow { - answer: "ok".into() - }] - ); - assert_eq!(envelope.unwrap(), body); - } else { - assert!(matches!(rows, Err(Error::InvalidResponse))); - assert!(matches!(envelope, Err(Error::InvalidResponse))); + match expected { + Expected::Rows => { + assert_eq!( + rows.unwrap(), + vec![QueryRow { + answer: "ok".into() + }] + ); + assert_eq!(envelope.unwrap(), body); + } + Expected::Undecodable => { + assert!(matches!(rows, Err(Error::InvalidResponse(_)))); + assert!(matches!(envelope, Err(Error::InvalidResponse(_)))); + } + Expected::Failed => { + assert!(matches!(rows, Err(Error::UnexpectedResponse))); + assert!(matches!(envelope, Err(Error::UnexpectedResponse))); + } } } @@ -131,7 +144,7 @@ async fn server_result_limits_allow_smaller_pages_without_retrying_other_failure .await; let connection = Connection::parse(&server.uri()).unwrap(); let error = execute_read( - &Client::no_redirect_for_test(), + &Client::for_test(), &connection, "SELECT 1", &BTreeMap::new(), @@ -145,21 +158,23 @@ async fn server_result_limits_allow_smaller_pages_without_retrying_other_failure } } -#[test] -fn insert_timeout_environment_controls_transport() { - for value in ["1", "3", "0", "invalid"] { - let result = std::process::Command::new(std::env::current_exe().unwrap()) - .args(["--exact", "insert_timeout_environment_child"]) - .env("LITELLM_TEST_INSERT_TIMEOUT", value) - .env("CLICKHOUSE_INSERT_TIMEOUT_SECONDS", value) - .output() - .unwrap(); - assert!( - result.status.success(), - "{}", - String::from_utf8_lossy(&result.stdout) - ); - } +#[rstest] +#[case::times_out("1")] +#[case::completes("3")] +#[case::zero("0")] +#[case::not_a_number("invalid")] +fn insert_timeout_environment_controls_transport(#[case] value: &str) { + let result = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "insert_timeout_environment_child"]) + .env("LITELLM_TEST_INSERT_TIMEOUT", value) + .env("CLICKHOUSE_INSERT_TIMEOUT_SECONDS", value) + .output() + .unwrap(); + assert!( + result.status.success(), + "{}", + String::from_utf8_lossy(&result.stdout) + ); } #[tokio::test] @@ -174,7 +189,7 @@ async fn insert_timeout_environment_child() { .mount(&server) .await; let result = insert_encoded_rows( - &Client::no_redirect_for_test(), + &Client::for_test(), &Connection::parse(&server.uri()).unwrap(), "traces", "otel_traces", @@ -183,7 +198,7 @@ async fn insert_timeout_environment_child() { ) .await; match value.as_str() { - "1" => assert!(matches!(result, Err(Error::Transport))), + "1" => assert!(matches!(result, Err(Error::Transport(error)) if error.dispatched())), "3" => assert!(result.is_ok()), _ => assert!(matches!( result, @@ -191,3 +206,29 @@ async fn insert_timeout_environment_child() { )), } } + +#[rstest] +#[tokio::test] +async fn unreachable_server_keeps_the_transport_cause() { + let address = std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap(); + let connection = Connection::parse(&format!("http://{address}")).unwrap(); + let error = execute_read( + &Client::for_test(), + &connection, + "SELECT 1", + &BTreeMap::new(), + ) + .await + .unwrap_err(); + let cause = std::error::Error::source(&error) + .and_then(|cause| cause.downcast_ref::()) + .expect("transport cause"); + assert!(!cause.dispatched()); + assert_ne!( + litellm_error::Report::new(&error).to_string(), + error.to_string() + ); +} diff --git a/litellm-rust/crates/testkit/Cargo.toml b/litellm-rust/crates/testkit/Cargo.toml index 97efb753384..ef3c3e6cd7e 100644 --- a/litellm-rust/crates/testkit/Cargo.toml +++ b/litellm-rust/crates/testkit/Cargo.toml @@ -8,6 +8,7 @@ publish = false [dependencies] flate2.workspace = true +litellm-http.workspace = true reqwest.workspace = true serde.workspace = true semver.workspace = true diff --git a/litellm-rust/crates/testkit/src/agent/configure.rs b/litellm-rust/crates/testkit/src/agent/configure.rs index a41bec9379a..6f1b970e8fb 100644 --- a/litellm-rust/crates/testkit/src/agent/configure.rs +++ b/litellm-rust/crates/testkit/src/agent/configure.rs @@ -5,10 +5,13 @@ use semver::Version; use crate::Error; -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray, strum::Display)] pub enum Wire { + #[strum(serialize = "chat completions")] ChatCompletions, + #[strum(serialize = "messages")] Messages, + #[strum(serialize = "responses")] Responses, } diff --git a/litellm-rust/crates/testkit/src/error.rs b/litellm-rust/crates/testkit/src/error.rs index 03520d827f1..8b747c28bda 100644 --- a/litellm-rust/crates/testkit/src/error.rs +++ b/litellm-rust/crates/testkit/src/error.rs @@ -15,7 +15,7 @@ pub enum Error { Request { url: String, #[source] - source: reqwest::Error, + source: litellm_http::Error, }, #[error("{url} answered with status {status}")] Status { url: String, status: u16 }, @@ -47,10 +47,14 @@ pub enum Error { expected: String, reported: String, }, - #[error("{agent} cannot talk to the gateway over {wire:?}")] + #[error("{agent} cannot talk to the gateway over {wire}")] UnsupportedWire { agent: &'static str, wire: Wire }, - #[error("agent did not finish within {0:?}")] + #[error("agent did not finish within {} seconds", .0.as_secs_f64())] Timeout(std::time::Duration), - #[error("io failure")] - Io(#[from] io::Error), + #[error("could not install the binary into {}", dir.display())] + Install { dir: PathBuf, source: io::Error }, + #[error("could not run {}", binary.display())] + Spawn { binary: PathBuf, source: io::Error }, + #[error("could not write agent files under {}", home.display())] + WriteFiles { home: PathBuf, source: io::Error }, } diff --git a/litellm-rust/crates/testkit/src/install/fetch.rs b/litellm-rust/crates/testkit/src/install/fetch.rs index 73008f7a0da..59a61bfaa0a 100644 --- a/litellm-rust/crates/testkit/src/install/fetch.rs +++ b/litellm-rust/crates/testkit/src/install/fetch.rs @@ -38,9 +38,9 @@ impl Fetch for HttpFetch { (Some(token), true) => request.bearer_auth(token), _ => request, }; - let request_error = |source| Error::Request { + let request_error = |source: reqwest::Error| Error::Request { url: url.to_owned(), - source, + source: source.into(), }; let response = request.send().await.map_err(request_error)?; let status = response.status(); diff --git a/litellm-rust/crates/testkit/src/install/mod.rs b/litellm-rust/crates/testkit/src/install/mod.rs index 1104bcec102..73f7deec864 100644 --- a/litellm-rust/crates/testkit/src/install/mod.rs +++ b/litellm-rust/crates/testkit/src/install/mod.rs @@ -52,7 +52,13 @@ impl Installer { version: version.clone(), binary: binary.clone(), }; - if fs::try_exists(&binary).await? && probe_version(&binary, version).await.is_ok() { + let install_error = |source| Error::Install { + dir: dir.clone(), + source, + }; + if fs::try_exists(&binary).await.map_err(install_error)? + && probe_version(&binary, version).await.is_ok() + { return Ok(installed); } @@ -61,21 +67,23 @@ impl Installer { verify_sha256(&release.asset, &release.sha256, &archive)?; let contents = extract_binary(&release.packaging, &archive)?; - fs::create_dir_all(&dir).await?; + fs::create_dir_all(&dir).await.map_err(install_error)?; let staging = dir.join(format!( ".{}.{}.{}.partial", agent.binary(), std::process::id(), STAGING_COUNTER.fetch_add(1, Ordering::Relaxed) )); - fs::write(&staging, contents).await?; - fs::set_permissions(&staging, std::fs::Permissions::from_mode(0o755)).await?; - fs::rename(&staging, &binary).await?; + fs::write(&staging, contents).await.map_err(install_error)?; + fs::set_permissions(&staging, std::fs::Permissions::from_mode(0o755)) + .await + .map_err(install_error)?; + fs::rename(&staging, &binary).await.map_err(install_error)?; match probe_version(&binary, version).await { Ok(()) => Ok(installed), Err(error) => { - fs::remove_file(&binary).await?; + fs::remove_file(&binary).await.map_err(install_error)?; Err(error) } } @@ -98,7 +106,11 @@ async fn probe_version(binary: &Path, expected: &Version) -> Result<(), Error> { .env("DISABLE_AUTOUPDATER", "1") .stdin(Stdio::null()) .output() - .await?; + .await + .map_err(|source| Error::Spawn { + binary: binary.to_owned(), + source, + })?; let stdout = String::from_utf8_lossy(&output.stdout); if stdout .split_whitespace() diff --git a/litellm-rust/crates/testkit/src/session.rs b/litellm-rust/crates/testkit/src/session.rs index 6b06e756cca..cd5eb08ae2b 100644 --- a/litellm-rust/crates/testkit/src/session.rs +++ b/litellm-rust/crates/testkit/src/session.rs @@ -28,7 +28,11 @@ impl Session { ) -> Result { let home = home.into(); let spec = agent.configure(&installed.version, &settings, &home)?; - spec.write_files(&home)?; + spec.write_files(&home) + .map_err(|source| Error::WriteFiles { + home: home.clone(), + source, + })?; Ok(Self { binary: installed.binary.clone(), home, @@ -55,7 +59,11 @@ impl Session { .output(); let output = timeout(limit, child) .await - .map_err(|_| Error::Timeout(limit))??; + .map_err(|_| Error::Timeout(limit))? + .map_err(|source| Error::Spawn { + binary: self.binary.clone(), + source, + })?; let parsed = agent.parse(&self.version, &String::from_utf8_lossy(&output.stdout)); let failed_silently = !output.status.success() && parsed.errors.is_empty(); Ok(Outcome { diff --git a/litellm-rust/crates/token-counter-fast/Cargo.toml b/litellm-rust/crates/token-counter-fast/Cargo.toml index 5127dc17f22..070bec58519 100644 --- a/litellm-rust/crates/token-counter-fast/Cargo.toml +++ b/litellm-rust/crates/token-counter-fast/Cargo.toml @@ -13,6 +13,7 @@ tokenizers.workspace = true unicode-normalization-alignments = "0.1.12" [dev-dependencies] +litellm-error.workspace = true rand.workspace = true rstest.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/token-counter-fast/src/error.rs b/litellm-rust/crates/token-counter-fast/src/error.rs index d2a706191ac..2cd8d5c810c 100644 --- a/litellm-rust/crates/token-counter-fast/src/error.rs +++ b/litellm-rust/crates/token-counter-fast/src/error.rs @@ -2,16 +2,25 @@ use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { - #[error("failed to load tokenizer: {0}")] + #[error("failed to load tokenizer")] Load(#[source] tokenizers::Error), - #[error("failed to load tokenizer: tiktoken rank file: {0}")] + #[error("failed to load tokenizer from a tiktoken rank file")] Ranks(#[source] RankError), #[error("failed to load tokenizer: Unicode character classes are unavailable")] UnicodeClasses, - #[error("tokenization failed: {0}")] + #[error("tokenization failed")] Encode(#[source] tokenizers::Error), } +impl Error { + pub fn is_load(&self) -> bool { + match self { + Self::Load(_) | Self::Ranks(_) | Self::UnicodeClasses => true, + Self::Encode(_) => false, + } + } +} + #[derive(Debug, ThisError)] pub enum RankError { #[error("rank {0} is reserved")] @@ -20,8 +29,8 @@ pub enum RankError { MissingByte(u8), #[error("line without a rank: {0:?}")] MissingRank(String), - #[error("token is not base64: {0}")] + #[error("token is not base64")] InvalidToken(#[source] base64::DecodeError), - #[error("rank is not an integer: {0}")] + #[error("rank is not an integer")] InvalidRank(#[source] std::num::ParseIntError), } diff --git a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs index f2d4de38f8c..11182e84e8d 100644 --- a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs @@ -254,8 +254,12 @@ mod tests { .err() .expect("invalid rank file"); assert!(matches!(error, Error::Ranks(_))); - assert!(error.to_string().starts_with(&format!( - "failed to load tokenizer: tiktoken rank file: {detail}" - ))); + assert!( + litellm_error::Report::new(&error) + .to_string() + .starts_with(&format!( + "failed to load tokenizer from a tiktoken rank file: {detail}" + )) + ); } } diff --git a/litellm-rust/crates/token-counter-huggingface/src/error.rs b/litellm-rust/crates/token-counter-huggingface/src/error.rs index e7f6260321b..b4e09db253f 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/error.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/error.rs @@ -2,10 +2,21 @@ use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { - #[error("failed to load tokenizer: {0}")] + #[error("failed to load tokenizer")] Load(#[source] tokenizers::Error), - #[error("tokenization failed: {0}")] + #[error("invalid encoding JSON")] + EncodingJson(#[source] serde_json::Error), + #[error("tokenization failed")] Encode(#[source] tokenizers::Error), - #[error("token decoding failed: {0}")] + #[error("token decoding failed")] Decode(#[source] tokenizers::Error), } + +impl Error { + pub fn is_load(&self) -> bool { + match self { + Self::Load(_) | Self::EncodingJson(_) => true, + Self::Encode(_) | Self::Decode(_) => false, + } + } +} diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs index 0acea0532f8..c3af9b42767 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/lib.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -13,11 +13,11 @@ pub use tokenizers::{ }; pub fn encoding_from_json(json: &str) -> Result { - serde_json::from_str(json).map_err(|error| Error::Load(error.into())) + serde_json::from_str(json).map_err(Error::EncodingJson) } pub fn encoding_to_json(encoding: &Encoding) -> Result { - serde_json::to_string(encoding).map_err(|error| Error::Load(error.into())) + serde_json::to_string(encoding).map_err(Error::EncodingJson) } pub struct HuggingFaceTokenizer { diff --git a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml index d1789a693a3..bca7345dcd4 100644 --- a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml +++ b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] base64.workspace = true +litellm-error.workspace = true once_cell = "1.21.3" rustc-hash = "2.1.3" thiserror.workspace = true diff --git a/litellm-rust/crates/token-counter-tiktoken/src/error.rs b/litellm-rust/crates/token-counter-tiktoken/src/error.rs index e28cbbfb620..024be00f048 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/error.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/error.rs @@ -1,5 +1,50 @@ +use std::{io, num::ParseIntError}; + +use litellm_error::Source; use thiserror::Error as ThisError; +use tiktoken_rs::{DecodeKeyError, EncodeError}; #[derive(Debug, ThisError)] -#[error("unsupported tokenizer: {0}")] -pub struct UnsupportedTokenizer(pub String); +pub enum Error { + #[error("unsupported tokenizer: {0}")] + UnsupportedTokenizer(String), + #[error("could not read tiktoken rank file {file}")] + Read { + file: &'static str, + source: io::Error, + }, + #[error("invalid tiktoken rank file")] + Ranks(#[from] RankError), + #[error("could not build the tiktoken encoder")] + Encoder(#[source] Source), + #[error("tokenization failed")] + Encode(#[source] EncodeError), + #[error("token decoding failed")] + Decode(#[source] DecodeKeyError), +} + +impl Error { + pub fn is_load(&self) -> bool { + match self { + Self::UnsupportedTokenizer(_) + | Self::Read { .. } + | Self::Ranks(_) + | Self::Encoder(_) => true, + Self::Encode(_) | Self::Decode(_) => false, + } + } +} + +#[derive(Debug, ThisError)] +pub enum RankError { + #[error("line without a rank: {0:?}")] + MissingRank(String), + #[error("token is not base64")] + InvalidToken(#[source] base64::DecodeError), + #[error("rank is not an integer")] + InvalidRank(#[source] ParseIntError), + #[error("a rank is assigned to more than one token")] + DuplicateRank, + #[error("byte 0x{0:02X} has no token")] + MissingByte(u8), +} diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs index 30d0f820da3..c4cb786d9e4 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -5,8 +5,8 @@ mod ranks; use std::collections::HashSet; -pub use error::UnsupportedTokenizer; -pub use ranks::{LoadError, Vocabulary}; +pub use error::{Error, RankError}; +pub use ranks::Vocabulary; pub struct TiktokenTokenizer { encoder: &'static tiktoken_rs::CoreBPE, @@ -22,7 +22,7 @@ impl TiktokenTokenizer { pub fn from_cached_ranks( name: &str, load: impl FnOnce(&str) -> std::io::Result, - ) -> Result { + ) -> Result { let (loaded, name) = ranks::load(name, load)?; Ok(Self { encoder: &loaded.bpe, @@ -32,7 +32,7 @@ impl TiktokenTokenizer { } /// The encodings tiktoken-rs embeds, for hosts without the packaged rank files. - pub fn from_name(name: &str) -> Result { + pub fn from_name(name: &str) -> Result { let (encoder, name) = match name { "cl100k_base" => (tiktoken_rs::cl100k_base_singleton(), "cl100k_base"), "o200k_base" => (tiktoken_rs::o200k_base_singleton(), "o200k_base"), @@ -41,7 +41,7 @@ impl TiktokenTokenizer { "p50k_edit" => (tiktoken_rs::p50k_edit_singleton(), "p50k_edit"), "r50k_base" => (tiktoken_rs::r50k_base_singleton(), "r50k_base"), "gpt2" => (tiktoken_rs::r50k_base_singleton(), "gpt2"), - _ => return Err(UnsupportedTokenizer(name.to_owned())), + _ => return Err(Error::UnsupportedTokenizer(name.to_owned())), }; Ok(Self { encoder, @@ -62,12 +62,12 @@ impl TiktokenTokenizer { self.encoder.encode_ordinary(text) } - pub fn encode_special(&self, text: &str, allowed: &[String]) -> Result, String> { + pub fn encode_special(&self, text: &str, allowed: &[String]) -> Result, Error> { let allowed = allowed.iter().map(String::as_str).collect(); self.encoder .encode(text, &allowed) .map(|(ids, _)| ids) - .map_err(|error| error.to_string()) + .map_err(Error::Encode) } pub fn special_tokens(&self) -> HashSet { @@ -92,17 +92,13 @@ impl TiktokenTokenizer { (stable, completions) } - pub fn decode_bytes(&self, ids: &[u32]) -> Result, String> { - self.encoder - .decode_bytes(ids) - .map_err(|error| error.to_string()) + pub fn decode_bytes(&self, ids: &[u32]) -> Result, Error> { + self.encoder.decode_bytes(ids).map_err(Error::Decode) } - pub fn decode(&self, ids: &[u32]) -> Result { - self.encoder - .decode_bytes(ids) + pub fn decode(&self, ids: &[u32]) -> Result { + self.decode_bytes(ids) .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) - .map_err(|error| error.to_string()) } pub fn name(&self) -> &str { @@ -161,7 +157,8 @@ mod tests { #[test] fn unsupported_encoding_preserves_its_name() { - let Err(UnsupportedTokenizer(name)) = TiktokenTokenizer::from_name("unknown-encoding") + let Err(Error::UnsupportedTokenizer(name)) = + TiktokenTokenizer::from_name("unknown-encoding") else { panic!("unknown encoding must be rejected"); }; diff --git a/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs index 192e6f65a0b..36c30e711c1 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs @@ -1,10 +1,10 @@ use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_error::Source; use once_cell::sync::OnceCell; use rustc_hash::FxHashMap; -use thiserror::Error; use tiktoken_rs::{CoreBPE, O200K_BASE_PAT_STR, Rank}; -use crate::UnsupportedTokenizer; +use crate::error::{Error, RankError}; const CL100K: &str = "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"; const O200K: &str = "fb374d419588a4632f3f557e76b4b70aebbca790"; @@ -76,22 +76,30 @@ impl Vocabulary { } } -#[derive(Debug, Error)] -pub enum LoadError { - #[error(transparent)] - Unsupported(#[from] UnsupportedTokenizer), - #[error("failed to load tiktoken ranks: {0}")] - Ranks(String), -} - /// Loads `name` once per process. The returned name is the one requested (`gpt2` stays /// `gpt2`, as `tiktoken.get_encoding("gpt2").name` does), while `gpt2` and `r50k_base` share /// one cached encoder. pub(super) fn load( name: &str, load_file: impl FnOnce(&str) -> std::io::Result, -) -> Result<(&'static Loaded, &'static str), LoadError> { - let (requested, canonical, file, cache) = match name { +) -> Result<(&'static Loaded, &'static str), Error> { + let (requested, canonical, file, cache) = encoding(name)?; + let loaded = cache.get_or_try_init(|| { + let ranks = load_file(file).map_err(|source| Error::Read { file, source })?; + build(canonical, &ranks) + })?; + Ok((loaded, requested)) +} + +type Encoding = ( + &'static str, + &'static str, + &'static str, + &'static OnceCell, +); + +fn encoding(name: &str) -> Result { + Ok(match name { "cl100k_base" => ("cl100k_base", "cl100k_base", CL100K, &CL100K_ENCODER), "o200k_base" => ("o200k_base", "o200k_base", O200K, &O200K_ENCODER), "o200k_harmony" => ("o200k_harmony", "o200k_harmony", O200K, &HARMONY_ENCODER), @@ -99,16 +107,11 @@ pub(super) fn load( "p50k_edit" => ("p50k_edit", "p50k_edit", P50K, &EDIT_ENCODER), "r50k_base" => ("r50k_base", "r50k_base", P50K, &R50K_ENCODER), "gpt2" => ("gpt2", "r50k_base", P50K, &R50K_ENCODER), - _ => return Err(UnsupportedTokenizer(name.to_owned()).into()), - }; - let loaded = cache.get_or_try_init(|| { - let ranks = load_file(file).map_err(|error| LoadError::Ranks(error.to_string()))?; - build(canonical, &ranks) - })?; - Ok((loaded, requested)) + _ => return Err(Error::UnsupportedTokenizer(name.to_owned())), + }) } -fn build(name: &str, ranks: &str) -> Result { +fn build(name: &str, ranks: &str) -> Result { let parsed = ranks .lines() .map(parse_rank) @@ -122,9 +125,11 @@ fn build(name: &str, ranks: &str) -> Result { .collect::>() .len() != encoder.len() - || (0..=u8::MAX).any(|byte| !encoder.contains_key(&[byte][..])) { - return Err(LoadError::Ranks("invalid vocabulary ranks".into())); + return Err(RankError::DuplicateRank.into()); + } + if let Some(byte) = (0..=u8::MAX).find(|byte| !encoder.contains_key(&[*byte][..])) { + return Err(RankError::MissingByte(byte).into()); } let (pattern, specials): (&str, &[(&str, Rank)]) = match name { "cl100k_base" => ( @@ -185,9 +190,12 @@ fn build(name: &str, ranks: &str) -> Result { .chain(special_tokens.values()) .copied() .max() - .ok_or_else(|| LoadError::Ranks("empty vocabulary".into()))?; - let bpe = CoreBPE::new(encoder.clone(), special_tokens.clone(), pattern) - .map_err(|error| LoadError::Ranks(error.to_string()))?; + .unwrap_or_default(); + let bpe = CoreBPE::new(encoder.clone(), special_tokens.clone(), pattern).map_err(|error| { + Error::Encoder(Source::from( + Box::::from(error), + )) + })?; Ok(Loaded { bpe, vocabulary: Vocabulary { @@ -198,16 +206,12 @@ fn build(name: &str, ranks: &str) -> Result { }) } -fn parse_rank(line: &str) -> Result<(Vec, Rank), LoadError> { +fn parse_rank(line: &str) -> Result<(Vec, Rank), RankError> { let (token, rank) = line .split_once(' ') - .ok_or_else(|| LoadError::Ranks("missing rank".into()))?; - let bytes = STANDARD - .decode(token) - .map_err(|error| LoadError::Ranks(error.to_string()))?; - let rank = rank - .parse() - .map_err(|error: std::num::ParseIntError| LoadError::Ranks(error.to_string()))?; + .ok_or_else(|| RankError::MissingRank(line.to_owned()))?; + let bytes = STANDARD.decode(token).map_err(RankError::InvalidToken)?; + let rank = rank.parse().map_err(RankError::InvalidRank)?; Ok((bytes, rank)) } @@ -226,79 +230,117 @@ mod tests { ) } - #[test] - fn packaged_encodings_match_embedded_encodings_and_reuse_successful_loads() { - for name in [ + #[rstest] + #[case::cl100k_base("cl100k_base", None)] + #[case::o200k_base("o200k_base", None)] + #[case::o200k_harmony("o200k_harmony", None)] + #[case::p50k_base("p50k_base", None)] + #[case::p50k_edit("p50k_edit", None)] + #[case::r50k_base_shared_with_gpt2("r50k_base", Some("gpt2"))] + fn packaged_encodings_load_once_and_match_embedded_encodings( + #[case] name: &str, + #[case] alias: Option<&str>, + ) { + assert!( + TiktokenTokenizer::from_cached_ranks(name, |_| { + Err(std::io::Error::other("unreadable vocabulary")) + }) + .is_err() + ); + let loads = std::sync::atomic::AtomicUsize::new(0); + let barrier = std::sync::Barrier::new(4); + let encoders = std::thread::scope(|scope| { + let tasks: Vec<_> = (0..4) + .map(|_| { + scope.spawn(|| { + barrier.wait(); + TiktokenTokenizer::from_cached_ranks(name, |file| { + loads.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + read_packaged_ranks(file) + }) + .unwrap() + }) + }) + .collect(); + tasks + .into_iter() + .map(|task| task.join().unwrap()) + .collect::>() + }); + assert_eq!(loads.into_inner(), 1); + let actual = &encoders[0]; + let expected = TiktokenTokenizer::from_name(name).unwrap(); + assert_eq!(actual.special_tokens(), expected.special_tokens()); + let specials: Vec<_> = expected.special_tokens().into_iter().collect(); + let special_text = specials.join(" "); + assert_eq!( + actual.encode_special(&special_text, &specials).unwrap(), + expected.encode_special(&special_text, &specials).unwrap() + ); + assert_eq!(actual.encode(&special_text), expected.encode(&special_text)); + let cached = + TiktokenTokenizer::from_cached_ranks(name, |_| panic!("reloaded cached vocabulary")) + .unwrap(); + assert_eq!(cached.encode("cached"), expected.encode("cached")); + assert_eq!(cached.name(), name); + assert!(expected.vocabulary().is_none()); + assert_vocabulary_lookups(name, actual); + if let Some(alias) = alias { + let shared = TiktokenTokenizer::from_cached_ranks(alias, |_| { + panic!("reloaded the shared vocabulary") + }) + .unwrap(); + assert_eq!(shared.name(), alias); + assert_vocabulary_lookups(alias, &shared); + } + } + + fn uncached(name: &str) -> TiktokenTokenizer { + static BUILT: std::sync::Mutex> = + std::sync::Mutex::new(Vec::new()); + let (requested, canonical, file, _) = encoding(name).unwrap(); + let mut built = BUILT.lock().unwrap(); + let loaded = match built.iter().find(|(built, _)| *built == requested) { + Some((_, loaded)) => *loaded, + None => { + let ranks = read_packaged_ranks(file).unwrap(); + let loaded: &'static Loaded = + Box::leak(Box::new(build(canonical, &ranks).unwrap())); + built.push((requested, loaded)); + loaded + } + }; + TiktokenTokenizer { + encoder: &loaded.bpe, + vocabulary: Some(&loaded.vocabulary), + name: requested, + } + } + + #[rstest] + fn packaged_ranks_encode_like_embedded_encodings( + #[values( "cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base", - "gpt2", - ] { - if name != "gpt2" { - assert!( - TiktokenTokenizer::from_cached_ranks(name, |_| { - Err(std::io::Error::other("unreadable vocabulary")) - }) - .is_err() - ); - } - let loads = std::sync::atomic::AtomicUsize::new(0); - let barrier = std::sync::Barrier::new(4); - let encoders = std::thread::scope(|scope| { - let tasks: Vec<_> = (0..4) - .map(|_| { - scope.spawn(|| { - barrier.wait(); - TiktokenTokenizer::from_cached_ranks(name, |file| { - loads.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - read_packaged_ranks(file) - }) - .unwrap() - }) - }) - .collect(); - tasks - .into_iter() - .map(|task| task.join().unwrap()) - .collect::>() - }); - assert_eq!(loads.into_inner(), usize::from(name != "gpt2")); - let actual = &encoders[0]; - let expected = TiktokenTokenizer::from_name(name).unwrap(); - assert_eq!(actual.special_tokens(), expected.special_tokens()); - let specials: Vec<_> = expected.special_tokens().into_iter().collect(); - let special_text = specials.join(" "); - assert_eq!( - actual.encode_special(&special_text, &specials).unwrap(), - expected.encode_special(&special_text, &specials).unwrap() - ); - for text in [ - "", - "café 漢字 ع 🙂", - "a\r\nb\t ", - " hello 123456789", - &special_text, - ] { - let ids = expected.encode(text); - assert_eq!(actual.encode(text), ids, "{name}: {text:?}"); - assert_eq!(actual.count_tokens(text), ids.len(), "{name}: {text:?}"); - assert_eq!( - actual.decode_bytes(&ids).unwrap(), - expected.decode_bytes(&ids).unwrap() - ); - } - let cached = TiktokenTokenizer::from_cached_ranks(name, |_| { - panic!("reloaded cached vocabulary") - }) - .unwrap(); - assert_eq!(cached.encode("cached"), expected.encode("cached")); - assert_eq!(cached.name(), name); - assert!(expected.vocabulary().is_none()); - assert_vocabulary_lookups(name, actual); - } + "gpt2" + )] + name: &str, + #[values("", "café 漢字 ع 🙂", "a\r\nb\t ", " hello 123456789")] text: &str, + ) { + let actual = uncached(name); + let expected = TiktokenTokenizer::from_name(name).unwrap(); + + let ids = expected.encode(text); + assert_eq!(actual.encode(text), ids); + assert_eq!(actual.count_tokens(text), ids.len()); + assert_eq!( + actual.decode_bytes(&ids).unwrap(), + expected.decode_bytes(&ids).unwrap() + ); } /// The token-level lookups tiktoken's Python `Encoding` exposes, checked against the diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml index 67e5bd60537..01ce4b6bac8 100644 --- a/litellm-rust/crates/token-counter/Cargo.toml +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -14,6 +14,7 @@ tiktoken = ["dep:litellm-token-counter-tiktoken"] [dependencies] indexmap = { version = "2.14.0", features = ["serde"] } itoa = "1.0" +litellm-error.workspace = true litellm-token-counter-fast = { workspace = true, optional = true } litellm-token-counter-huggingface = { workspace = true, optional = true } litellm-token-counter-tiktoken = { workspace = true, optional = true } diff --git a/litellm-rust/crates/token-counter/src/counter.rs b/litellm-rust/crates/token-counter/src/counter.rs index 6370d9d74f8..f3c16f93b84 100644 --- a/litellm-rust/crates/token-counter/src/counter.rs +++ b/litellm-rust/crates/token-counter/src/counter.rs @@ -1,12 +1,12 @@ use serde::Serialize; -use crate::Error; use crate::python_json; use crate::tools::format_function_definitions; use crate::types::{ ContentItem, CountableContentBlock, CountableRequest, Message, MessageContent, TextValue, ToolChoice, ToolDefinition, }; +use crate::{Error, Unsupported}; const TOKENS_PER_MESSAGE: usize = 3; const TOKENS_PER_NAME: usize = 1; @@ -53,7 +53,7 @@ impl TokenCounter { self.count_optional_text_value(request.query.as_ref())? + self.count_optional_text_value(request.documents.as_ref())? } else { - return Err(Error::MissingInput); + return Err(Unsupported::MissingInput.into()); }; Ok(InputTokenCount { model: request.model.clone(), @@ -96,7 +96,7 @@ impl TokenCounter { TextValue::Number(number) => match (number.as_i64(), number.as_u64()) { (Some(number), _) => self.count_text(itoa::Buffer::new().format(number)), (_, Some(number)) => self.count_text(itoa::Buffer::new().format(number)), - _ => Err(Error::FloatText), + _ => Err(Unsupported::FloatText.into()), }, TextValue::Text(text) => self.count_text(text), TextValue::List(items) => items @@ -143,7 +143,9 @@ impl TokenCounter { None => Ok(0), } } - ContentItem::Block(CountableContentBlock::Unsupported) => Err(Error::ContentBlock), + ContentItem::Block(CountableContentBlock::Unsupported) => { + Err(Unsupported::ContentBlock.into()) + } } } diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs index 7a74138c074..89a97423c48 100644 --- a/litellm-rust/crates/token-counter/src/error.rs +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -1,39 +1,66 @@ use std::string::FromUtf8Error; +use litellm_error::Source; use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { - #[error("unsupported tokenizer: {0}")] - UnsupportedTokenizer(String), - #[error("failed to load tokenizer: {0}")] - Load(#[source] Box), - #[error("failed to load tokenizer: tiktoken rank file: {0}")] - Ranks(#[source] Box), - #[error("failed to load tokenizer: Unicode character classes are unavailable")] - UnicodeClasses, - #[error("unsupported by the rust token counter: request body could not be parsed: {0}")] - RequestParse(#[source] serde_json::Error), - #[error("unsupported by the rust token counter: request has no countable input")] - MissingInput, - #[error( - "unsupported by the rust token counter: float text values are counted by the python path" - )] - FloatText, - #[error( - "unsupported by the rust token counter: content block type is counted by the python path" - )] - ContentBlock, - #[error("unsupported by the rust token counter: array parameter without items")] - ArrayItems, - #[error("unsupported by the rust token counter: text value could not be serialized: {0}")] - JsonSerialization(#[source] serde_json::Error), - #[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")] - JsonUtf8(#[source] FromUtf8Error), - #[error("tokenization failed: {0}")] - Encode(#[source] Box), - #[error("token decoding failed: {0}")] - Decode(String), - #[error("token counting task failed: {0}")] - Task(String), + #[cfg(feature = "fast")] + #[error(transparent)] + Fast(#[from] litellm_token_counter_fast::Error), + #[cfg(feature = "huggingface")] + #[error(transparent)] + HuggingFace(#[from] litellm_token_counter_huggingface::Error), + #[cfg(feature = "tiktoken")] + #[error(transparent)] + Tiktoken(#[from] litellm_token_counter_tiktoken::Error), + #[error("unsupported by the rust token counter")] + Unsupported(#[from] Unsupported), + #[error("token counting task failed")] + Task(#[source] Source), +} + +impl Error { + pub fn is_unsupported(&self) -> bool { + match self { + Self::Unsupported(_) => true, + #[cfg(feature = "fast")] + Self::Fast(_) => false, + #[cfg(feature = "huggingface")] + Self::HuggingFace(_) => false, + #[cfg(feature = "tiktoken")] + Self::Tiktoken(_) => false, + Self::Task(_) => false, + } + } + + pub fn is_load(&self) -> bool { + match self { + #[cfg(feature = "fast")] + Self::Fast(error) => error.is_load(), + #[cfg(feature = "huggingface")] + Self::HuggingFace(error) => error.is_load(), + #[cfg(feature = "tiktoken")] + Self::Tiktoken(error) => error.is_load(), + Self::Unsupported(_) | Self::Task(_) => false, + } + } +} + +#[derive(Debug, ThisError)] +pub enum Unsupported { + #[error("request body could not be parsed")] + RequestParse(#[source] serde_json::Error), + #[error("request has no countable input")] + MissingInput, + #[error("float text values are counted by the python path")] + FloatText, + #[error("content block type is counted by the python path")] + ContentBlock, + #[error("array parameter without items")] + ArrayItems, + #[error("text value could not be serialized")] + JsonSerialization(#[source] serde_json::Error), + #[error("serialized text value is not UTF-8")] + JsonUtf8(#[source] FromUtf8Error), } diff --git a/litellm-rust/crates/token-counter/src/fast.rs b/litellm-rust/crates/token-counter/src/fast.rs index d1f270dcdc9..3ef0021e355 100644 --- a/litellm-rust/crates/token-counter/src/fast.rs +++ b/litellm-rust/crates/token-counter/src/fast.rs @@ -1,4 +1,3 @@ -use litellm_token_counter_fast::Error as BackendError; pub use litellm_token_counter_fast::FastTokenizer; use crate::{Error, TextCodec, TokenCounter, Tokenizer}; @@ -53,17 +52,6 @@ impl Tokenizer for FastTokenizer { } } -impl From for Error { - fn from(error: BackendError) -> Self { - match error { - BackendError::Load(source) => Self::Load(source), - BackendError::Ranks(source) => Self::Ranks(source.into()), - BackendError::UnicodeClasses => Self::UnicodeClasses, - BackendError::Encode(source) => Self::Encode(source), - } - } -} - #[cfg(all(test, feature = "huggingface", feature = "tiktoken"))] mod tests { use super::*; diff --git a/litellm-rust/crates/token-counter/src/huggingface.rs b/litellm-rust/crates/token-counter/src/huggingface.rs index 43613fac8ea..5730e9ce873 100644 --- a/litellm-rust/crates/token-counter/src/huggingface.rs +++ b/litellm-rust/crates/token-counter/src/huggingface.rs @@ -1,4 +1,3 @@ -use litellm_token_counter_huggingface::Error as BackendError; pub use litellm_token_counter_huggingface::{ AddedToken, EncodeInput, Encoding, HuggingFaceTokenizer, InputSequence, PaddingDirection, PaddingParams, PaddingStrategy, TruncationDirection, TruncationParams, encoding_from_json, @@ -34,13 +33,3 @@ impl TextCodec for HuggingFaceTokenizer { HuggingFaceTokenizer::name(self) } } - -impl From for Error { - fn from(error: BackendError) -> Self { - match error { - BackendError::Load(source) => Self::Load(source), - BackendError::Encode(source) => Self::Encode(source), - BackendError::Decode(source) => Self::Decode(source.to_string()), - } - } -} diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs index 84bf35ca2ba..c58b8cf372f 100644 --- a/litellm-rust/crates/token-counter/src/lib.rs +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -19,6 +19,6 @@ pub mod huggingface; pub mod tiktoken; pub use counter::{InputTokenCount, TokenCounter}; -pub use error::Error; +pub use error::{Error, Unsupported}; pub use tokenizer::{TextCodec, Tokenizer}; pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/python_json.rs b/litellm-rust/crates/token-counter/src/python_json.rs index 00a55ad300c..b5c074a24c7 100644 --- a/litellm-rust/crates/token-counter/src/python_json.rs +++ b/litellm-rust/crates/token-counter/src/python_json.rs @@ -6,8 +6,8 @@ use std::io::{self, Write}; use serde::Serialize; use serde_json::ser::{Formatter, Serializer}; -use super::Error; use super::types::TextValue; +use super::{Error, Unsupported}; pub(super) fn dumps(value: &TextValue) -> Result { let mut output = Vec::with_capacity(serialized_len(value)?); @@ -16,9 +16,9 @@ pub(super) fn dumps(value: &TextValue) -> Result { &mut output, PythonFormatter, )) - .map_err(Error::JsonSerialization)?; + .map_err(Unsupported::JsonSerialization)?; debug_assert_eq!(output.len(), output.capacity()); - String::from_utf8(output).map_err(Error::JsonUtf8) + String::from_utf8(output).map_err(|error| Unsupported::JsonUtf8(error).into()) } fn serialized_len(value: &TextValue) -> Result { @@ -29,7 +29,7 @@ fn serialized_len(value: &TextValue) -> Result { TextValue::Number(number) => match (number.as_i64(), number.as_u64()) { (Some(number), _) => Ok(unsigned_len(number.unsigned_abs()) + usize::from(number < 0)), (_, Some(number)) => Ok(unsigned_len(number)), - _ => Err(Error::FloatText), + _ => Err(Unsupported::FloatText.into()), }, TextValue::Text(text) => Ok(quoted_len(text)), TextValue::List(items) => items @@ -147,9 +147,9 @@ mod tests { let value = serde_json::from_str(input).expect("fixture parses"); let error = dumps(&value).expect_err("floats are declined"); - assert_eq!( - error.to_string(), - "unsupported by the rust token counter: float text values are counted by the python path" - ); + assert!(matches!( + error, + crate::Error::Unsupported(crate::Unsupported::FloatText) + )); } } diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index b95621f3ccc..136280b6115 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -1,4 +1,3 @@ -use litellm_token_counter_tiktoken::{LoadError, UnsupportedTokenizer}; pub use litellm_token_counter_tiktoken::{TiktokenTokenizer, Vocabulary, encoding_for_model}; use crate::{Error, TextCodec, TokenCounter, Tokenizer}; @@ -23,25 +22,10 @@ impl TextCodec for TiktokenTokenizer { } fn decode(&self, ids: &[u32], _skip_special_tokens: bool) -> Result { - TiktokenTokenizer::decode(self, ids).map_err(|error| Error::Decode(error.to_string())) + TiktokenTokenizer::decode(self, ids).map_err(Error::from) } fn name(&self) -> &str { TiktokenTokenizer::name(self) } } - -impl From for Error { - fn from(error: UnsupportedTokenizer) -> Self { - Self::UnsupportedTokenizer(error.0) - } -} - -impl From for Error { - fn from(error: LoadError) -> Self { - match error { - LoadError::Unsupported(error) => error.into(), - LoadError::Ranks(source) => Self::Ranks(source.into()), - } - } -} diff --git a/litellm-rust/crates/token-counter/src/tools.rs b/litellm-rust/crates/token-counter/src/tools.rs index 22150a4e4b1..a5a3fdbfabf 100644 --- a/litellm-rust/crates/token-counter/src/tools.rs +++ b/litellm-rust/crates/token-counter/src/tools.rs @@ -3,8 +3,8 @@ use std::fmt::Write; -use super::Error; use super::types::{EnumValue, Schema, SchemaType, ToolDefinition}; +use super::{Error, Unsupported}; pub(super) fn format_function_definitions(tools: &[ToolDefinition]) -> Result { ToolFormatter::new(tools.len()).format(tools) @@ -117,7 +117,7 @@ impl ToolFormatter { None => self.output.push_str("number"), }, "array" => { - let items = props.items.as_deref().ok_or(Error::ArrayItems)?; + let items = props.items.as_deref().ok_or(Unsupported::ArrayItems)?; self.write_type(items, indent)?; self.output.push_str("[]"); } @@ -298,7 +298,7 @@ mod tests { assert!(matches!( format_function_definitions(&tools), - Err(Error::ArrayItems) + Err(Error::Unsupported(Unsupported::ArrayItems)) )); } } diff --git a/litellm-rust/crates/token-counter/src/types.rs b/litellm-rust/crates/token-counter/src/types.rs index c1236f94d9a..7e1513d7706 100644 --- a/litellm-rust/crates/token-counter/src/types.rs +++ b/litellm-rust/crates/token-counter/src/types.rs @@ -5,7 +5,7 @@ use serde::de::{MapAccess, SeqAccess, Visitor}; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::Number; -use super::Error; +use super::{Error, Unsupported}; /// The parts of a request body the host's budget counter reads. Chat and /// Anthropic Messages bodies carry `messages`; completions carry `prompt`; @@ -33,7 +33,7 @@ pub struct CountableRequest { impl CountableRequest { pub fn parse(body: &[u8]) -> Result { - serde_json::from_slice(body).map_err(Error::RequestParse) + serde_json::from_slice(body).map_err(|error| Unsupported::RequestParse(error).into()) } } diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 542bd4a1fc4..cc18bad5402 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -2,7 +2,7 @@ use rstest::rstest; #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] use litellm_token_counter::TokenCounter; -use litellm_token_counter::{CountableRequest, Error}; +use litellm_token_counter::{CountableRequest, Error, Unsupported}; #[cfg(any(feature = "fast", feature = "huggingface"))] mod json { @@ -100,7 +100,12 @@ mod json { let request = CountableRequest::parse(body).expect("shape parses"); assert!(matches!( counter(load).count_request(&request), - Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) + Err(Error::Unsupported( + Unsupported::MissingInput + | Unsupported::FloatText + | Unsupported::ContentBlock + | Unsupported::ArrayItems + )) )); } @@ -139,7 +144,7 @@ mod json { } fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { - assert!(matches!(load("{}"), Err(Error::Load(_)))); + assert!(load("{}").is_err_and(|error| error.is_load())); } macro_rules! json_backend_tests { @@ -206,7 +211,7 @@ mod json { let counter = $loader(&tokenizer.to_string(false).unwrap()).unwrap(); let request = CountableRequest::parse(br#"{"prompt":"hello"}"#).unwrap(); let error = counter.count_request(&request).unwrap_err(); - assert!(matches!(error, Error::Encode(_))); + assert!(!error.is_load() && !error.is_unsupported()); assert_eq!(error.source().unwrap().to_string(), expected.to_string()); } @@ -250,7 +255,7 @@ mod json { fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { assert!(matches!( CountableRequest::parse(body), - Err(Error::RequestParse(_)) + Err(Error::Unsupported(Unsupported::RequestParse(_))) )); } @@ -443,7 +448,7 @@ mod tiktoken { #[case] rank_file: &str, #[values(CL100K, O200K)] encoding: TiktokenEncoding, ) { - assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); + assert!((encoding.load)(rank_file).is_err_and(|error| error.is_load())); } } @@ -452,6 +457,7 @@ mod tiktoken { fn unsupported_encoding_reaches_the_counter_caller() { assert!(matches!( TokenCounter::from_tiktoken("unknown-encoding"), - Err(Error::UnsupportedTokenizer(name)) if name == "unknown-encoding" + Err(error) if error.is_load() + && litellm_error::Report::new(&error).to_string().contains("unknown-encoding") )); } diff --git a/litellm/__init__.py b/litellm/__init__.py index 2cff21bf29a..8dfe5adca1c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -681,6 +681,7 @@ moonshot_models: Set = set() publicai_models: Set = set() darkbloom_models: Set = set() openinfer_models: Final[Set[str]] = set() +zerogpu_models: Final[Set[str]] = set() v0_models: Set = set() morph_models: Set = set() opencode_go_models: Final[set[str]] = set() # mutable-ok: populated when the shared model catalog loads @@ -969,6 +970,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje darkbloom_models.add(key) elif value.get("litellm_provider") == "openinfer": openinfer_models.add(key) + elif value.get("litellm_provider") == "zerogpu": + zerogpu_models.add(key) elif value.get("litellm_provider") == "v0": v0_models.add(key) elif value.get("litellm_provider") == "morph": @@ -1142,6 +1145,7 @@ model_list = list( | publicai_models | darkbloom_models | openinfer_models + | zerogpu_models | v0_models | morph_models | opencode_go_models @@ -1261,6 +1265,7 @@ def _build_models_by_provider() -> dict: "publicai": publicai_models, "darkbloom": darkbloom_models, "openinfer": openinfer_models, + "zerogpu": zerogpu_models, "v0": v0_models, "morph": morph_models, "opencode_go": opencode_go_models, @@ -1449,6 +1454,7 @@ from .exceptions import ( Timeout, APIConnectionError, UnsupportedParamsError, + UnsupportedEndpointError as UnsupportedEndpointError, APIResponseValidationError, UnprocessableEntityError, InternalServerError, @@ -1738,6 +1744,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/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 0d6dfb5c0d4..2b083b03766 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1096,6 +1096,7 @@ openai_compatible_providers: Final[list[str]] = [ "llamafile", "lm_studio", "galadriel", + "zerogpu", "github_copilot", # GitHub Copilot Chat API "chatgpt", # ChatGPT subscription API "novita", @@ -1669,6 +1670,7 @@ MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id" +COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: Final = "_complexity_router_child_session_id" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " @@ -1736,6 +1738,7 @@ PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-" TYPESAFE_DANGER_ZONE_HEADER_PREFIX: Final = "typesafe-danger-zone-" INTERNAL_KWARG_PREFIX: Final = "_litellm_" CONTROL_OPTIONS_KEY: Final = f"{INTERNAL_KWARG_PREFIX}control" +SESSION_ID_FROM_TRACE_ID_PARAM_KEY: Final = f"{INTERNAL_KWARG_PREFIX}session_id_from_trace_id" AZURE_SPEECH_CUSTOM_LLM_PROVIDER: Final = "azure_speech" AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech" diff --git a/litellm/exceptions.py b/litellm/exceptions.py index b8e933bf046..cd5139f5f79 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -232,7 +232,7 @@ class BadRequestError(openai.BadRequestError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, - body: dict | None = None, + body: Mapping[str, object] | None = None, ): self.status_code = 400 self.message = f"litellm.BadRequestError: {message}" @@ -276,7 +276,7 @@ class ImageFetchError(BadRequestError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, - body: dict | None = None, + body: Mapping[str, object] | None = None, ): super().__init__( message=message, @@ -317,6 +317,10 @@ class MalformedRequestFieldError(BadRequestError): pass +class UnsupportedEndpointError(BadRequestError): + pass + + class UnprocessableEntityError(openai.UnprocessableEntityError): def __init__( self, @@ -556,6 +560,7 @@ class ContextWindowExceededError(BadRequestError): llm_provider, response: httpx.Response | None = None, litellm_debug_info: str | None = None, + body: Mapping[str, object] | None = None, ): self.status_code = 400 self.model = model @@ -567,6 +572,7 @@ class ContextWindowExceededError(BadRequestError): llm_provider=self.llm_provider, response=response, litellm_debug_info=self.litellm_debug_info, + body=body, ) # Call the base class constructor with the parameters it needs # set after, to make it clear the raised error is a context window exceeded error @@ -685,7 +691,7 @@ class ContentPolicyViolationError(BadRequestError): response: httpx.Response | None = None, litellm_debug_info: str | None = None, provider_specific_fields: dict | None = None, - body: dict | None = None, + body: Mapping[str, object] | None = None, ): self.status_code = 400 self.message = f"litellm.ContentPolicyViolationError: {message}" @@ -1054,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 9c88e737a40..84156919774 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -284,6 +284,12 @@ def image_generation( extra_headers: Final = kwargs.get("extra_headers", None) headers: Final[dict] = kwargs.get("headers", None) or {} base_model: Final = kwargs.get("base_model", None) + call_kwargs: Final = cast(Mapping[str, object], kwargs) # cast-ok: image_generation's **kwargs is untyped + max_retries: Final = call_kwargs.get("max_retries") + is_router_call: Final = any( + isinstance(call_metadata := call_kwargs.get(k), Mapping) and "model_group" in call_metadata + for k in ("metadata", "litellm_metadata") + ) if extra_headers is not None: headers.update(extra_headers) model_response: ImageResponse = litellm.utils.ImageResponse() @@ -550,6 +556,7 @@ def image_generation( optional_params["extra_headers"] = extra_headers # Forward OpenAI organization if present (set by proxy pre-call utils) organization: Final[str | None] = kwargs.get("organization", None) + sdk_max_retries: Final = 0 if is_router_call else max_retries model_response = openai_chat_completions.image_generation( model=model, prompt=prompt, @@ -557,7 +564,9 @@ def image_generation( api_key=api_key or dynamic_api_key, api_base=api_base, logging_obj=litellm_logging_obj, - optional_params=optional_params, + optional_params=( + optional_params if sdk_max_retries is None else {**optional_params, "max_retries": sdk_max_retries} + ), model_response=model_response, organization=organization, aimg_generation=aimg_generation, @@ -1053,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, @@ -1093,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/opentelemetry_utils/gen_ai_semconv.py b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py index 035e5c5db10..84c4b2e58ab 100644 --- a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py +++ b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py @@ -35,6 +35,7 @@ from enum import Enum from typing import TYPE_CHECKING, Final, Protocol from litellm.integrations.otel.mappers.utils import drop_none +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.safe_json_dumps import safe_dumps if TYPE_CHECKING: @@ -148,8 +149,11 @@ class OTELGenAISemconvMixin: """Map a LiteLLM ``call_type`` to spec ``gen_ai.operation.name``. Substring match (e.g. ``aembedding`` -> ``embeddings``); defaults to - ``chat``. + ``chat``. A pass-through call's operation, stored under ``PASSTHROUGH_OPERATION_KEY``, wins. """ + passthrough_operation: Final = kwargs.get(PASSTHROUGH_OPERATION_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # kwargs is an untyped dict + if isinstance(passthrough_operation, GenAIOperation): + return passthrough_operation.value call_type: Final = kwargs.get("call_type", "") or "" match call_type: case s if "embedding" in s: diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index e09226ac9cd..4c5e65ff07d 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -600,6 +600,7 @@ class OpenTelemetryV2(CustomLogger): trace=call.trace, session_id=call.session_id, metadata_keys=tuple(self.config.baggage_metadata_keys), + operation=call.operation, ) end_time_ns: Final = to_ns(end_time) if carrier is not None and carrier.span is not None: diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index b96ed55872d..a9a57de04ac 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -49,7 +49,7 @@ from litellm.constants import ( LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, SESSION_ID_GENERATED_METADATA_KEY, ) -from litellm.integrations.otel.model.semconv import resolve_operation +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation, resolve_operation from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls from litellm.integrations.otel.model.utils import as_str, as_str_mapping, to_seconds from litellm.types.utils import InternalCallOrigin @@ -262,6 +262,7 @@ class LLMCallEvent: # never hands off to a provider, so ``upstream_started`` stays False even # though the call completed and is worth a span. cache_hit: bool + operation: GenAIOperation # A best-effort ``"{operation} {model}"`` name known at ``pre_call`` time. The # span is renamed from the typed payload at close (``finish_span``); this only # needs to be reasonable for a span that never gets closed (a leak). @@ -275,7 +276,7 @@ class LLMCallEvent: def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent: raw_payload: Final = kwargs.get("standard_logging_object") payload: Final = cast("StandardLoggingPayload", raw_payload) if raw_payload else None - operation: Final = resolve_operation(as_str(kwargs.get("call_type"))) + operation: Final = resolve_operation(as_str(kwargs.get("call_type")), kwargs.get(PASSTHROUGH_OPERATION_KEY)) model: Final = as_str(kwargs.get("model")) or "" trace: Final = caller_trace_controls(kwargs) return cls( @@ -288,6 +289,7 @@ class LLMCallEvent: upstream_start_seconds=_epoch_seconds(kwargs.get("api_call_start_time")), purpose=internal_call_origin(payload, kwargs), cache_hit=bool(payload.get("cache_hit")) if payload else False, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown + operation=operation, provisional_span_name=f"{operation.value} {model}".strip(), time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs), completion_start_seconds=_epoch_seconds(kwargs.get("completion_start_time")), diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 21b2f44f077..7a62ed6aed9 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -464,6 +464,7 @@ class LLMCallSpanData: session_id: str | None = None, *, metadata_keys: tuple[str, ...] = (), + operation: GenAIOperation | None = None, ) -> LLMCallSpanData: params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {}) cache_hit: Final = as_bool(payload.get("cache_hit")) is True # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown @@ -486,13 +487,13 @@ class LLMCallSpanData: # no prompt/response text. finish_reasons: Final = _finish_reasons(choices_out) call_type: Final = as_str(payload.get("call_type")) - operation: Final = resolve_operation(call_type) + resolved_operation: Final = operation or resolve_operation(call_type) embedding_output: Final = ( - EmbeddingOutput.from_response(response) if operation is GenAIOperation.EMBEDDINGS else None + EmbeddingOutput.from_response(response) if resolved_operation is GenAIOperation.EMBEDDINGS else None ) cost_lookup_failed: Final = payload.get("response_cost_failure_debug_info") is not None and not cache_hit return cls( - operation=operation, + operation=resolved_operation, provider=resolve_provider(as_str(payload.get("custom_llm_provider"))), request_model=context.request_model, response_model=context.response_model, diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index b7f955a75ed..8d186a2715a 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -34,6 +34,7 @@ class GenAIOperation(str, Enum): LITELLM_VECTOR_STORE_FILE_MANAGEMENT = "litellm.vector_store_file_management" LITELLM_RESPONSES_MANAGEMENT = "litellm.responses_management" LITELLM_MODERATION = "litellm.moderation" + LITELLM_PASSTHROUGH = "litellm.passthrough" class GenAIOutputType(str, Enum): @@ -482,14 +483,21 @@ def resolve_provider(custom_llm_provider: str | None) -> str: return mapped.value if mapped is not None else custom_llm_provider -def resolve_operation(call_type: str | None) -> GenAIOperation: +PASSTHROUGH_OPERATION_KEY: Final = "passthrough_gen_ai_operation" + + +def resolve_operation(call_type: str | None, passthrough_operation: object = None) -> GenAIOperation: """Map a litellm ``call_type`` to a ``gen_ai.operation.name`` value. An unmapped call type still falls back to ``chat`` so every series keeps an operation label, but it logs at debug rather than falling through silently: a new call type mislabelled as ``chat`` mixes its latency and cost into everyone's chat charts, which is invisible until someone reads the numbers. + + A pass-through call's operation, stored under ``PASSTHROUGH_OPERATION_KEY``, wins over its call type. """ + if isinstance(passthrough_operation, GenAIOperation): + return passthrough_operation if not call_type: return GenAIOperation.CHAT mapped: Final = _OPERATION_BY_CALL_TYPE.get(call_type.lower()) diff --git a/litellm/integrations/otel/mount.py b/litellm/integrations/otel/mount.py index c0d42eeb5c5..9d6c5fca75c 100644 --- a/litellm/integrations/otel/mount.py +++ b/litellm/integrations/otel/mount.py @@ -22,6 +22,7 @@ from litellm.integrations.otel.model.config import is_otel_v2_enabled if TYPE_CHECKING: from fastapi import FastAPI from opentelemetry.trace import Span, TracerProvider + from starlette.types import ASGIApp, Receive, Scope, Send # Routes excluded from server-span tracing by default: high-frequency pollers and # static UI/docs assets, none of which are LLM traffic. Entries are substring-matched @@ -44,6 +45,33 @@ _DEFAULT_EXCLUDED_ROUTES: Final = ( ) _DEFAULT_EXCLUDED_URLS: Final = ",".join(_DEFAULT_EXCLUDED_ROUTES) + +class _RequestTraceExclusionMiddleware: + def __init__(self, app: "ASGIApp", excluded_urls: str) -> None: + from opentelemetry.util.http import parse_excluded_urls + + self.app = app + self.excluded_urls = parse_excluded_urls(excluded_urls) + + async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None: + if scope["type"] not in ("http", "websocket"): + await self.app(scope, receive, send) + return + from opentelemetry.instrumentation.asgi import ( # pyright: ignore[reportMissingTypeStubs] # upstream ships no py.typed marker + get_host_port_url_tuple, # pyright: ignore[reportUnknownVariableType] # upstream URL helper has no annotations + ) + + from litellm.integrations.otel.plumbing.context import ( + request_trace_scope, + request_traces_excluded, + ) + + url: Final = str(get_host_port_url_tuple(scope)[2]) # pyright: ignore[reportUnknownArgumentType] # upstream constructs a str URL but leaves its return type unannotated + excluded: Final = request_traces_excluded() or self.excluded_urls.url_disabled(url) + with request_trace_scope(excluded): + await self.app(scope, receive, send) + + # Passthrough routes are catch-alls (e.g. "/openai/{endpoint:path}"), so the # default OTel server-span name "{method} {route}" collapses every upstream # endpoint into "POST /openai/{endpoint:path}". The hook below renames those spans @@ -129,6 +157,8 @@ def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | No if app.middleware_stack is not None: verbose_logger.debug("Skipping OTel V2 FastAPI instrumentation: the app has already started") return + if getattr(app, "_is_instrumented_by_opentelemetry", False): + return # Lazy: only the V2-enabled path needs the optional # ``opentelemetry-instrumentation-fastapi`` package, which is not part of the @@ -151,11 +181,7 @@ def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | No claim_server_span(span) _passthrough_span_name_hook(span, scope) - excluded_urls: Final = ( - os.environ.get("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS") - if "OTEL_PYTHON_FASTAPI_EXCLUDED_URLS" in os.environ - else _DEFAULT_EXCLUDED_URLS - ) + excluded_urls: Final = os.environ.get("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", _DEFAULT_EXCLUDED_URLS) FastAPIInstrumentor.instrument_app( app, tracer_provider=tracer_provider, @@ -166,6 +192,12 @@ def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | No # template in their name, which can't be rewritten from a hook. exclude_spans=["receive", "send"], ) + instrumented_stack: Final = app.build_middleware_stack + + def build_middleware_stack() -> "ASGIApp": + return _RequestTraceExclusionMiddleware(instrumented_stack(), excluded_urls) + + app.build_middleware_stack = build_middleware_stack # rebind-ok: wrap the complete ASGI stack if get_global_response_propagator() is None: set_global_response_propagator(UpstreamFirstTraceResponsePropagator()) except Exception as e: diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index d0c419abe87..b328b281dde 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -30,6 +30,21 @@ if TYPE_CHECKING: _PROPAGATOR: Final = TraceContextTextMapPropagator() _W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate")) +_request_traces_excluded: Final[ContextVar[bool]] = ContextVar("litellm_otel_request_traces_excluded", default=False) + + +def request_traces_excluded() -> bool: + return _request_traces_excluded.get() + + +@contextmanager +def request_trace_scope(excluded: bool) -> Generator[None]: + token: Final = _request_traces_excluded.set(excluded) + try: + yield + finally: + _request_traces_excluded.reset(token) + # The request's root span — the FastAPI-owned SERVER span — captured ONCE when the # proxy first resolves it, so request-level spans (the LLM call, guardrails) can diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 72d51471ba2..1197ae055d0 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -26,6 +26,7 @@ from litellm.integrations.opentelemetry import ( ) from litellm.integrations.otel.model.metadata import time_to_first_chunk_seconds from litellm.integrations.otel.model.semconv import ( + PASSTHROUGH_OPERATION_KEY, Error, GenAI, Metric, @@ -280,7 +281,9 @@ class GenAIMetricRecorder: def _common_attributes(self, kwargs: Mapping[str, Any]) -> dict: params: Final = kwargs.get("litellm_params") or {} common_attrs: Final[dict] = { - GenAI.OPERATION_NAME: resolve_operation(kwargs.get("call_type")).value, + GenAI.OPERATION_NAME: resolve_operation( + kwargs.get("call_type"), kwargs.get(PASSTHROUGH_OPERATION_KEY) + ).value, **_provider_attributes(params.get("custom_llm_provider")), GenAI.REQUEST_MODEL: kwargs.get("model"), "gen_ai.framework": "litellm", diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 9736192a816..a369e352609 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -35,7 +35,8 @@ from opentelemetry.sdk.trace.export import ( from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( InMemorySpanExporter, ) -from opentelemetry.trace import Span, SpanContext, SpanKind, Status, StatusCode, Tracer +from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult +from opentelemetry.trace import Link, Span, SpanContext, SpanKind, Status, StatusCode, Tracer, TraceState from opentelemetry.util.re import parse_env_headers from opentelemetry.util.types import Attributes, AttributeValue @@ -68,6 +69,7 @@ from litellm.integrations.otel.model.semconv import ( from litellm.integrations.otel.model.spans import LiteLLMSpanKind from litellm.integrations.otel.plumbing.context import ( request_destinations, + request_traces_excluded, suppressed_backends, ) from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls @@ -1499,6 +1501,28 @@ def _spec_processor( return _MessageContentFilter(scoped) if strip_message_content else scoped +class _RequestExclusionSampler(Sampler): + def __init__(self, delegate: Sampler) -> None: + self.delegate = delegate + + def should_sample( + self, + parent_context: Context | None, + trace_id: int, + name: str, + kind: SpanKind | None = None, + attributes: Attributes = None, + links: Sequence[Link] | None = None, + trace_state: TraceState | None = None, + ) -> SamplingResult: + if request_traces_excluded(): + return SamplingResult(Decision.DROP) + return self.delegate.should_sample(parent_context, trace_id, name, kind, attributes, links, trace_state) + + def get_description(self) -> str: + return f"RequestExclusion({self.delegate.get_description()})" + + def build_tracer_provider( config: OpenTelemetryV2Config, exporter: SpanExporter | None = None, @@ -1531,6 +1555,7 @@ def build_tracer_provider( destination asked for. """ provider: Final = TracerProvider(resource=build_resource(config)) + provider.sampler = _RequestExclusionSampler(provider.sampler) if baggage_processor is None: baggage_processor = LiteLLMBaggageSpanProcessor(allowed_keys=config.baggage_promoted_keys) provider.add_span_processor(baggage_processor) 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/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 1c816b22ef1..47ec8331c83 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -2036,14 +2036,6 @@ def _map_azure_exception( litellm_debug_info=extra_information, response=getattr(original_exception, "response", None), ) - elif "This model's maximum context length is" in error_str: - raise ContextWindowExceededError( - message=f"AzureException ContextWindowExceededError - {message}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) elif "DeploymentNotFound" in error_str: raise NotFoundError( message=f"AzureException NotFoundError - {message}", @@ -2085,6 +2077,18 @@ def _map_azure_exception( response=getattr(original_exception, "response", None), body=getattr(original_exception, "body", None), ) + elif getattr(original_exception, "status_code", None) in (None, 400) and ( + azure_error_code == "context_length_exceeded" + or ExceptionCheckers.is_error_str_context_window_exceeded(error_str) + ): + raise ContextWindowExceededError( + message=f"AzureException ContextWindowExceededError - {message}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + body=getattr(original_exception, "body", None), + ) elif "invalid_request_error" in error_str and getattr(original_exception, "status_code", None) in (None, 400): raise BadRequestError( message=f"AzureException BadRequestError - {message}", diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index be432d1fdd4..eb4cf421a19 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -19,17 +19,11 @@ ones on key conflicts, and the caller backfills ``litellm_provider`` with the provider it requested. If no capability rule matches, model-info resolution misses as if no rules existed. -LEGACY-SCHEMA SHIM (temporary, until the new-schema JSON reaches main): released -proxies fetch this JSON remotely from main, whose block still ships the old schema -where a rule mixes ``litellm_provider`` with capability keys and may inherit a -parent's ``model_info`` via ``extends``. Such a legacy rule is tolerated rather -than skipped: ``extends`` is resolved once at install time (single level, against -raw parents), and the resolved rule acts as BOTH kinds, a routing rule (its -``litellm_provider`` participates in first-hit inference) and a capability rule -(its full ``model_info``, provider included, participates in the union). New-schema -rules never mix the two and never use ``extends``. A rule whose -``litellm_provider`` is not a string is invalid and is warned about and skipped -(a warning rather than a crash, for the same remote-fetch reason). +A rule whose ``model_info`` carries ``litellm_provider`` alongside any other key, +or carries it as a non-string, is neither kind: it is warned about and skipped +rather than raising, since released proxies fetch this JSON remotely. A rule +carrying ``extends`` is skipped the same way. Mixed rules and ``extends`` are the +old rule format, so their warning points the operator at a current cost map. Rules are only consulted after exact and case-insensitive lookups miss, so an exact cost-map entry always takes precedence over any rule. @@ -62,37 +56,8 @@ NAME_FIELD: Final = "name" PATTERN_FIELD: Final = "pattern" MODEL_INFO_FIELD: Final = "model_info" PROVIDER_KEY: Final = "litellm_provider" -LEGACY_EXTENDS_FIELD: Final = "extends" FILL_MISSING_FOR_PROVIDERS_FIELD: Final = "fill_missing_for_providers" - - -def _resolve_legacy_extends(rules: list) -> list: - """Expand legacy ``extends`` inheritance so each rule's ``model_info`` is self-contained. - - Compatibility shim for the old remote schema: single level, resolved against each - parent's raw ``model_info``, with the child's own keys winning on conflict. Non-dict - rules and dangling parents pass through unchanged; new-schema rules carry no - ``extends`` and are untouched. - """ - base_by_name: Final = { - rule[NAME_FIELD]: rule[MODEL_INFO_FIELD] - for rule in rules - if isinstance(rule, dict) - and isinstance(rule.get(NAME_FIELD), str) - and isinstance(rule.get(MODEL_INFO_FIELD), dict) - } - - def resolved(rule: object) -> object: - if not isinstance(rule, dict): - return rule - parent_name: Final = rule.get(LEGACY_EXTENDS_FIELD) - own_info: Final = rule.get(MODEL_INFO_FIELD) - parent_info: Final = base_by_name.get(parent_name) if isinstance(parent_name, str) else None - if parent_info is None or not isinstance(own_info, dict): - return rule - return {**rule, MODEL_INFO_FIELD: {**parent_info, **own_info}} - - return [resolved(rule) for rule in rules] +EXTENDS_FIELD: Final = "extends" @dataclass(frozen=True, slots=True) @@ -127,19 +92,33 @@ def _parse_fill_missing_for_providers(rule: Mapping[str, object], pattern_label: return frozenset(raw_fill_missing_for_providers) -def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: +def _warn_old_format_rule(rule_label: str, reason: str) -> None: + verbose_logger.warning( + "LiteLLM: skipping fallback generalization rule %s, written in the old cost map rule format (%s). " + "Point LITELLM_MODEL_COST_MAP_URL at a current cost map, or with LITELLM_LOCAL_MODEL_COST_MAP set, " + "replace the bundled backup with a current one.", + rule_label, + reason, + ) + + +def _compile_rule(rule: object) -> _CompiledRule | None: if not isinstance(rule, dict): - return () + return None pattern: Final = rule.get(PATTERN_FIELD) model_info: Final = rule.get(MODEL_INFO_FIELD) + rule_label: Final = f"{rule.get(NAME_FIELD, pattern)}" if not isinstance(pattern, str) or not isinstance(model_info, dict): verbose_logger.warning( "LiteLLM: skipping malformed fallback generalization rule %s (needs string '%s' and dict '%s').", - rule.get(NAME_FIELD, pattern), + rule_label, PATTERN_FIELD, MODEL_INFO_FIELD, ) - return () + return None + if EXTENDS_FIELD in rule: + _warn_old_format_rule(rule_label, f"'{EXTENDS_FIELD}' is no longer supported") + return None try: compiled: Final = re.compile(pattern, re.IGNORECASE) except re.error as e: @@ -148,37 +127,28 @@ def _compile_rule(rule: object) -> tuple[_CompiledRule, ...]: pattern, e, ) - return () + return None fill_missing_for_providers: Final = _parse_fill_missing_for_providers(rule, pattern) if fill_missing_for_providers is None: - return () + return None if PROVIDER_KEY not in model_info: - return ( - _CapabilityRule( - pattern=compiled, - model_info=model_info, - fill_missing_for_providers=fill_missing_for_providers, - ), + return _CapabilityRule( + pattern=compiled, + model_info=model_info, + fill_missing_for_providers=fill_missing_for_providers, ) provider: Final = model_info[PROVIDER_KEY] if not isinstance(provider, str): verbose_logger.warning( - "LiteLLM: skipping invalid fallback generalization rule %s: '%s' in '%s' must be a string.", - rule.get(NAME_FIELD, pattern), + "LiteLLM: skipping malformed fallback generalization rule %s ('%s' must be a string).", + rule_label, PROVIDER_KEY, - MODEL_INFO_FIELD, ) - return () - if len(model_info) == 1: - return (_RoutingRule(pattern=compiled, provider=provider),) - return ( - _RoutingRule(pattern=compiled, provider=provider), - _CapabilityRule( - pattern=compiled, - model_info=model_info, - fill_missing_for_providers=fill_missing_for_providers, - ), - ) + return None + if len(model_info) != 1: + _warn_old_format_rule(rule_label, f"'{PROVIDER_KEY}' must be the only key in '{MODEL_INFO_FIELD}'") + return None + return _RoutingRule(pattern=compiled, provider=provider) class _FallbackGeneralizations: @@ -192,7 +162,7 @@ class _FallbackGeneralizations: def set_rules(self, rules: list | None) -> None: installed: Final = rules if isinstance(rules, list) else [] - compiled: Final = tuple(kind for rule in _resolve_legacy_extends(installed) for kind in _compile_rule(rule)) + compiled: Final = tuple(kind for kind in map(_compile_rule, installed) if kind is not None) self.rules = installed self.routing_rules = tuple(rule for rule in compiled if isinstance(rule, _RoutingRule)) self.capability_rules = tuple(rule for rule in compiled if isinstance(rule, _CapabilityRule)) @@ -224,10 +194,7 @@ class _FallbackGeneralizations: ) if not matched: return None - fill_missing: Final[Mapping[str, object]] = { - key: value for model_info in matched for key, value in model_info.items() if key != PROVIDER_KEY - } - return fill_missing or None + return {key: value for model_info in matched for key, value in model_info.items()} or None _registry: Final = _FallbackGeneralizations() @@ -236,10 +203,8 @@ _registry: Final = _FallbackGeneralizations() def set_fallback_generalizations(rules: list | None) -> None: """Install the active rule list, compiling and classifying each rule. - Legacy ``extends`` inheritance is resolved here, once, before classification; - a legacy rule mixing ``litellm_provider`` with capability keys installs as both - kinds. Malformed and invalid-regex rules are warned about and skipped. Called - once when the model cost map is loaded (and again on any reload). + Malformed, mixed-kind, and invalid-regex rules are warned about and skipped. + Called once when the model cost map is loaded (and again on any reload). """ _registry.set_rules(rules) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index ca8e2af18e6..ba964d200ec 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -6,7 +6,7 @@ from typing import Final, cast from pydantic import TypeAdapter, ValidationError -from litellm.constants import CONTROL_OPTIONS_KEY +from litellm.constants import CONTROL_OPTIONS_KEY, SESSION_ID_FROM_TRACE_ID_PARAM_KEY from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.types.litellm_params import MAX_CONTROL_INT_DIGITS, ControlOptions @@ -212,6 +212,9 @@ def get_litellm_params( # Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining) _meta: Final = resolved_metadata or {} + session_id_from_trace_id: Final = ( + litellm_session_id is None and not _meta.get("session_id") and bool(_meta.get("trace_id")) + ) if litellm_session_id is None: litellm_session_id = _meta.get("session_id") or _meta.get("trace_id") if litellm_trace_id is None: @@ -254,6 +257,7 @@ def get_litellm_params( or (get_base_model_from_litellm_call_metadata(metadata=metadata) if metadata else None), "litellm_trace_id": litellm_trace_id, "litellm_session_id": litellm_session_id, + **({SESSION_ID_FROM_TRACE_ID_PARAM_KEY: True} if session_id_from_trace_id else {}), "hf_model_name": hf_model_name, "custom_prompt_dict": custom_prompt_dict, "litellm_metadata": litellm_metadata, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 48e8a44b9a4..3b8157b2249 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) 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/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index e022839c5c7..c831e92c09d 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -2,7 +2,7 @@ import re from collections.abc import Mapping from typing import Final -from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import SESSION_ID_FROM_TRACE_ID_PARAM_KEY, SESSION_ID_GENERATED_METADATA_KEY _HTTP_HEADER_NAME_PATTERN: Final = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") _FORBIDDEN_AFFINITY_HEADERS: Final = frozenset( @@ -53,36 +53,51 @@ def _get_provider_affinity_header_name(litellm_params: object | None, default_he return validate_provider_affinity_header_name(header) -def get_stable_session_id(litellm_params: object | None) -> str | None: - if litellm_params is None: - return None - - direct_session_id: Final = _get_value(litellm_params, "session_id") - if direct_session_id: - return str(direct_session_id) - - metadata_values: Final[tuple[object, ...]] = tuple( +def _metadata_values(litellm_params: object) -> tuple[object, ...]: + return tuple( value for key in ("metadata", "litellm_metadata") if (value := _get_value(litellm_params, key)) is not None ) + + +def _non_generated_litellm_session_id(litellm_params: object) -> str | None: has_generated_session_id: Final = any( - isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY) - for metadata in metadata_values + _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) for metadata in _metadata_values(litellm_params) ) - litellm_session_id: Final = _get_value(litellm_params, "litellm_session_id") - if litellm_session_id and not has_generated_session_id: - return str(litellm_session_id) + return str(litellm_session_id) if litellm_session_id and not has_generated_session_id else None - for metadata in metadata_values: - if ( - isinstance(metadata, Mapping) - and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY) - and (value := metadata.get("session_id")) + +def _metadata_session_id(litellm_params: object) -> str | None: + for metadata in _metadata_values(litellm_params): + if not _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) and ( + value := _get_value(metadata, "session_id") ): return str(value) return None +def get_caller_litellm_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None or _get_value(litellm_params, SESSION_ID_FROM_TRACE_ID_PARAM_KEY): + return None + return _non_generated_litellm_session_id(litellm_params) + + +def get_requested_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None: + return None + direct_session_id: Final = _get_value(litellm_params, "session_id") + return str(direct_session_id) if direct_session_id else _metadata_session_id(litellm_params) + + +def get_stable_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None: + return None + direct_session_id: Final = _get_value(litellm_params, "session_id") + if direct_session_id: + return str(direct_session_id) + return _non_generated_litellm_session_id(litellm_params) or _metadata_session_id(litellm_params) + + def add_provider_affinity_header( headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers 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/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 611051bb2dd..032951be74a 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -3,6 +3,7 @@ from abc import abstractmethod from typing import TYPE_CHECKING, Final, cast import litellm +from litellm.constants import STREAM_SSE_DONE_STRING if TYPE_CHECKING: import httpx @@ -142,7 +143,8 @@ class BaseModelResponseIterator: def _handle_string_chunk(self, str_line: str) -> GenericStreamingChunk | ModelResponseStream: # chunk is a str at this point stripped_json_chunk: Final = BaseModelResponseIterator._string_to_dict_parser(str_line=str_line) - if "[DONE]" in str_line: + stripped_sse_chunk: Final = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(str_line) + if stripped_sse_chunk is not None and stripped_sse_chunk.strip() == STREAM_SSE_DONE_STRING: return GenericStreamingChunk( text="", is_finished=True, diff --git a/litellm/llms/bedrock/chat/native_output_fallback.py b/litellm/llms/bedrock/chat/native_output_fallback.py index 1b8ce5fcda4..6892cd37beb 100644 --- a/litellm/llms/bedrock/chat/native_output_fallback.py +++ b/litellm/llms/bedrock/chat/native_output_fallback.py @@ -1,11 +1,14 @@ from collections.abc import Callable, Coroutine +from functools import partial from typing import Final, TypeAlias from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.common_utils import ( NativeOutputFallback, is_bedrock_native_output_rejection, + is_bedrock_strict_tools_rejection, warn_bedrock_native_output_rejected, + warn_bedrock_strict_tools_dropped, ) from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper @@ -14,12 +17,13 @@ NativeOutputResponse: TypeAlias = ModelResponse | CustomStreamWrapper NativeOutputAttemptResult: TypeAlias = NativeOutputResponse | Coroutine[object, object, NativeOutputResponse] NativeOutputAttempt: TypeAlias = Callable[[NativeOutputFallback | None], NativeOutputAttemptResult] NativeOutputFallbackFactory: TypeAlias = Callable[[], NativeOutputFallback | None] +_RetryFor: TypeAlias = Callable[[BaseLLMException], NativeOutputFallback | None] -def _fallback_for( - error: BaseLLMException, fallback: NativeOutputFallbackFactory, model: str +def _native_output_retry( + error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str ) -> NativeOutputFallback | None: - if not is_bedrock_native_output_rejection(error.status_code, error.message): + if fallback is None or not is_bedrock_native_output_rejection(error.status_code, error.message): return None retry: Final = fallback() if retry is not None: @@ -27,16 +31,38 @@ def _fallback_for( return retry +def _strict_tools_retry( + error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str +) -> NativeOutputFallback | None: + if fallback is None or not is_bedrock_strict_tools_rejection(error.status_code, error.message): + return None + retry: Final = fallback() + if retry is not None: + warn_bedrock_strict_tools_dropped(model) + return retry + + +def _fallback_for( + error: BaseLLMException, + fallback: NativeOutputFallbackFactory | None, + strict_tools_fallback: NativeOutputFallbackFactory | None, + model: str, +) -> NativeOutputFallback | None: + native_output_retry: Final = _native_output_retry(error, fallback, model) + if native_output_retry is not None: + return native_output_retry + return _strict_tools_retry(error, strict_tools_fallback, model) + + async def _await_with_native_output_fallback( first_attempt: Coroutine[object, object, NativeOutputResponse], attempt: NativeOutputAttempt, - fallback: NativeOutputFallbackFactory, - model: str, + retry_for: _RetryFor, ) -> NativeOutputResponse: try: return await first_attempt except BaseLLMException as error: - retry: Final = _fallback_for(error, fallback, model) + retry: Final = retry_for(error) if retry is None: raise second_attempt: Final = attempt(retry) @@ -47,16 +73,20 @@ def call_with_native_output_fallback( attempt: NativeOutputAttempt, fallback: NativeOutputFallbackFactory | None, model: str, + strict_tools_fallback: NativeOutputFallbackFactory | None = None, ) -> NativeOutputAttemptResult: - if fallback is None: + if fallback is None and strict_tools_fallback is None: return attempt(None) + retry_for: Final = partial( + _fallback_for, fallback=fallback, strict_tools_fallback=strict_tools_fallback, model=model + ) try: first_attempt: Final = attempt(None) except BaseLLMException as error: - retry: Final = _fallback_for(error, fallback, model) + retry: Final = retry_for(error) if retry is None: raise return attempt(retry) if isinstance(first_attempt, Coroutine): - return _await_with_native_output_fallback(first_attempt, attempt, fallback, model) + return _await_with_native_output_fallback(first_attempt, attempt, retry_for) return first_attempt diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 28888f1a23a..fde995c7907 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -362,6 +362,60 @@ class NativeOutputFallback: litellm_params: Mapping[str, object] +_BEDROCK_STRICT_TOOLS_REJECTION: Final = re.compile( + r"Too many strict tools \(\d+\)" + r"|Schemas contains too many (?:optional parameters|parameters with union types)" + r"|[Cc]ompiled grammar (?:is too large|size)" + r"|tools\.\d+\.custom(?:\.strict)?: " + r"|Grammar compilation timed out" +) +_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) +_TOOL_SPEC_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def is_bedrock_strict_tools_rejection(status_code: int, error_text: str) -> bool: + return status_code == 400 and _BEDROCK_STRICT_TOOLS_REJECTION.search(error_text) is not None + + +def warn_bedrock_strict_tools_dropped(model: str) -> None: + verbose_logger.warning( + "Bedrock rejected the strict tool schemas for model=%s, so litellm sent the request again without `strict` " + "because drop_params is set.", + model, + ) + + +def _without_strict(tool: Mapping[str, object]) -> Mapping[str, object]: + spec_key: Final = next((key for key in ("function", "toolSpec") if key in tool), None) + if spec_key is None: + return tool + spec: Final = _TOOL_SPEC_ADAPTER.validate_python(tool[spec_key]) + return {**tool, spec_key: {key: value for key, value in spec.items() if key != "strict"}} + + +def _sends_strict(tool: Mapping[str, object], supports_strict_tools: bool) -> bool: + if "toolSpec" in tool: + return "strict" in _TOOL_SPEC_ADAPTER.validate_python(tool["toolSpec"]) + if "function" not in tool: + return False + return supports_strict_tools and bool(_TOOL_SPEC_ADAPTER.validate_python(tool["function"]).get("strict")) + + +def bedrock_strict_tools_fallback( + model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object] +) -> NativeOutputFallback | None: + tools: Final = _TOOLS_ADAPTER.validate_python(optional_params.get("tools", ())) + supports_strict_tools: Final = bedrock_converse_supports_strict_tools( + bedrock_capability_model(model, litellm_params.get("base_model")) + ) + if not any(_sends_strict(tool, supports_strict_tools) for tool in tools): + return None + return NativeOutputFallback( + optional_params={**optional_params, "tools": [_without_strict(tool) for tool in tools]}, + litellm_params=litellm_params, + ) + + def requests_output_config_format(optional_params: Mapping[str, object]) -> bool: match optional_params: case {"output_config": {"format": object()}}: 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/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index a9f95c1f9b9..bf9c119f6f2 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -206,6 +206,10 @@ class BedrockCountTokensConfig(BaseAWSLLM): encoded_body: Final = base64.b64encode(json.dumps(body_data).encode()).decode() return {"input": {"invokeModel": {"body": encoded_body}}} + @staticmethod + def count_tokens_model_id(model: str) -> str: + return get_bedrock_base_model(model).removeprefix("bedrock/") + def get_bedrock_count_tokens_endpoint( self, model: str, @@ -225,12 +229,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): Returns: Complete endpoint URL for CountTokens API """ - # Use existing LiteLLM function to get the base model ID (removes region prefix) - model_id = get_bedrock_base_model(model) - - # Remove bedrock/ prefix if present - model_id = model_id.removeprefix("bedrock/") # Remove "bedrock/" prefix - encoded_model_id: Final = self.encode_model_id(model_id=model_id) + encoded_model_id: Final = self.encode_model_id(model_id=self.count_tokens_model_id(model)) base_url, _ = self.get_runtime_endpoint( api_base=api_base, diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index ec7fbfa92da..cd2ea6b9ccc 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -3,6 +3,7 @@ Constants and helpers for ChatGPT subscription OAuth. """ import base64 +import hashlib import json import os import platform @@ -17,6 +18,7 @@ import httpx from pydantic import JsonValue, TypeAdapter from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils +from litellm.litellm_core_utils.provider_affinity import get_caller_litellm_session_id, get_requested_session_id from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ProviderSpecificHeader @@ -395,22 +397,31 @@ def _normalize_litellm_params(litellm_params: object) -> Mapping[str, object]: return MappingProxyType({}) +def _prompt_cache_session_id(prompt_cache_key: object) -> str | None: + if not isinstance(prompt_cache_key, str) or not prompt_cache_key: + return None + return ( + prompt_cache_key + if _safe_header_value(prompt_cache_key.strip()) == prompt_cache_key + else hashlib.sha256(prompt_cache_key.encode("utf-8", "surrogatepass")).hexdigest() + ) + + def get_chatgpt_session_id(litellm_params: object) -> str | None: params: Final = _normalize_litellm_params(litellm_params) - for key in ("litellm_session_id", "session_id"): - value = params.get(key) - if value: - return str(value) - metadata: Final = _normalize_litellm_params(params.get("metadata")) - if metadata: - value = metadata.get("session_id") - if value: - return str(value) - for key in ("litellm_trace_id", "litellm_call_id"): - value = params.get(key) - if value: - return str(value) - return None + return ( + get_caller_litellm_session_id(params) + or get_requested_session_id(params) + or _prompt_cache_session_id(params.get("prompt_cache_key")) + or next( + ( + str(value) + for key in ("litellm_session_id", "litellm_trace_id", "litellm_call_id") + if (value := params.get(key)) + ), + None, + ) + ) def ensure_chatgpt_session_id(litellm_params: object) -> str: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ef009e70137..afa12b58b29 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6,6 +6,7 @@ import ssl from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache, partial, reduce +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -62,7 +63,7 @@ from litellm.litellm_core_utils.realtime_errors import ( ) from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base -from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -572,6 +573,22 @@ def _unwrap_provider_error(error: Exception) -> Exception: return error +def _merge_messages_header_sources(*sources: Mapping[str, object] | None) -> Mapping[str, object]: + """Later sources win per header, except ``anthropic-beta``, which unions every source so client betas + never drop a deployment's configured ones.""" + entries: Final = tuple(chain.from_iterable((source or {}).items() for source in sources)) + beta_values: Final = tuple( + cast("str | Sequence[str]", value) # cast-ok: the merge helper skips non-str list entries + for name, value in entries + if name.lower() == "anthropic-beta" and isinstance(value, str | Sequence) + ) + merged_beta: Final = reduce(merge_anthropic_beta_headers, beta_values, "") + return MappingProxyType( + {name: value for name, value in entries if name.lower() != "anthropic-beta"} + | ({"anthropic-beta": merged_beta} if merged_beta else {}) + ) + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -2481,13 +2498,8 @@ class BaseLLMHTTPHandler: forwarded_headers: Final = kwargs.get("headers", None) # Also check for extra_headers in kwargs (from config or direct calls) extra_headers_from_kwargs: Final = kwargs.get("extra_headers", None) - # Merge all header sources: forwarded < extra_headers < provider_specific - merged_headers: Final[Mapping[str, object]] = MappingProxyType( - { - **(forwarded_headers or {}), - **(extra_headers_from_kwargs or {}), - **(provider_specific_headers or {}), - } + merged_headers: Final = _merge_messages_header_sources( + forwarded_headers, extra_headers_from_kwargs, provider_specific_headers ) request_params: Final = reduce( delete_nested_value, 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/common_utils.py b/litellm/llms/gemini/common_utils.py index 9634c580858..88c6494e3f7 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -18,6 +18,25 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse +# 2.5 Flash has no configurable imageSize and supports ten aspect ratios, excluding these four +# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/2-5-flash-image +_GEMINI_2_5_FLASH_MODEL_FAMILY: Final = "2.5-flash" +_GEMINI_2_5_FLASH_UNSUPPORTED_ASPECT_RATIOS: Final = frozenset(("1:4", "4:1", "1:8", "8:1")) + +# 3.1 Flash supports the 512 tier, while other Gemini image models start at 1K +# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/3-1-flash-image +_GEMINI_512_IMAGE_SIZE_MODEL: Final = "gemini-3.1-flash-image" + +# Flash Lite supports only 1K, even when the requested dimensions are larger +# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/3-1-flash-lite-image +_GEMINI_1K_ONLY_IMAGE_SIZE_MODEL: Final = "gemini-3.1-flash-lite-image" + +# LiteLLM chooses the nearest resolution tier using the side length of a square with the same total number of pixels +# as the requested image +_GEMINI_IMAGE_SIZE_1K_THRESHOLD: Final = (512 + 1024) // 2 +_GEMINI_IMAGE_SIZE_2K_THRESHOLD: Final = (1024 + 2048) // 2 +_GEMINI_IMAGE_SIZE_4K_THRESHOLD: Final = (2048 + 4096) // 2 + GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, "1:4": 1 / 4, @@ -105,8 +124,8 @@ def map_openai_size_to_gemini_image_config(size: str, model: str) -> dict[str, s return None width, height = dimensions - image_config: Final = {"aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height)} - image_size: Final = _map_dimensions_to_gemini_image_size(width, height) + image_config: Final = {"aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height, model)} + image_size: Final = _map_dimensions_to_gemini_image_size(width, height, model) if is_gemini_image_model(model): if supports_gemini_image_size(model): image_config["imageSize"] = image_size @@ -123,7 +142,7 @@ def supports_gemini_image_size(model: str) -> bool: return bool(value) except Exception: pass - return "2.5-flash" not in model + return _GEMINI_2_5_FLASH_MODEL_FAMILY not in model def is_gemini_image_model(model: str) -> bool: @@ -333,24 +352,36 @@ def _parse_openai_image_size(size: str) -> tuple[int, int] | None: return width, height -def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int) -> str: - if (width, height) in GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO: - return GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO[(width, height)] +def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int, model: str) -> str: + aspect_ratios: Final = tuple( + ratio + for ratio in GEMINI_IMAGE_ASPECT_RATIOS + if _GEMINI_2_5_FLASH_MODEL_FAMILY not in model or ratio not in _GEMINI_2_5_FLASH_UNSUPPORTED_ASPECT_RATIOS + ) + known_ratio: Final = GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO.get((width, height)) + if known_ratio is not None and known_ratio in aspect_ratios: + return known_ratio requested_ratio: Final = width / height return min( - GEMINI_IMAGE_ASPECT_RATIOS, + aspect_ratios, key=lambda aspect_ratio: abs(math.log(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio] / requested_ratio)), ) -def _map_dimensions_to_gemini_image_size(width: int, height: int) -> str: - effective_square_side: Final = math.sqrt(width * height) - if effective_square_side < 768: - return "512" - if effective_square_side < 1536: +def _map_dimensions_to_gemini_image_size(width: int, height: int, model: str) -> str: + if _GEMINI_1K_ONLY_IMAGE_SIZE_MODEL in model: return "1K" - if effective_square_side < 3072: + + effective_square_side: Final = math.sqrt(width * height) + if effective_square_side < _GEMINI_IMAGE_SIZE_1K_THRESHOLD: + if is_gemini_image_model(model) and _GEMINI_512_IMAGE_SIZE_MODEL not in model: + return "1K" + else: + return "512" + if effective_square_side < _GEMINI_IMAGE_SIZE_2K_THRESHOLD: + return "1K" + if effective_square_side < _GEMINI_IMAGE_SIZE_4K_THRESHOLD: return "2K" return "4K" 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/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 1f130f84283..476e22fc088 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -78,6 +78,8 @@ class HuggingFaceRerankConfig(BaseRerankConfig): "documents", "top_n", "return_documents", + "truncate", + "truncation_direction", ] def map_cohere_rerank_params( @@ -104,6 +106,10 @@ class HuggingFaceRerankConfig(BaseRerankConfig): optional_rerank_params["return_text"] = v elif k == "top_n" and v is not None: optional_rerank_params["top_n"] = v + elif k == "truncate" and isinstance(v, bool): + optional_rerank_params["truncate"] = v + elif k == "truncation_direction" and v is not None: + optional_rerank_params["truncation_direction"] = v elif k == "documents" and v is not None: optional_rerank_params["texts"] = v elif k == "query" and v is not None: @@ -151,7 +157,6 @@ class HuggingFaceRerankConfig(BaseRerankConfig): # HuggingFace API expects return_text parameter, corresponding to our return_documents parameter request_body: Final = { "raw_scores": False, - "truncate": False, "truncation_direction": "Right", } diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 2c0acacff23..2fb9402ed03 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -6,16 +6,17 @@ Why separate file? Make it easy to see how transformation works Docs - https://docs.mistral.ai/api/ """ -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from itertools import accumulate, count from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload import httpx +from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import ( - handle_messages_with_content_list_to_str_conversion, + convert_content_list_to_str, strip_litellm_internal_message_fields, strip_none_values_from_message, ) @@ -29,8 +30,8 @@ from litellm.router_utils.reasoning_effort_capability import ( nearest_declared_reasoning_effort, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.mistral import MistralThinkingBlock, MistralToolCallMessage -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.mistral import MistralTextBlock, MistralThinkingBlock, MistralToolCallMessage +from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionTextObject from litellm.types.utils import ModelResponse, ModelResponseStream from litellm.utils import convert_to_model_response_object, supports_reasoning @@ -58,6 +59,10 @@ def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider: _NO_TOOL_CALL_TURN: Final[tuple[int, frozenset[str]]] = (-1, frozenset()) +def _is_json_array(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a decoded JSON array is a list + return isinstance(value, list) + + class MistralConfig(OpenAIGPTConfig): """ Reference: https://docs.mistral.ai/api/ @@ -118,11 +123,14 @@ class MistralConfig(OpenAIGPTConfig): def custom_llm_provider(self) -> str: return "mistral" - def get_supported_openai_params(self, model: str) -> list[str]: - is_magistral: Final = "magistral" in model.lower() - accepts_reasoning_effort: Final = is_magistral or supports_reasoning( + def _accepts_reasoning(self, model: str) -> bool: + return "magistral" in model.lower() or supports_reasoning( model=model, custom_llm_provider=self.custom_llm_provider ) + + def get_supported_openai_params(self, model: str) -> list[str]: + is_magistral: Final = "magistral" in model.lower() + accepts_reasoning_effort: Final = self._accepts_reasoning(model) return [ "stream", "temperature", @@ -260,25 +268,32 @@ class MistralConfig(OpenAIGPTConfig): The above statement is not valid now. Need to plan to remove all the #1,2,3 Mistral API supports content as a list. """ - messages = [self._strip_output_only_fields(m) for m in messages] + keeps_reasoning: Final = self._accepts_reasoning(model) + final_position: Final = len(messages) - 1 + replayed_messages: Final = [ + MistralConfig._with_replayed_reasoning(m, keeps_reasoning, position == final_position) + for position, m in enumerate(messages) + ] ## 1. If 'image_url' or 'file' in content, then transform with base class and mistral-specific handling - for m in messages: + for m in replayed_messages: _content_block = m.get("content") if _content_block and isinstance(_content_block, list): if any(c.get("type") in ["image_url", "file"] for c in _content_block): if is_async: - return self._transform_messages_async(messages, model) + return self._transform_messages_async(replayed_messages, model) else: - messages = self._transform_messages_sync(messages, model) - return messages + return self._transform_messages_sync(replayed_messages, model) ## 2. If content is list, then convert to string - messages = handle_messages_with_content_list_to_str_conversion(messages) + flattened_messages: Final = tuple( + MistralConfig._with_flattened_content(m, keeps_reasoning) for m in replayed_messages + ) ## 3. Handle name in message prepared_messages: Final = tuple( - MistralConfig._handle_tool_call_message(MistralConfig._handle_name_in_message(m)) for m in messages + MistralConfig._handle_tool_call_message(MistralConfig._handle_name_in_message(m)) + for m in flattened_messages ) empty_flags: Final = tuple(map(MistralConfig._is_empty_assistant_message, prepared_messages)) tool_call_turns: Final = accumulate( @@ -442,17 +457,163 @@ class MistralConfig(OpenAIGPTConfig): return cleaned_tools - @classmethod - def _strip_output_only_fields(cls, message: AllMessageValues) -> AllMessageValues: - """ - LiteLLM attaches output-only fields (``reasoning_content``, ``thinking_blocks``, - ``provider_specific_fields``) to assistant responses. Mistral's input schema forbids - unknown fields, so replaying them verbatim in a follow-up turn triggers a - 422 ``extra_forbidden``. Drop them before the request is sent. - """ + @staticmethod + def _is_mistral_think_chunk(part: object) -> bool: + match part: + case {"type": "thinking", "thinking": list()}: + return True + case _: + return False + + @staticmethod + def _text_item_text(item: object) -> str | None: + match item: + case {"type": "text", "text": str(text)}: + return text + case _: + return None + + @staticmethod + def _sendable_think_chunk( + chunk: object, items: Sequence[object], may_stay_open: bool + ) -> MistralThinkingBlock | None: + text_items: Final = [ + MistralTextBlock(type="text", text=text) + for text in map(MistralConfig._text_item_text, items) + if text is not None + ] + if not text_items: + return None + match chunk: + case {"closed": False} if not may_stay_open: + return MistralThinkingBlock(type="thinking", thinking=text_items) + case {"closed": bool(closed)}: + return MistralThinkingBlock(type="thinking", thinking=text_items, closed=closed) + case _: + return MistralThinkingBlock(type="thinking", thinking=text_items) + + @staticmethod + def _sendable_part(part: object, may_stay_open: bool) -> object | None: + match part: + case {"type": "thinking", "thinking": object() as items} if _is_json_array(items): + return MistralConfig._sendable_think_chunk(part, items, may_stay_open) + case _: + return part + + @staticmethod + def _with_sendable_think_chunks( + message: AllMessageValues, content_parts: tuple[object, ...], may_end_open: bool + ) -> AllMessageValues: + last_position: Final = len(content_parts) - 1 + sendable_parts: Final = [ + sendable + for position, part in enumerate(content_parts) + if (sendable := MistralConfig._sendable_part(part, may_end_open and position == last_position)) is not None + ] + return cast( # cast-ok: Mistral's ThinkChunk is outside the OpenAI assistant content union + AllMessageValues, {**message, "content": sendable_parts or None} + ) + + @staticmethod + def _thinking_block_text(block: object) -> str: + match block: + case {"type": "thinking", "thinking": str(thinking)}: + return thinking + case _: + return "" + + @staticmethod + def _replayed_reasoning_text(message: ChatCompletionAssistantMessage) -> str | None: + reasoning_content: Final = message.get("reasoning_content") + if isinstance(reasoning_content, str) and reasoning_content.strip(): + return reasoning_content + thinking_blocks: Final = message.get("thinking_blocks") + thinking_text: Final = ( + "".join(map(MistralConfig._thinking_block_text, thinking_blocks)) + if isinstance(thinking_blocks, list) + else "" + ) + return thinking_text if thinking_text.strip() else None + + @staticmethod + def _assistant_content_parts(message: ChatCompletionAssistantMessage) -> tuple[object, ...] | None: + match message.get("content"): + case None | "": + return () + case str(text): + return (ChatCompletionTextObject(type="text", text=text),) + case list(parts) | tuple(parts): + return tuple(parts) + case _: + return None + + @staticmethod + def _is_prefix_turn(message: object) -> bool: + match message: + case {"prefix": True}: + return True + case _: + return False + + @staticmethod + def _with_replayed_reasoning( + message: AllMessageValues, keeps_reasoning: bool, is_final_turn: bool + ) -> AllMessageValues: if message["role"] != "assistant": return message - return strip_litellm_internal_message_fields(message) + stripped_message: Final = strip_litellm_internal_message_fields(message) + content_parts: Final = MistralConfig._assistant_content_parts(message) + if not keeps_reasoning: + return MistralConfig._without_think_chunks(stripped_message, content_parts) + if content_parts is not None and any(map(MistralConfig._is_mistral_think_chunk, content_parts)): + return MistralConfig._with_sendable_think_chunks( + stripped_message, content_parts, is_final_turn and MistralConfig._is_prefix_turn(message) + ) + reasoning_text: Final = MistralConfig._replayed_reasoning_text(message) + has_turn_output: Final = bool(content_parts) or bool(message.get("tool_calls")) + if reasoning_text is None or content_parts is None or not has_turn_output: + return stripped_message + think_chunk: Final = MistralThinkingBlock( + type="thinking", thinking=[MistralTextBlock(type="text", text=reasoning_text)], closed=True + ) + return cast( # cast-ok: Mistral's ThinkChunk is outside the OpenAI assistant content union + AllMessageValues, {**stripped_message, "content": [think_chunk, *content_parts]} + ) + + @staticmethod + def _without_think_chunks(message: AllMessageValues, content_parts: tuple[object, ...] | None) -> AllMessageValues: + if content_parts is None or not any(map(MistralConfig._is_mistral_think_chunk, content_parts)): + return message + kept_parts: Final = [part for part in content_parts if not MistralConfig._is_mistral_think_chunk(part)] + return cast( # cast-ok: an assistant turn takes a content list or no content + AllMessageValues, {**message, "content": kept_parts or None} + ) + + @staticmethod + def _with_flattened_content(message: AllMessageValues, keeps_reasoning: bool) -> AllMessageValues: + flattened_text: Final = convert_content_list_to_str(message) + match message: + case {"role": "assistant", "content": list(content)} if keeps_reasoning and any( + map(MistralConfig._is_mistral_think_chunk, content) + ): + text_parts: Final = ( + (ChatCompletionTextObject(type="text", text=flattened_text),) if flattened_text else () + ) + think_chunks: Final = tuple(filter(MistralConfig._is_mistral_think_chunk, content)) + sent_chunks: Final = tuple( + sendable + for chunk in think_chunks + if (sendable := MistralConfig._sendable_part(chunk, not text_parts)) is not None + ) + return cast( # cast-ok: Mistral's ThinkChunk is outside the OpenAI assistant content union + AllMessageValues, {**message, "content": [*sent_chunks, *text_parts]} + ) + case _ if flattened_text: + return cast( # cast-ok: every message role takes str content + AllMessageValues, {**message, "content": flattened_text} + ) + case _: + return message @classmethod def _handle_name_in_message(cls, message: AllMessageValues) -> AllMessageValues: diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 0c13b1cd1d3..1e70141f939 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -5,12 +5,15 @@ Ollama /chat/completion calls handled in llm_http_handler.py """ from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Any, Final, Protocol, TypedDict +import httpx from typing_extensions import NotRequired, ReadOnly import litellm from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.types.mcp import has_header from litellm.types.utils import EmbeddingResponse @@ -82,6 +85,15 @@ def _process_ollama_embedding_response( return model_response +def _ollama_embedding_headers(api_key: str | None, headers: Mapping[str, str] | None) -> Mapping[str, str]: + return MappingProxyType( + { + **({"Authorization": f"Bearer {api_key}"} if api_key and not has_header(headers, "Authorization") else {}), + **(headers or {}), + } + ) + + async def ollama_aembeddings( api_base: str, model: str, @@ -90,10 +102,19 @@ async def ollama_aembeddings( optional_params: Mapping[str, object], logging_obj: object, encoding: TokenEncoder | None, -): + *, + api_key: str | None = None, + headers: Mapping[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, +) -> EmbeddingResponse: data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = await litellm.module_level_aclient.post(url=ensure_api_route(api_base, "api/embed"), json=data) + response: Final = await litellm.module_level_aclient.post( + url=ensure_api_route(api_base, "api/embed"), + json=data, + headers=dict(_ollama_embedding_headers(api_key, headers)), + timeout=timeout, + ) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( @@ -114,10 +135,19 @@ def ollama_embeddings( model_response: EmbeddingResponse, logging_obj: object, encoding: TokenEncoder | None = None, -): + *, + api_key: str | None = None, + headers: Mapping[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, +) -> EmbeddingResponse: data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = litellm.module_level_client.post(url=ensure_api_route(api_base, "api/embed"), json=data) + response: Final = litellm.module_level_client.post( + url=ensure_api_route(api_base, "api/embed"), + json=data, + headers=dict(_ollama_embedding_headers(api_key, headers)), + timeout=timeout, + ) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( 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/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 07822bb17b3..9cb074a177e 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -499,8 +499,13 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools - final_request_params: Final = dict( - ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params) + final_request_params: Final = cast( # cast-ok: typed dict spread is inferred as Unknown by basedpyright + dict[str, object], + { + "model": model, + "input": replay_safe_input, + **response_api_optional_request_params, + }, ) return final_request_params @@ -1209,8 +1214,13 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools - data: Final = dict( - ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params) + data: Final = cast( # cast-ok: typed dict spread is inferred as Unknown by basedpyright + dict[str, object], + { + "model": model, + "input": replay_safe_input, + **response_api_optional_request_params, + }, ) return url, data diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 0034ed2b9d4..2b71707239b 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -58,6 +58,12 @@ "max_completion_tokens": "max_tokens" } }, + "zerogpu": { + "base_url": "https://api.zerogpu.ai/v1", + "api_key_env": "ZEROGPU_API_KEY", + "api_base_env": "ZEROGPU_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] + }, "chutes": { "base_url": "https://llm.chutes.ai/v1/", "api_key_env": "CHUTES_API_KEY", diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 27f2aedb300..0fa80cb2cad 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1,3 +1,4 @@ +import math import re from collections.abc import Callable, Iterable, Mapping, Sequence from copy import deepcopy @@ -18,6 +19,11 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.gemini.common_utils import ( + GEMINI_IMAGE_ASPECT_RATIOS, + GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO, + map_openai_size_to_gemini_image_config, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ( VERTEX_AI_PROVIDER_METADATA_FIELDS, @@ -35,6 +41,31 @@ GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/" +# https://cloud.google.com/vertex-ai/generative-ai/docs/image/configure-aspect-ratio +_IMAGEN_ASPECT_RATIOS: Final = MappingProxyType( + {"1:1": 1.0, "3:4": 3 / 4, "4:3": 4 / 3, "9:16": 9 / 16, "16:9": 16 / 9} +) + + +def map_vertex_image_size(size: object, model: str, *, imagen: bool = False) -> Mapping[str, str] | None: + if not isinstance(size, str) or (image_config := map_openai_size_to_gemini_image_config(size, model)) is None: + return None + + width, height = map(int, size.lower().split("x")) + requested_ratio: Final = width / height + aspect_ratio: Final = ( + min(_IMAGEN_ASPECT_RATIOS, key=lambda ratio: abs(math.log(_IMAGEN_ASPECT_RATIOS[ratio] / requested_ratio))) + if imagen + else image_config["aspectRatio"] + ) + if ( + not math.isclose(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio], requested_ratio) + and GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO.get((width, height)) != aspect_ratio + ): + verbose_logger.warning("Vertex AI size %s is approximated with aspectRatio %s", size, aspect_ratio) + return MappingProxyType({"aspectRatio": aspect_ratio} if imagen else image_config) + + def gemini_video_metadata_from_openai(video_metadata: Mapping[str, object]) -> dict[str, object]: return { gemini_key: video_metadata[openai_key] diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 6d071363823..0b90f74ea6d 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -10,7 +10,7 @@ from httpx._types import RequestFiles import litellm from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -67,7 +67,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): return {} size: Final = image_edit_optional_params.get("size") - return {"aspectRatio": self._map_size_to_aspect_ratio(size or "")} + return dict(map_vertex_image_size(size, model) or {"aspectRatio": "1:1"}) def _resolve_vertex_project(self) -> str | None: return ( @@ -176,11 +176,11 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): contents: Final[dict[str, object]] = {"role": "USER", "parts": parts} # Add image-specific configuration - image_config: Final = ( - {"aspect_ratio": image_edit_optional_request_params["aspectRatio"]} - if "aspectRatio" in image_edit_optional_request_params - else None - ) + image_config: Final = { + field: image_edit_optional_request_params[param] + for param, field in (("aspectRatio", "aspect_ratio"), ("imageSize", "image_size")) + if param in image_edit_optional_request_params + } generation_config: Final[dict[str, object]] = { key: value for key, value in (("response_modalities", ["IMAGE"]), ("image_config", image_config)) if value @@ -222,17 +222,6 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): model_response.data = cast(list[OpenAIImage], data_list) return model_response - def _map_size_to_aspect_ratio(self, size: str) -> str: - """Map OpenAI size format to Gemini aspect ratio format""" - aspect_ratio_map: Final = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[HttpxPartType]: images: Final[list[FileTypes]] = image if isinstance(image, list) else [image] return [ diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index deaa32f7308..87b0dfdbcbc 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -15,7 +15,7 @@ import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -70,13 +70,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): if "n" in filtered_params: mapped_params["sampleCount"] = filtered_params["n"] - if "size" in filtered_params: - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(filtered_params["size"]) - if "mask" in filtered_params: mapped_params["mask"] = filtered_params["mask"] - return mapped_params + size_params: Final = map_vertex_image_size(filtered_params.get("size"), model, imagen=True) or {} + return {**mapped_params, **size_params} def _resolve_vertex_project(self) -> str | None: return ( @@ -189,6 +187,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): "sampleCount": sample_count, "editMode": edit_mode, "editConfig": {"baseSteps": base_steps}, + **( + {"aspectRatio": image_edit_optional_request_params["aspectRatio"]} + if "aspectRatio" in image_edit_optional_request_params + else {} + ), } # Set default values for Vertex AI-specific parameters (not configurable by users via OpenAI API) @@ -239,17 +242,6 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): model_response.data = cast(list[OpenAIImage], data_list) return model_response - def _map_size_to_aspect_ratio(self, size: str) -> str: - """Map OpenAI size format to Imagen aspect ratio format""" - aspect_ratio_map: Final = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _prepare_reference_images( self, image: FileTypes | list[FileTypes], diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index f6df772dd56..e1049ef8ece 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -14,7 +14,7 @@ from litellm.llms.gemini.common_utils import ( get_gemini_image_web_search_requests, map_gemini_image_tools_params, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -89,8 +89,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): if k == "n": mapped_params["candidate_count"] = v elif k == "size": - # Map OpenAI size format to Gemini aspectRatio - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) + continue elif k in ("aspectRatio", "aspect_ratio"): mapped_params["aspectRatio"] = v elif k in ("imageSize", "image_size"): @@ -103,22 +102,23 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): elif k not in ("tools", "web_search_options", "imageConfig"): mapped_params[k] = v + if "size" in non_default_params and "size" not in optional_params: + image_params: Final = _JSON_OBJECT.validate_python(mapped_params) + image_config: Final = _JSON_OBJECT.validate_python(image_params.get("imageConfig", {})) + size_params: Final = { + key: value + for key, value in ( + map_vertex_image_size(non_default_params["size"], model) or {"aspectRatio": "1:1"} + ).items() + if key not in image_config + } + return map_gemini_image_tools_params( + _JSON_OBJECT.validate_python(non_default_params), {**size_params, **image_params} + ) + mapped_params = map_gemini_image_tools_params(non_default_params, mapped_params) return mapped_params - def _map_size_to_aspect_ratio(self, size: str) -> str: - """ - Map OpenAI size format to Gemini aspect ratio format - """ - aspect_ratio_map: Final = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _resolve_vertex_project(self) -> str | None: return ( getattr(self, "_vertex_project", None) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 927c7b74e67..daf87b1ec58 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -11,6 +11,7 @@ from litellm.llms.base_llm.image_generation.transformation import ( ) from litellm.llms.vertex_ai.common_utils import ( get_vertex_base_url, + map_vertex_image_size, pop_vertex_request_labels, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM @@ -72,25 +73,14 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): mapped_params["sampleCount"] = v elif k == "size": # Map OpenAI size format to Imagen aspectRatio - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) + mapped_params["aspectRatio"] = ( + map_vertex_image_size(v, model, imagen=True) or {"aspectRatio": "1:1"} + )["aspectRatio"] else: mapped_params[k] = v return mapped_params - def _map_size_to_aspect_ratio(self, size: str) -> str: - """ - Map OpenAI size format to Imagen aspect ratio format - """ - aspect_ratio_map: Final = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _resolve_vertex_project(self) -> str | None: return ( getattr(self, "_vertex_project", None) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 29253359003..363e5596cb9 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -1,5 +1,7 @@ # What is this? ## Handler file for calling claude-3 on vertex ai +import base64 +import json from collections.abc import Mapping from typing import TYPE_CHECKING, Final @@ -19,6 +21,19 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) +_VERTEX_LABELS_HEADER: Final = "X-Vertex-AI-Labels" +_LABELS: Final = TypeAdapter(Mapping[object, object]) + + +def _vertex_labels_header(labels: object, header_names: Mapping[str, object]) -> str | None: + if not isinstance(labels, dict) or any(name.lower() == _VERTEX_LABELS_HEADER.lower() for name in header_names): + return None + string_labels: Final = { + key: value + for key, value in _LABELS.validate_python(labels).items() + if isinstance(key, str) and isinstance(value, str) + } + return base64.b64encode(json.dumps(string_labels).encode()).decode() if string_labels else None class VertexAIError(Exception): @@ -102,6 +117,12 @@ class VertexAIAnthropicConfig(AnthropicConfig): litellm_params: dict, headers: dict, ) -> dict: + params: Final = _NON_DEFAULT_PARAMS.validate_python(optional_params) + labels_header: Final = _vertex_labels_header(params.get("labels"), _NON_DEFAULT_PARAMS.validate_python(headers)) + if "labels" in params: + del optional_params["labels"] # rebind-ok: labels travel in the X-Vertex-AI-Labels header, not the body + if labels_header is not None: + headers[_VERTEX_LABELS_HEADER] = labels_header # rebind-ok: outbound headers are returned by mutation data: Final = super().transform_request( model=model, messages=messages, diff --git a/litellm/main.py b/litellm/main.py index 9f412e1f24b..41c4302499e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -123,6 +123,7 @@ from litellm.llms.bedrock.common_utils import ( BedrockModelInfo, NativeOutputFallback, bedrock_route_for_request, + bedrock_strict_tools_fallback, deployment_model_info, without_bedrock_native_output_opt_in, without_bedrock_route_prefix, @@ -206,7 +207,12 @@ from litellm.utils import ( from ._logging import verbose_logger from .caching.caching import disable_cache, enable_cache, update_cache -from .litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params, safe_deep_copy +from .litellm_core_utils.core_helpers import ( + as_str_mapping, + drops_unsupported_params, + normalize_drop_params, + safe_deep_copy, +) from .litellm_core_utils.fallback_utils import ( async_completion_with_fallbacks, completion_with_fallbacks, @@ -4223,6 +4229,7 @@ def _complete_sagemaker(ctx: CompletionDispatchContext) -> _CompletionDispatchRe _ADDITIONAL_DROP_PARAMS_ADAPTER: Final = TypeAdapter(list[str]) _OPTIONAL_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OLLAMA_EMBEDDING_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str]) _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) @@ -4378,6 +4385,11 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu else None ), converse_model, + strict_tools_fallback=( + partial(bedrock_strict_tools_fallback, converse_model, converse_params, dict(litellm_params)) + if drops_unsupported_params(litellm_params.get("drop_params")) + else None + ), ) elif bedrock_route == "converse_like": model = model.replace("converse_like/", "") @@ -6112,7 +6124,18 @@ def completion( optional_params, rs_val = strip_reasoning_summary_aliases_from_optional_params(optional_params) - if isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort: + if custom_llm_provider == "openai" and isinstance(resolved_reasoning_effort, dict): + from litellm.responses.utils import openai_bridge_reasoning + + optional_params["reasoning_effort"] = openai_bridge_reasoning( + requested=resolved_reasoning_effort, + effort_kept="reasoning_effort" in optional_params, + summary_alias=rs_val, + model=model, + drop_params=drop_unsupported_reasoning, + additional_drop_params=routing_kwargs.get("additional_drop_params"), + ) + elif isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort: optional_params["reasoning_effort"] = resolved_reasoning_effort elif rs_val is not None: eff: Final = optional_params.get("reasoning_effort", resolved_reasoning_effort) @@ -7309,6 +7332,9 @@ def embedding( logging_obj=logging, optional_params=optional_params, model_response=EmbeddingResponse(), + api_key=api_key, + headers=_OLLAMA_EMBEDDING_HEADERS_ADAPTER.validate_python(headers), + timeout=timeout, ) elif custom_llm_provider == "sagemaker": response = sagemaker_llm.embedding( @@ -8375,7 +8401,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, @@ -8537,10 +8563,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) @@ -8550,8 +8582,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 8207c3122f9..3e0e86012a4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4039,6 +4039,7 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure_ai/gpt-oss-120b": { + "deprecation_date": "2027-09-15", "input_cost_per_token": 1.5e-07, "output_cost_per_token": 6e-07, "litellm_provider": "azure_ai", @@ -4243,6 +4244,54 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, @@ -4904,6 +4953,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -4922,6 +4976,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -5009,6 +5068,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5048,6 +5112,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5113,6 +5182,11 @@ "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.1-codex": { @@ -5188,6 +5262,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5344,6 +5423,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -5428,6 +5512,11 @@ "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/global/gpt-5.1-codex": { @@ -5680,6 +5769,11 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5717,6 +5811,11 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5754,6 +5853,11 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5791,6 +5895,11 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6203,6 +6312,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -6221,6 +6335,11 @@ "mode": "chat", "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -6710,6 +6829,11 @@ "default_reasoning_effort": "none", "input_cost_per_token_batches": 6.25e-07, "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_minimal_reasoning_effort": true, "cache_read_input_token_cost_batches": 6.25e-08 @@ -6797,6 +6921,11 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6835,6 +6964,11 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6906,6 +7040,11 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_batches": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6944,6 +7083,11 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_batches": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6981,6 +7125,11 @@ "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -7016,6 +7165,11 @@ "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -7051,6 +7205,11 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/responses" @@ -7112,6 +7271,11 @@ "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/gpt-5-chat": { @@ -7459,6 +7623,11 @@ "output_cost_per_token": 1.4e-05, "output_cost_per_token_batches": 7e-06, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -7498,6 +7667,11 @@ "output_cost_per_token": 1.4e-05, "output_cost_per_token_batches": 7e-06, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -7567,6 +7741,11 @@ "mode": "responses", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/responses" @@ -7599,9 +7778,9 @@ "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -7637,9 +7816,9 @@ "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -7996,9 +8175,9 @@ "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8043,9 +8222,9 @@ "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8095,9 +8274,9 @@ "output_cost_per_token_priority": 4e-05, "output_cost_per_token_above_272k_tokens_priority": 6e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -8157,9 +8336,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8220,9 +8399,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8283,9 +8462,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_flex": 6e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8346,9 +8525,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_flex": 6e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8409,9 +8588,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8472,9 +8651,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -8519,9 +8698,9 @@ "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8567,9 +8746,9 @@ "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8615,9 +8794,9 @@ "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -8663,9 +8842,9 @@ "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -8718,7 +8897,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8770,7 +8949,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8800,6 +8979,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -8815,9 +9041,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8863,9 +9089,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8896,6 +9122,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/us/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, + "cache_read_input_token_cost": 5.5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "input_cost_per_token_above_272k_tokens": 2.2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "output_cost_per_token_above_272k_tokens": 8.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/us/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, @@ -8911,9 +9184,9 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8944,6 +9217,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/eu/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_above_272k_tokens": 2.4e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_272k_tokens": 9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/eu/gpt-6.1-sol": { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, @@ -8959,9 +9279,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8992,6 +9312,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/apac/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_above_272k_tokens": 2.4e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_272k_tokens": 9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/apac/gpt-6.1-sol": { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, @@ -9007,9 +9374,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -9134,9 +9501,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9188,9 +9555,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9243,9 +9610,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, "output_cost_per_token_priority": 2.64e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9298,9 +9665,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-06, "output_cost_per_token_priority": 2.64e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9345,9 +9712,9 @@ "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9393,9 +9760,9 @@ "output_cost_per_token": 5.5e-07, "output_cost_per_token_above_272k_tokens": 8.25e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -9448,7 +9815,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_priority": 3.3e-05, "output_cost_per_token_priority": 2.2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -9536,9 +9903,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9590,9 +9957,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9645,9 +10012,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, "output_cost_per_token_priority": 2.64e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9700,9 +10067,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-06, "output_cost_per_token_priority": 2.64e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9760,9 +10127,9 @@ "output_cost_per_token_above_272k_tokens_priority": 9e-05, "output_cost_per_token_batches": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9814,9 +10181,9 @@ "output_cost_per_token_above_272k_tokens_batches": 2.475e-05, "output_cost_per_token_batches": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9868,9 +10235,9 @@ "output_cost_per_token_above_272k_tokens_batches": 2.475e-05, "output_cost_per_token_batches": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9921,9 +10288,9 @@ "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9981,9 +10348,9 @@ "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10034,9 +10401,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10086,9 +10453,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10139,9 +10506,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10191,9 +10558,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10239,9 +10606,9 @@ "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/batch", @@ -10282,9 +10649,9 @@ "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/batch", @@ -10322,9 +10689,9 @@ "mode": "chat", "output_cost_per_token": 4.5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10372,9 +10739,9 @@ "mode": "chat", "output_cost_per_token": 4.5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10423,9 +10790,9 @@ "mode": "chat", "output_cost_per_token": 1.25e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10470,9 +10837,9 @@ "mode": "chat", "output_cost_per_token": 1.25e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10934,6 +11301,11 @@ "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -10965,6 +11337,11 @@ "mode": "chat", "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -10996,9 +11373,9 @@ "mode": "responses", "output_cost_per_token": 4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -11065,6 +11442,11 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/batch", @@ -11095,6 +11477,11 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/batch", @@ -11124,6 +11511,11 @@ "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -11155,6 +11547,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "output_cost_per_token_batches": 2.2e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": false, @@ -11268,6 +11665,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11305,6 +11707,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11428,6 +11835,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -11446,6 +11858,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -11533,6 +11950,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11572,6 +11994,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11608,6 +12035,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11673,6 +12105,11 @@ "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.1-codex": { @@ -11800,6 +12237,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11832,6 +12274,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11898,6 +12345,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": false, @@ -11918,6 +12370,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": false, @@ -11935,6 +12392,7 @@ "output_cost_per_second": 0.0001 }, "azure_ai/Cohere-embed-v3-english": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 512, @@ -11945,6 +12403,7 @@ "supports_embedding_image_input": true }, "azure_ai/Cohere-embed-v3-multilingual": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 512, @@ -11955,6 +12414,7 @@ "supports_embedding_image_input": true }, "azure_ai/Codestral-2501": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 3e-07, "litellm_provider": "azure_ai", "max_input_tokens": 256000, @@ -11974,6 +12434,7 @@ "supports_native_streaming": true }, "azure_ai/FLUX-1.1-pro": { + "deprecation_date": "2027-01-05", "litellm_provider": "azure_ai", "max_input_tokens": 5000, "mode": "image_generation", @@ -11984,6 +12445,7 @@ ] }, "azure_ai/FLUX.1-Kontext-pro": { + "deprecation_date": "2027-01-05", "litellm_provider": "azure_ai", "max_input_tokens": 5000, "mode": "image_generation", @@ -11994,6 +12456,7 @@ ] }, "azure_ai/flux.2-pro": { + "deprecation_date": "2027-01-20", "litellm_provider": "azure_ai", "max_input_tokens": 32000, "mode": "image_generation", @@ -12004,6 +12467,7 @@ ] }, "azure_ai/FLUX.2-flex": { + "deprecation_date": "2027-01-20", "input_cost_per_pixel": 5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 32000, @@ -12368,6 +12832,7 @@ "supports_tool_choice": true }, "azure_ai/Llama-3.3-70B-Instruct": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 7.1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12379,6 +12844,7 @@ "supports_tool_choice": true }, "azure_ai/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "deprecation_date": "2026-11-19", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -12391,6 +12857,7 @@ "supports_vision": true }, "azure_ai/Llama-4-Scout-17B-16E-Instruct": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 10000000, @@ -12624,6 +13091,7 @@ ] }, "azure_ai/cohere-command-a": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12696,6 +13164,7 @@ "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models" }, "azure_ai/cohere-rerank-v4.0-pro": { + "deprecation_date": "2027-09-13", "input_cost_per_query": 0.0025, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -12706,6 +13175,7 @@ "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-cohere-rerank-4-0-in-microsoft-foundry/4477076" }, "azure_ai/cohere-rerank-v4.0-fast": { + "deprecation_date": "2027-09-13", "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -12716,6 +13186,7 @@ "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-cohere-rerank-4-0-in-microsoft-foundry/4477076" }, "azure_ai/deepseek-v3.2": { + "deprecation_date": "2027-11-01", "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12730,6 +13201,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3.2-speciale": { + "deprecation_date": "2027-11-01", "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12800,6 +13272,7 @@ "supports_tool_choice": true }, "azure_ai/embed-v-4-0": { + "deprecation_date": "2027-06-01", "input_cost_per_token": 1.2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12818,6 +13291,7 @@ "input_cost_per_image_token": 4.7e-07 }, "azure_ai/grok-4": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 262000, @@ -12906,6 +13380,7 @@ "supports_web_search": true }, "azure_ai/grok-4-1-fast-non-reasoning": { + "deprecation_date": "2027-05-19", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -12919,6 +13394,7 @@ "supports_web_search": true }, "azure_ai/grok-4-1-fast-reasoning": { + "deprecation_date": "2027-05-19", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -12933,6 +13409,7 @@ "supports_web_search": true }, "azure_ai/grok-code-fast-1": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 256000, @@ -13100,7 +13577,7 @@ "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { - "deprecation_date": "2026-12-09", + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -14619,11 +15096,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14638,6 +15116,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -14647,11 +15126,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14665,7 +15145,8 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -14787,11 +15268,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14806,6 +15288,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -14815,11 +15298,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14833,7 +15317,8 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -15197,6 +15682,7 @@ "supports_tool_choice": true }, "cerebras/llama-3.3-70b": { + "deprecation_date": "2026-02-16", "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15207,6 +15693,7 @@ "supports_tool_choice": true }, "cerebras/llama3.1-70b": { + "deprecation_date": "2025-01-17", "input_cost_per_token": 6e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15217,6 +15704,7 @@ "supports_tool_choice": true }, "cerebras/llama3.1-8b": { + "deprecation_date": "2026-05-27", "input_cost_per_token": 1e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15241,6 +15729,7 @@ "supports_tool_choice": true }, "cerebras/qwen-3-32b": { + "deprecation_date": "2026-02-16", "input_cost_per_token": 4e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15552,6 +16041,7 @@ "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "deprecation_date": "2027-04-08", "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, @@ -15570,7 +16060,8 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, "input_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_batches": 7.5e-06 + "output_cost_per_token_batches": 7.5e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "claude-opus-4-5-20251101": { "supports_web_search": true, @@ -15764,7 +16255,6 @@ "us": 1.1 }, "supports_output_config": true, - "supports_speed": true, "prompt_cache_min_tokens": 2048, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, @@ -15807,7 +16297,6 @@ "us": 1.1 }, "supports_output_config": true, - "supports_speed": true, "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { @@ -20520,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, @@ -20532,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, @@ -20548,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, @@ -20560,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, @@ -20578,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, @@ -20601,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, @@ -20611,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, @@ -20624,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, @@ -20634,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, @@ -20647,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, @@ -20657,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, @@ -20671,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, @@ -20681,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, @@ -20695,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, @@ -20706,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, @@ -20721,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, @@ -20732,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, @@ -20748,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, @@ -20759,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, @@ -20855,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, @@ -20865,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, @@ -20878,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, @@ -20888,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, @@ -20901,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, @@ -21057,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", @@ -21087,13 +21592,19 @@ "supports_video_input": true }, "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", @@ -21107,13 +21618,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", @@ -21129,7 +21646,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, @@ -21139,12 +21658,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", @@ -21154,13 +21670,20 @@ "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 }, "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", @@ -21209,13 +21732,17 @@ "supports_tool_choice": true }, "databricks/databricks-gemini-3-8-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": "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_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21233,13 +21760,17 @@ "supports_vision": true }, "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": "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_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21258,9 +21789,9 @@ }, "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_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, @@ -21268,8 +21799,8 @@ "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, + "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21530,7 +22061,9 @@ "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, @@ -21540,6 +22073,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": [ @@ -21613,8 +22147,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, @@ -21624,6 +22161,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": [ @@ -21640,10 +22178,13 @@ "supports_vision": true }, "databricks/databricks-gpt-5-6-terra": { - "cache_creation_input_token_cost": 3.12501e-06, - "cache_read_input_token_cost": 2.4997e-07, - "input_cost_per_token": 2.49998e-06, - "input_dbu_cost_per_token": 3.5714e-05, + "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, @@ -21651,8 +22192,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": "chat", - "output_cost_per_token": 1.500002e-05, - "output_dbu_cost_per_token": 0.000214286, + "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": [ "text", @@ -21668,10 +22210,13 @@ "supports_vision": true }, "databricks/databricks-gpt-5-6-luna": { - "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.0003e-07, - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, + "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, "max_output_tokens": 128000, @@ -21679,8 +22224,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": "chat", - "output_cost_per_token": 5.99998e-06, - "output_dbu_cost_per_token": 8.5714e-05, + "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": [ "text", @@ -21698,7 +22244,9 @@ "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, @@ -21708,6 +22256,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": [ @@ -21727,7 +22276,8 @@ "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, @@ -21735,8 +22285,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", @@ -22055,6 +22606,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, @@ -34478,6 +35048,7 @@ "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini-tts": { + "deprecation_date": "2027-01-06", "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", @@ -35095,6 +35666,7 @@ "cache_read_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.25e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, @@ -35107,6 +35679,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, @@ -35184,6 +35763,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, @@ -35351,6 +35937,7 @@ "cache_read_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2027-04-01", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -35359,6 +35946,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, @@ -35410,6 +36004,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, @@ -35525,6 +36126,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, @@ -35574,6 +36182,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, @@ -35615,6 +36230,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, @@ -35650,15 +36272,15 @@ "gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_12hr": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, - "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, "cache_creation_input_token_cost_above_12hr_priority": 5e-05, - "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -35702,6 +36324,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": { @@ -35786,6 +36417,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": { @@ -35828,11 +36468,11 @@ "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_12hr": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, - "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, @@ -35878,6 +36518,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": { @@ -35961,6 +36610,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": { @@ -36042,6 +36700,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": { @@ -36125,6 +36792,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": { @@ -36207,6 +36883,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": { @@ -36444,6 +37129,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", @@ -36491,6 +37184,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": { @@ -36553,6 +37254,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": { @@ -36607,6 +37316,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": { @@ -36658,6 +37375,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": { @@ -36717,6 +37442,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": [ @@ -36774,6 +37507,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": [ @@ -36823,6 +37564,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": { @@ -36875,6 +37624,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": { @@ -36929,6 +37686,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": { @@ -36983,6 +37748,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": { @@ -37034,6 +37807,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": { @@ -37075,6 +37856,7 @@ "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, + "deprecation_date": "2027-04-01", "input_cost_per_token": 2e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, @@ -37085,6 +37867,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": { @@ -37122,6 +37912,7 @@ "supports_minimal_reasoning_effort": false }, "gpt-5-pro": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -37130,6 +37921,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, @@ -37172,6 +37970,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, @@ -37222,6 +38027,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, @@ -37332,6 +38144,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, @@ -37401,6 +38220,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -37413,6 +38233,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, @@ -37464,6 +38291,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, @@ -37501,6 +38335,7 @@ "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_batches": 2.5e-09, "cache_read_input_token_cost_flex": 2.5e-09, + "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, @@ -37511,6 +38346,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, @@ -37559,6 +38401,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, @@ -38210,6 +39059,7 @@ }, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -38230,6 +39080,7 @@ }, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -38249,6 +39100,7 @@ "search_context_size_medium": 0.005 }, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -40996,9 +41848,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/devstral-small-2-25-12", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41010,9 +41867,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41024,9 +41886,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41331,6 +42198,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, @@ -41547,13 +42430,13 @@ "mode": "embedding" }, "mistral/mistral-large-4": { - "cache_read_input_token_cost": 6.8e-08, - "input_cost_per_token": 6.8e-07, + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.36e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 2.09e-06, + "output_cost_per_token": 4.18e-06, "reasoning_effort_levels": [ "none", "high" @@ -41568,13 +42451,13 @@ "supports_vision": true }, "mistral/mistral-large-4-0": { - "cache_read_input_token_cost": 6.8e-08, - "input_cost_per_token": 6.8e-07, + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.36e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 2.09e-06, + "output_cost_per_token": 4.18e-06, "reasoning_effort_levels": [ "none", "high" @@ -41721,8 +42604,13 @@ "max_output_tokens": 8191, "mode": "chat", "output_cost_per_token": 3e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41921,6 +42809,7 @@ }, "moonshot/kimi-k2.5": { "cache_read_input_token_cost": 1e-07, + "deprecation_date": "2026-08-31", "input_cost_per_token": 6e-07, "litellm_provider": "moonshot", "max_input_tokens": 262144, @@ -41951,7 +42840,31 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3.1-pro": { + "cache_creation_input_token_cost": 2e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://platform.kimi.ai/docs/pricing/chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, "moonshot/kimi-k3": { + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "moonshot", @@ -41966,6 +42879,7 @@ ], "source": "https://platform.kimi.ai/docs/pricing/chat-k3", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -41973,6 +42887,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-128k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -41985,6 +42900,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-128k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -41998,6 +42914,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-32k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 1e-06, "litellm_provider": "moonshot", "max_input_tokens": 32768, @@ -42010,6 +42927,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-32k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 1e-06, "litellm_provider": "moonshot", "max_input_tokens": 32768, @@ -42023,6 +42941,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-8k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-07, "litellm_provider": "moonshot", "max_input_tokens": 8192, @@ -42035,6 +42954,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-8k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-07, "litellm_provider": "moonshot", "max_input_tokens": 8192, @@ -42048,6 +42968,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-auto": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -43240,6 +44161,7 @@ "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -43252,6 +44174,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, @@ -43360,6 +44289,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, @@ -43430,6 +44366,7 @@ "supports_vision": false }, "o3-pro": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -43438,6 +44375,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, @@ -43475,6 +44419,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, @@ -47290,13 +48241,13 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 4.5e-06, + "output_cost_per_token": 7.5e-06, "source": "https://www.together.ai/pricing", "supports_prompt_caching": true }, @@ -47806,6 +48757,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 9e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, + "deprecation_date": "2027-04-08", "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, @@ -47824,7 +48776,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -52901,15 +53853,21 @@ "xai/grok-4.7": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 8e-07, + "cache_read_input_token_cost_batches": 4e-07, "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 3.2e-06, + "input_cost_per_token_batches": 1.6e-06, "litellm_provider": "xai", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, + "output_cost_per_token_above_200k_tokens_batches": 9.6e-06, + "output_cost_per_token_batches": 4.8e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -56990,6 +57948,7 @@ ] }, "gpt-4o-mini-transcribe-2025-12-15": { + "deprecation_date": "2027-02-26", "input_cost_per_audio_token": 1.25e-06, "input_cost_per_second": 5e-05, "input_cost_per_token": 1.25e-06, @@ -64527,6 +65486,7 @@ ] }, "xai/grok-imagine-image-pro": { + "deprecation_date": "2026-05-15", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -64933,9 +65893,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -65279,6 +66244,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", @@ -65364,6 +66503,7 @@ ] }, "scaleway/glm-5.2": { + "deprecation_date": "2026-11-16", "input_cost_per_token": 1.8e-06, "litellm_provider": "scaleway", "max_input_tokens": 256000, @@ -67922,6 +69062,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-mini": { @@ -67936,6 +69081,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { @@ -67974,6 +69124,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5-codex": { @@ -67998,6 +69153,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5-nano": { @@ -68010,6 +69170,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5-pro": { @@ -68020,6 +69185,11 @@ "mode": "responses", "output_cost_per_token": 0.000132, "output_cost_per_token_batches": 6.6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.1-codex-max": { @@ -68044,6 +69214,11 @@ "output_cost_per_token": 1.54e-05, "output_cost_per_token_batches": 7.7e-06, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.2-codex": { @@ -68062,6 +69237,11 @@ "mode": "responses", "output_cost_per_token": 0.0001848, "output_cost_per_token_batches": 9.24e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.3-codex": { @@ -68074,6 +69254,11 @@ "mode": "responses", "output_cost_per_token": 1.54e-05, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-mini": { @@ -68089,6 +69274,11 @@ "output_cost_per_token": 4.95e-06, "output_cost_per_token_batches": 2.475e-06, "output_cost_per_token_priority": 9.9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-nano": { @@ -68101,6 +69291,11 @@ "mode": "chat", "output_cost_per_token": 1.375e-06, "output_cost_per_token_batches": 6.875e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-pro": { @@ -68115,6 +69310,11 @@ "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-astra": { @@ -68131,6 +69331,11 @@ "mode": "chat", "output_cost_per_token": 6e-05, "output_cost_per_token_above_272k_tokens": 9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supported_endpoints": [ "/v1/chat/completions", @@ -68175,9 +69380,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_above_272k_tokens": 9e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -68223,9 +69428,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -68277,6 +69482,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/o3-2025-04-16": { @@ -68290,6 +69500,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/o3-deep-research": { @@ -68312,6 +69527,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/o4-mini-2025-04-16": { @@ -68325,6 +69545,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/text-embedding-3-large": { @@ -68418,6 +69643,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-mini": { @@ -68432,6 +69662,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { @@ -68470,6 +69705,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5-codex": { @@ -68494,6 +69734,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5-nano": { @@ -68506,6 +69751,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5-pro": { @@ -68516,6 +69766,11 @@ "mode": "responses", "output_cost_per_token": 0.000132, "output_cost_per_token_batches": 6.6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.1-codex-max": { @@ -68540,6 +69795,11 @@ "output_cost_per_token": 1.54e-05, "output_cost_per_token_batches": 7.7e-06, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.2-codex": { @@ -68558,6 +69818,11 @@ "mode": "responses", "output_cost_per_token": 0.0001848, "output_cost_per_token_batches": 9.24e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.3-codex": { @@ -68570,6 +69835,11 @@ "mode": "responses", "output_cost_per_token": 1.54e-05, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-mini": { @@ -68585,6 +69855,11 @@ "output_cost_per_token": 4.95e-06, "output_cost_per_token_batches": 2.475e-06, "output_cost_per_token_priority": 9.9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-nano": { @@ -68597,6 +69872,11 @@ "mode": "chat", "output_cost_per_token": 1.375e-06, "output_cost_per_token_batches": 6.875e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-pro": { @@ -68611,6 +69891,11 @@ "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/o1-mini": { @@ -68657,7 +69942,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2026-10-31", + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -71002,6 +72287,7 @@ "xai/grok-code-fast-1": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "deprecation_date": "2026-05-15", "input_cost_per_image_token": 1e-06, "input_cost_per_token": 1e-06, "input_cost_per_token_above_200k_tokens": 2e-06, @@ -71898,7 +73184,7 @@ "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, - "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost": 1e-07, "deprecation_date": "2027-09-29", "input_cost_per_token": 2e-06, "litellm_provider": "azure_ai", @@ -72505,15 +73791,15 @@ "gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, - "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, - "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, @@ -72557,6 +73843,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": { @@ -74474,14 +75769,14 @@ "supports_vision": true }, "wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2.5e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 943718, "mode": "chat", - "output_cost_per_token": 7.5e-07, - "source": "https://openrouter.ai/api/v1/models/nvidia/nemotron-3.5-super-vl/endpoints", + "output_cost_per_token": 1.5e-06, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -76056,6 +77351,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}, @@ -76188,7 +77492,7 @@ "supports_xhigh_reasoning_effort": true }, "baseten/decision-models/baseten-elephant-9b-latest": { - "input_cost_per_token": 0.0, + "input_cost_per_token": 3.5e-08, "litellm_provider": "baseten", "max_input_tokens": 16384, "mode": "decisions", @@ -76197,7 +77501,7 @@ "supports_vision": true }, "baseten/decision-models/baseten-mammoth-27b-latest": { - "input_cost_per_token": 0.0, + "input_cost_per_token": 1e-07, "litellm_provider": "baseten", "max_input_tokens": 16384, "mode": "decisions", @@ -76229,13 +77533,13 @@ "supports_vision": true }, "refuseless/glm-5.3-abliterated": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 1.25e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.25e-06, + "output_cost_per_token": 2e-06, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -76250,13 +77554,13 @@ ] }, "refuseless/glm-5.3-flash-abliterated": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 2.5e-07, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 1.2e-07, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8e-07, + "output_cost_per_token": 5e-07, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -76271,13 +77575,13 @@ ] }, "refuseless/deepseek-v4.1-flash-abliterated": { - "cache_read_input_token_cost": 4.5e-08, - "input_cost_per_token": 2.2e-07, + "cache_read_input_token_cost": 3.5e-08, + "input_cost_per_token": 1.4e-07, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 130000, "mode": "chat", - "output_cost_per_token": 7.5e-07, + "output_cost_per_token": 5e-07, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -76333,6 +77637,15 @@ "max" ] }, + "fireworks_ai/accounts/fireworks/models/rune": { + "cache_read_input_token_cost": 8.4e-09, + "input_cost_per_token": 4.2e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://api.fireworks.ai/v1/serverless/models?format=nested" + }, "bedrock_mantle/openai.gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, @@ -76370,5 +77683,614 @@ "supports_vision": true, "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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", + "input_cost_per_token": 1.05e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 2.1e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-spot": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.05e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-spot": { + "cache_read_input_token_cost": 1.82e-07, + "input_cost_per_token": 9.8e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.08e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-1017-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 1ab03fb3fdc..d73ac19bae1 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/_types.py b/litellm/proxy/_types.py index dd91bb6eef3..5ce6cd93a22 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3461,12 +3461,12 @@ class LiteLLM_SpendLogs_ResponseObject(LiteLLMPydanticObjectBase): class TokenCountRequest(LiteLLMPydanticObjectBase): model: str prompt: str | None = None - messages: list[dict] | None = None + messages: list[dict[Any, Any]] | None = None """ Anthropic token counting endpoint uses /messages """ - contents: list[dict] | None = None + contents: list[dict[Any, Any]] | None = None """ Google /countTokens endpoint expects contents to be a list of dicts with the following structure: """ diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1d2045eb743..1d88e133cf9 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -18,7 +18,6 @@ from litellm.anthropic_interface.exceptions import ( AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.messages.context_management import ( AnthropicContextManagementError, @@ -49,6 +48,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) +from litellm.proxy.common_utils.token_count_log import rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.litellm_pre_call_utils import ( forwarded_anthropic_oauth_token, @@ -394,12 +394,17 @@ async def count_tokens( # Extract required fields model_name: Final = data.get("model") messages: Final = data.get("messages", []) + request_model: Final = model_name if isinstance(model_name, str) else None if not model_name: - raise HTTPException(status_code=400, detail="model parameter is required") + raise rejected_token_count( + HTTPException(status_code=400, detail="model parameter is required"), litellm_call_id, request_model + ) if not messages: - raise HTTPException(status_code=400, detail="messages parameter is required") + raise rejected_token_count( + HTTPException(status_code=400, detail="messages parameter is required"), litellm_call_id, request_model + ) auto_drop_params: Final = should_auto_drop_params_for_agentic_cli( request.headers.get("user-agent", ""), "drop_params" in data, proxy_config @@ -409,19 +414,18 @@ async def count_tokens( token_fields if isinstance(token_fields, ProxyException) else _token_count_request(token_fields) ) if isinstance(token_request, ProxyException): - return _anthropic_error_json_response(token_request, request) + return rejected_token_count( + _anthropic_error_json_response(token_request, request), litellm_call_id, request_model + ) token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", + litellm_call_id=litellm_call_id, client_oauth_token=forwarded_anthropic_oauth_token(request.headers), **_count_tokens_body_options(data, auto_drop_params), ) - except UnsupportedContentTypeError as e: - return _anthropic_error_json_response( - ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400), request - ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cb490b55352..99edfb940d7 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -3257,8 +3257,9 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e if valid_token.budget_throttle_pct is not None: valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct - if _end_user_object is not None: + if _end_user_object is not None or litellm.max_end_user_budget_id is not None: valid_token_dict.update(end_user_params) + if _end_user_object is not None: valid_token_dict["end_user_object_permission"] = _end_user_object.object_permission # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions 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/agents.py b/litellm/proxy/client/cli/commands/agents.py index 52e36964bc1..4d09e4d1e0f 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -34,7 +34,7 @@ from .claude_settings import ( load_json_or_empty, ) from .cmd_quoting import quote_for_cmd -from .config import claude_subscription_default +from .config import subscription_default from .pi import ( LITELLM_PROXY_API_KEY_ENV, PI_PROVIDER_NAME, @@ -83,9 +83,11 @@ _INSTALL_DOCS: Final[dict[str, str]] = { } _HIDDEN_AGENTS: Final = frozenset({"pi"}) -_SUBSCRIPTION_AGENTS: Final = frozenset({"claude"}) +_SUBSCRIPTION_LOGINS: Final = MappingProxyType({"claude": "claude.ai", "codex": "ChatGPT"}) +_SUBSCRIPTION_AGENTS: Final = frozenset(_SUBSCRIPTION_LOGINS) CODEX_PROXY_PROVIDER: Final = "litellm" +CODEX_SUBSCRIPTION_PROVIDER: Final = "lite_subscription" CODEX_HOME_ENV: Final = "CODEX_HOME" CODEX_MODEL_CATALOG_FILENAME: Final = "litellm-models.json" _CODEX_PREFLIGHT_TIMEOUT_SECONDS: Final = 10.0 @@ -143,6 +145,10 @@ def build_agent_env( travels in the x-litellm-api-key header through ANTHROPIC_CUSTOM_HEADERS, and no ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY is left in the environment, since either one makes Claude Code stop using the subscription (and its connectors). + With subscription=True the OpenAI profile keeps Codex on its ChatGPT login: the key goes to + LITELLM_PROXY_API_KEY, which Codex sends in the x-litellm-api-key header, and OPENAI_API_KEY and + OPENAI_BASE_URL are left exactly as the caller had them, so the key is never offered to Codex as an + OpenAI credential and a child process never pairs the caller's own OpenAI key with the proxy URL. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -161,21 +167,24 @@ def build_agent_env( if ENABLE_GATEWAY_MODEL_DISCOVERY_ENV not in env: env[ENABLE_GATEWAY_MODEL_DISCOVERY_ENV] = ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE if PROFILE_OPENAI in profiles: - env[OPENAI_BASE_URL_ENV] = root + "/v1" - env[OPENAI_API_KEY_ENV] = api_key + if subscription: + env[LITELLM_PROXY_API_KEY_ENV] = api_key + else: + env[OPENAI_BASE_URL_ENV] = root + "/v1" + env[OPENAI_API_KEY_ENV] = api_key if PROFILE_LITELLM in profiles: env[LITELLM_PROXY_API_KEY_ENV] = api_key return env -def codex_proxy_provider(base_url: str) -> Mapping[str, str | bool]: +def codex_proxy_provider(base_url: str, *, subscription: bool = False) -> Mapping[str, str | bool]: return MappingProxyType( { "name": "LiteLLM proxy", "base_url": base_url.rstrip("/") + "/v1", "wire_api": "responses", "supports_websockets": False, - "requires_openai_auth": False, + "requires_openai_auth": subscription, } ) @@ -304,7 +313,7 @@ def subscription_launch_args( return (_SETTINGS_FLAG, str(overlay)), rest -def _codex_proxy_args(base_url: str) -> list[str]: +def _codex_proxy_args(base_url: str, subscription: bool) -> list[str]: """Codex `-c` overrides that point it at the proxy. Codex ignores OPENAI_BASE_URL (it always dials api.openai.com), so the env @@ -312,24 +321,35 @@ def _codex_proxy_args(base_url: str) -> list[str]: inline; supports_websockets=false forces the HTTP/SSE Responses transport because the proxy does not speak the Responses WebSocket protocol. The key is read from OPENAI_API_KEY, which build_agent_env already exports. + With subscription=True the provider is `lite_subscription` with requires_openai_auth, so Codex + signs in with its own ChatGPT login, and the key travels in the x-litellm-api-key header from + LITELLM_PROXY_API_KEY. It sets no env_key, which would win over the login. A `-c` override cannot + remove keys, so the id is one no config.toml defines: `lite configure codex` writes `litellm`, and + the cookbook tells users to write `litellm_subscription`, whose env_key or headers would otherwise apply. """ - provider: Final = f"model_providers.{CODEX_PROXY_PROVIDER}" + provider_id: Final = CODEX_SUBSCRIPTION_PROVIDER if subscription else CODEX_PROXY_PROVIDER + provider: Final = f"model_providers.{provider_id}" + credential: Final = ( + f"{provider}.env_http_headers={{{json.dumps(LITELLM_KEY_HEADER)}={json.dumps(LITELLM_PROXY_API_KEY_ENV)}}}" + if subscription + else f'{provider}.env_key="{OPENAI_API_KEY_ENV}"' + ) return [ "-c", - f'model_provider="{CODEX_PROXY_PROVIDER}"', + f'model_provider="{provider_id}"', *( argument - for key, value in codex_proxy_provider(base_url).items() + for key, value in codex_proxy_provider(base_url, subscription=subscription).items() for argument in ("-c", f"{provider}.{key}={json.dumps(value)}") ), "-c", - f'{provider}.env_key="{OPENAI_API_KEY_ENV}"', + credential, "-c", f"{provider}.http_headers={{}}", ] -_PROXY_ARGS: Final[dict[str, Callable[[str], list[str]]]] = { +_PROXY_ARGS: Final[dict[str, Callable[[str, bool], list[str]]]] = { "codex": _codex_proxy_args, } @@ -410,14 +430,14 @@ _Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]] _PREPARERS: Final[Mapping[str, _Preparer]] = MappingProxyType({"pi": prepare_pi, "codex": prepare_codex}) -def agent_launch_args(command: str, base_url: str) -> list[str]: +def agent_launch_args(command: str, base_url: str, *, subscription: bool = False) -> list[str]: """Extra CLI args an agent needs to actually honor the proxy. Claude Code and OpenCode respect the exported env vars, so they get nothing here; Codex needs its provider pointed via config overrides. """ builder: Final = _PROXY_ARGS.get(os.path.basename(command)) - return builder(base_url) if builder else [] + return builder(base_url, subscription) if builder else [] class ListedModel(LiteLLMBaseModel): @@ -889,11 +909,13 @@ def run_agent( (verify_subscription if subscription else verify)(base_url, api_key) env_before_sync: Final = base_env if base_env is not None else os.environ - synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify) + synced: Final = ( + _NO_EXTRA_ENV if subscription else sync_models(binary, env_before_sync, base_url, api_key, skip_verify) + ) if isinstance(synced, ModelSyncSkipped): warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}") - prepare: Final = preparers.get(os.path.basename(command[0])) + prepare: Final = None if subscription else preparers.get(os.path.basename(command[0])) prepared_args: Final = tuple(prepare(base_url, api_key, env_before_sync)) if prepare is not None else () env: Final = MappingProxyType( @@ -903,7 +925,11 @@ def run_agent( } ) synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else () - extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args) + extra_args: Final = ( + *agent_launch_args(command[0], base_url, subscription=subscription), + *synced_args, + *prepared_args, + ) settings_args, forwarded = ( subscription_launch_args( command, @@ -914,7 +940,7 @@ def run_agent( ), write_overlay, ) - if subscription and os.path.basename(command[0]) in _SUBSCRIPTION_AGENTS + if subscription and os.path.basename(command[0]) == "claude" else ((), tuple(command[1:])) ) if reattach_terminal is not None: @@ -948,10 +974,19 @@ def resolve_api_key(ctx: click.Context) -> str: _SKIP_VERIFY_HELP: Final = "Skip the pre-launch key check against the proxy." -_SUBSCRIPTION_HELP: Final = ( - "Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the " - "x-litellm-api-key header instead of replacing Claude Code's own credential. " - "Make it the default with `lite config set claude_subscription true`." +_SUBSCRIPTION_HELP: Final = MappingProxyType( + { + "claude": ( + "Keep using your claude.ai login (and its connectors). The LiteLLM key is sent in the " + "x-litellm-api-key header instead of replacing Claude Code's own credential. " + "Make it the default with `lite config set claude_subscription true`." + ), + "codex": ( + "Keep using your ChatGPT login. The LiteLLM key is sent in the " + "x-litellm-api-key header instead of replacing Codex's own credential. " + "Make it the default with `lite config set codex_subscription true`." + ), + } ) @@ -964,10 +999,8 @@ def _launch( api_key: Final = resolve_api_key(ctx) display_name, _profiles = agent_profile(binary) - use_subscription: Final = ( - subscription if subscription is not None else binary in _SUBSCRIPTION_AGENTS and claude_subscription_default() - ) - login_note: Final = " (keeping your claude.ai login)" if use_subscription else "" + use_subscription: Final = subscription if subscription is not None else subscription_default(binary) + login_note: Final = f" (keeping your {_SUBSCRIPTION_LOGINS[binary]} login)" if use_subscription else "" click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}{login_note}") try: @@ -989,7 +1022,7 @@ _Decorated: Final = TypeVar("_Decorated", bound=Callable[..., object] | click.Co def _subscription_option(binary: str) -> Callable[[_Decorated], _Decorated]: if binary not in _SUBSCRIPTION_AGENTS: return lambda func: func - return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP) + return click.option("--subscription/--no-subscription", default=None, help=_SUBSCRIPTION_HELP[binary]) def _make_agent_command(binary: str, display_name: str) -> click.Command: diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index b5302d1f561..e38f54adc4f 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -13,7 +13,6 @@ from rich.console import Console from rich.table import Table from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never -from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, SYSTEM_KEYRING, @@ -143,6 +142,9 @@ _TeamMapping: Final = TypeVar("_TeamMapping", bound=Mapping[str, object]) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_PROXY_KEY_CHECK_TIMEOUT_SECONDS: Final = 3 +_PROXY_REJECTED_KEY_STATUSES: Final = frozenset({401, 403}) + KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" KEYRING_ENABLE_HINT: Final = "keyring --enable (or unset PYTHON_KEYRING_BACKEND)" @@ -1033,10 +1035,12 @@ def print_token(ctx: click.Context): Designed to be used as Claude Code's `apiKeyHelper` (https://docs.claude.com/en/docs/claude-code/settings): stdout must - contain only the token, so all diagnostics go to stderr. The token - expires after `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h); a - `lite login --pkce` token renews itself here first, and once a token - has expired for good, run the same `lite login` command again. + contain only the token, so all diagnostics go to stderr. A token younger + than `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h) is printed without + asking the proxy. An older one is checked against the proxy and refused + only if the proxy rejects it. A `lite login --pkce` + token renews itself here first, and once a token has expired for good, + run the same `lite login` command again. """ vault: Final = context_secret_vault(ctx) token_data: Final = load_token(vault=vault) @@ -1055,7 +1059,7 @@ def print_token(ctx: click.Context): sys.exit(1) renews: Final = "refresh_token" in token_data - if not is_cli_token_fresh(token_data) and not renews: + if not renews and token_expired(token_data): click.echo("Token expired. Run 'lite login' again.", err=True) sys.exit(1) @@ -1111,8 +1115,32 @@ def whoami(ctx: click.Context): expires_at: Final = token_data.get("expires_at") if isinstance(expires_at, (int, float)): click.echo(_key_expiry_line(expires_at, renews="refresh_token" in token_data)) - elif age_hours > CLI_JWT_EXPIRATION_HOURS: - click.echo(f"Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") + elif token_expired(token_data): + click.echo("Token expired. Run 'lite login' again.") + + +def token_expired(token_data: Mapping[str, object]) -> bool: + if is_cli_token_fresh(token_data): + return False + if isinstance(token_data.get("expires_at"), (int, float)): + return True + return _proxy_rejects_key(token_data) + + +def _proxy_rejects_key(token_data: Mapping[str, object]) -> bool: + key: Final = token_data.get("key") + base_url: Final = token_data.get("base_url") + if not isinstance(key, str) or not isinstance(base_url, str) or not base_url: + return True + try: + response: Final = requests.get( + f"{base_url}/v1/models", + headers={"Authorization": f"Bearer {key}"}, + timeout=_PROXY_KEY_CHECK_TIMEOUT_SECONDS, + ) + except requests.RequestException: + return False + return response.status_code in _PROXY_REJECTED_KEY_STATUSES def _key_expiry_line(expires_at: float, renews: bool) -> str: diff --git a/litellm/proxy/client/cli/commands/config.py b/litellm/proxy/client/cli/commands/config.py index 539d88dce4a..1689a08c3de 100644 --- a/litellm/proxy/client/cli/commands/config.py +++ b/litellm/proxy/client/cli/commands/config.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.private_json import ensure_private_dir, write_pr HIDDEN_COMMANDS_KEY: Final = "hidden_commands" CLAUDE_SUBSCRIPTION_KEY: Final = "claude_subscription" +CODEX_SUBSCRIPTION_KEY: Final = "codex_subscription" PKCE_KEY: Final = "pkce" _config_adapter: Final[TypeAdapter[Mapping[str, str]]] = TypeAdapter(Mapping[str, str]) @@ -114,9 +115,15 @@ def _normalize_boolean(key: str) -> Callable[[str], str]: return normalize -def claude_subscription_default() -> bool: - """Whether `lite claude` keeps the claude.ai login unless told otherwise.""" - return get_config_value(CLAUDE_SUBSCRIPTION_KEY) == "true" +_SUBSCRIPTION_KEYS: Final[Mapping[str, str]] = MappingProxyType( + {"claude": CLAUDE_SUBSCRIPTION_KEY, "codex": CODEX_SUBSCRIPTION_KEY} +) + + +def subscription_default(agent: str) -> bool: + """Whether `lite ` keeps the agent's own subscription login unless told otherwise.""" + key: Final = _SUBSCRIPTION_KEYS.get(agent) + return key is not None and get_config_value(key) == "true" def pkce_default() -> bool: @@ -129,6 +136,7 @@ _NORMALIZERS: Final[Mapping[str, Callable[[str], str]]] = MappingProxyType( "base_url": normalize_base_url, HIDDEN_COMMANDS_KEY: _normalize_hidden_commands, CLAUDE_SUBSCRIPTION_KEY: _normalize_boolean(CLAUDE_SUBSCRIPTION_KEY), + CODEX_SUBSCRIPTION_KEY: _normalize_boolean(CODEX_SUBSCRIPTION_KEY), PKCE_KEY: _normalize_boolean(PKCE_KEY), } ) 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_log.py b/litellm/proxy/common_utils/token_count_log.py new file mode 100644 index 00000000000..04c07a29117 --- /dev/null +++ b/litellm/proxy/common_utils/token_count_log.py @@ -0,0 +1,133 @@ +import json +import re +from asyncio import CancelledError +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, TypeVar + +from fastapi import status + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.openai_error_payload import error_status_code +from litellm.types.token_count import TokenCountSource +from litellm.types.utils import TokenCountResponse + +_BARE_LOGFMT_VALUE: Final = re.compile(r"[\w.:/@+-]+", re.ASCII) +_JSON_ESCAPED_BRACES: Final = str.maketrans({"{": "\\u007b", "}": "\\u007d"}) + + +def _logfmt_value(value: str | None) -> str: + if value is None: + return "None" + if _BARE_LOGFMT_VALUE.fullmatch(value): + return value + return json.dumps(value).translate(_JSON_ESCAPED_BRACES) + + +_Counted: Final = TypeVar("_Counted") +_Rejection: Final = TypeVar("_Rejection") +_ProviderCount: Final = TypeVar("_ProviderCount", bound=Mapping[str, object]) + + +@dataclass(frozen=True, slots=True) +class _CountFields: + model_used: str | None + input_tokens: int | None + count_source: TokenCountSource | None + + +_FAILED_COUNT: Final = _CountFields(model_used=None, input_tokens=None, count_source=None) + + +@dataclass(frozen=True, slots=True) +class _CountRequest: + litellm_call_id: str + request_model: str | None + model_id: str | None + + +def deployment_model_id(deployment: Mapping[str, object] | None) -> str | None: + match deployment: + case {"model_info": {"id": str() as model_id}}: + return model_id + case _: + return None + + +def _log_token_count(request: _CountRequest, fields: _CountFields, status_code: int) -> None: + verbose_proxy_logger.info( + "token count - litellm_call_id=%s request_model=%s model_id=%s model_used=%s input_tokens=%s " + "count_source=%s status_code=%s", + _logfmt_value(request.litellm_call_id), + _logfmt_value(request.request_model), + _logfmt_value(request.model_id), + _logfmt_value(fields.model_used), + fields.input_tokens, + fields.count_source, + status_code, + extra=MappingProxyType( + { + "litellm_call_id": request.litellm_call_id, + "request_model": request.request_model, + "model_id": request.model_id, + "model_used": fields.model_used, + "input_tokens": fields.input_tokens, + "count_source": fields.count_source, + "status_code": status_code, + } + ), + ) + + +def rejected_token_count(rejection: _Rejection, litellm_call_id: str, request_model: str | None) -> _Rejection: + _log_token_count( + _CountRequest(litellm_call_id, request_model, model_id=None), + _FAILED_COUNT, + error_status_code(rejection, status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + return rejection + + +async def _logged( + count: Awaitable[_Counted], request: _CountRequest, fields_of: Callable[[_Counted], _CountFields] +) -> _Counted: + try: + counted: Final = await count + except CancelledError: + raise + except BaseException as e: + _log_token_count(request, _FAILED_COUNT, error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR)) + raise + _log_token_count(request, fields_of(counted), status.HTTP_200_OK) + return counted + + +def _token_count_response_fields(counted: TokenCountResponse) -> _CountFields: + return _CountFields( + model_used=counted.model_used, input_tokens=counted.total_tokens, count_source=counted.count_source + ) + + +async def logged_token_count( + count: Awaitable[TokenCountResponse], litellm_call_id: str, request_model: str, model_id: str | None +) -> TokenCountResponse: + return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), _token_count_response_fields) + + +async def logged_provider_token_count( + count: Awaitable[_ProviderCount], + litellm_call_id: str, + request_model: str, + model_id: str | None, + model_used: Callable[[], str], +) -> _ProviderCount: + def provider_fields(counted: _ProviderCount) -> _CountFields: + input_tokens: Final = counted.get("input_tokens") + return _CountFields( + model_used=model_used(), + input_tokens=input_tokens if isinstance(input_tokens, int) and not isinstance(input_tokens, bool) else None, + count_source="provider", + ) + + return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), provider_fields) 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/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 3d1a9ab4cfe..54976857584 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -21,6 +21,8 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, ) +from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER +from litellm.proxy.common_utils.token_count_log import rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse @@ -286,25 +288,35 @@ async def google_count_tokens(request: Request, model_name: str): from litellm.proxy.common_utils.http_parsing_utils import read_request_body from litellm.proxy.proxy_server import token_counter as internal_token_counter + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) data: Final = await read_request_body(request=request) try: token_request: Final = _google_count_tokens_request(model_name, data) except ValidationError: - return _google_error_response( - request, - ProxyException( - message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + return rejected_token_count( + _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", + type="invalid_request_error", + param=None, + code=400, + ), ), + litellm_call_id, + model_name, ) try: - token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + token_response: Final = await internal_token_counter( + request=token_request, call_endpoint=True, litellm_call_id=litellm_call_id + ) except ProxyException as error: return _google_error_response(request, error) except HTTPException as error: return _google_error_response(request, proxy_exception_from_http_exception(error, {})) except Exception as error: - log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + log_llm_api_exception(error, litellm_call_id) return _google_error_response( request, ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 567b8452ae2..e44d9304ac1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -45,11 +45,14 @@ from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.cost_calculator import get_usage_object from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, + split_media_for_tpm, ) from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import estimate_image_tokens, offload_token_count from litellm.llms.azure.ptu_capacity import normalized_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( @@ -98,6 +101,7 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + from litellm.router import Router from litellm.types.agents import AgentResponse from litellm.types.caching import RedisPipelineIncrementOperation @@ -135,6 +139,18 @@ class RateLimitedModel: return requested_limit if requested_limit is not None else limits.get(self.group) +def _resolve_media_models(router: "Router | None", model: str, team_id: str | None) -> tuple[str, ...]: + if router is None: + return (model,) + return tuple(dict.fromkeys(router.resolved_litellm_models(model, team_id, use_base_model=True))) or (model,) + + +def _resolve_media_models_via_proxy_router(model: str, team_id: str | None) -> tuple[str, ...]: + from litellm.proxy.proxy_server import llm_router + + return _resolve_media_models(llm_router, model, team_id) + + def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None: from litellm.proxy.proxy_server import llm_router @@ -862,6 +878,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): time_provider: Callable[[], datetime] | None = None, tag_rate_limit_resolver: TagRateLimitResolver = resolve_tag_rate_limits_from_db, model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router, + media_model_resolver: Callable[[str, str | None], tuple[str, ...]] = _resolve_media_models_via_proxy_router, ptu_team_ceiling_resolver: Callable[ [str, str], PTUTeamCeiling | None ] = _resolve_ptu_team_ceiling_via_proxy_router, @@ -872,6 +889,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._time_provider = time_provider or datetime.now self._tag_rate_limit_resolver = tag_rate_limit_resolver self._model_group_resolver = model_group_resolver + self._media_model_resolver = media_model_resolver self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver self._fail_closed_resolver = fail_closed_resolver self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver @@ -1133,11 +1151,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _estimate_tokens_for_request( self, - data: dict, + data: object, model: str | None = None, min_configured_tpm_limit: int | None = None, call_type: str | None = None, configured_output_tokens: int | None = None, + team_id: str | None = None, ) -> int: """ Estimate total tokens this request will consume so we can reserve them @@ -1164,6 +1183,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=team_id, + media_model_resolver=self._media_model_resolver, ) total_estimated: Final = estimated_input_tokens + max_tokens_estimate @@ -1184,6 +1205,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): call_type: str | None, configured_output_tokens: int | None, raw_estimate: int, + team_id: str | None = None, ) -> int: """The team PTU ceiling counts Azure normalized tokens, so its reservation weighs the output budget the way the ceiling does instead of the raw sum the other scopes reserve.""" @@ -1194,6 +1216,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=team_id, + media_model_resolver=self._media_model_resolver, ) normalized: Final = normalized_tokens( ceiling, prompt_tokens=estimated_input_tokens, completion_tokens=max_tokens_estimate @@ -1206,6 +1230,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit: int | None = None, call_type: str | None = None, configured_output_tokens: int | None = None, + include_media: bool = True, + team_id: str | None = None, + media_model_resolver: Callable[[str, str | None], tuple[str, ...]] | None = None, ) -> tuple[int, int]: """ Estimate input tokens and output (max_tokens) budget separately, so @@ -1229,10 +1256,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): floor entirely, so the reservation reflects what this tenant's model actually emits rather than one constant shared by every tenant. """ - if not isinstance(data, dict): + request_data: Final = as_str_mapping(data) + if not isinstance(data, dict) or request_data is None: return 0, 0 + request_model: Final = request_data.get("model") translated_data: Final = self._translate_google_genai_native_request(data, call_type) - estimable_data: Final = translated_data if translated_data is not None else data + estimable_data: Final = translated_data if translated_data is not None else request_data selected_fields: Final[tuple[object | None, object | None, object | None]] = ( (None, None, estimable_data.get("input")) if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES @@ -1249,10 +1278,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) messages, prompt, raw_input = selected_fields - input_text: Final = ( + original_input: Final = ( billable_input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input ) + media_source: Final = messages if messages is not None else original_input + projected, media = ( + split_media_for_tpm(media_source) + if include_media and (isinstance(messages, list) or call_type in RESPONSES_API_CALL_TYPES) + else (original_input, ()) + ) + input_text: Final = projected if call_type in RESPONSES_API_CALL_TYPES else original_input + fallback_model: Final = request_model if isinstance(request_model, str) else None + resolved_models: Final = ( + media_model_resolver(fallback_model or "", team_id) + if media and media_model_resolver is not None + else (fallback_model or "",) + ) + media_tokens: Final = sum( + self._estimate_media_block_tokens(block, resolved_models, fallback_model) for block in media + ) + total_chars: Final = ( len(get_str_from_messages(messages)) if isinstance(messages, list) and messages @@ -1267,7 +1313,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else 0 ) - estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 + estimated_input_tokens: Final = ( + max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 + ) + media_tokens explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type) generates_no_output: Final = ( @@ -1283,7 +1331,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): max_tokens_estimate: Final = ( 0 if generates_no_output - or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None) + or (explicit_max_tokens is None and estimated_input_tokens == 0 and configured_output_tokens is None) else explicit_max_tokens if explicit_max_tokens is not None else configured_output_tokens @@ -1293,6 +1341,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return estimated_input_tokens, max_tokens_estimate * self.get_output_candidate_count(data, call_type) + @staticmethod + def _estimate_media_block_tokens( + block: Mapping[str, object], models: tuple[str, ...], fallback_model: str | None + ) -> int: + if block.get("type") == "input_audio": + audio: Final = _callback_kwargs_mapping(block.get("input_audio")) + data: Final = audio.get("data") + return estimate_audio_tokens(data if isinstance(data, str) else None, models, fallback_model=fallback_model) + image: Final = _callback_kwargs_mapping(block.get("image_url")) + url: Final = image.get("url") + detail: Final = image.get("detail") + return estimate_image_tokens( + models, + url if isinstance(url, str) else "", + detail if isinstance(detail, str) else "auto", + fallback_model=fallback_model, + ) + def _is_redis_cluster(self) -> bool: """ Check if the dual cache is using Redis cluster. @@ -3958,7 +4024,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): except Exception: # noqa: BLE001 # tokenizer failures degrade to the cheap estimate if call_type in RERANK_API_CALL_TYPES and isinstance(selected_text, str): return max(0, len(selected_text) // DEFAULT_CHARS_PER_TOKEN) - estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data, call_type=call_type) + estimated_input_tokens, _ = self._estimate_input_and_output_tokens( + data=data, call_type=call_type, include_media=False + ) return estimated_input_tokens + audio_token_estimate async def _reserve_project_io_tokens_or_raise( @@ -4003,6 +4071,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data=data, min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, + include_media=False, ) raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type @@ -4360,12 +4429,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Post-call reconciliation refunds the over-reservation # delta when actual usage comes in below the floor. estimated_tokens: Final = max( - self._estimate_tokens_for_request( + await offload_token_count(self._estimate_tokens_for_request)( data=data, model=requested_model, min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=user_api_key_dict.team_id, ), 1, ) @@ -4383,13 +4453,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) capped_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data) - ptu_estimated_tokens: Final = self._estimate_ptu_tokens_for_request( + ptu_estimated_tokens: Final = await offload_token_count(self._estimate_ptu_tokens_for_request)( ceiling=stash.ptu_ceiling, data=capped_request, min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, raw_estimate=estimated_tokens, + team_id=user_api_key_dict.team_id, ) if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit: if not reuses_request_slot: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 314f14b6279..24a31723faa 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -32,6 +32,7 @@ from litellm.constants import ( BILLED_RETRY_ATTEMPT_METADATA_KEY, CALLER_SESSION_ID_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, + COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, @@ -210,13 +211,46 @@ def _caller_session_beside_trace_header( return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name) -def _drop_client_caller_session(request_body: Mapping[str, object]) -> None: +def _child_tier_session_metadata( + normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str +) -> Mapping[str, str | None]: + parent: Final = normalized_headers.get("x-parent-session-id") + if not parent or not _SESSION_ID_VALUE_RE.fullmatch(parent): + return {} + own_headers: Final = { + key: value + for key, value in normalized_headers.items() + if key not in ("x-parent-session-id", "x-litellm-trace-id") + } + active: Final = data.get(metadata_variable_name) + inactive: Final = data.get("litellm_metadata") + inactive_session: Final = ( + _caller_body_session_id(data, "litellm_metadata") + if metadata_variable_name == "metadata" + and not (isinstance(active, Mapping) and "session_id" in active) + and isinstance(inactive, Mapping) + else None + ) + own_session: Final = ( + get_chain_id_from_headers(own_headers) + or _get_anthropic_session_id_from_metadata(data.get("metadata")) + or _caller_body_session_id(data, metadata_variable_name) + or inactive_session + or _session_id_from_baggage(normalized_headers.get("baggage", "")) + ) + return {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: own_session if own_session != parent else None} + + +def _drop_client_session_metadata(request_body: Mapping[str, object]) -> None: for metadata_key in ("metadata", "litellm_metadata"): if isinstance(metadata := request_body.get(metadata_key), dict): - metadata.pop(CALLER_SESSION_ID_METADATA_KEY, None) + for key in (CALLER_SESSION_ID_METADATA_KEY, COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY): + cast(dict[str, object], metadata).pop(key, None) # cast-ok: request metadata is a str-keyed JSON object -_PROXY_SET_METADATA_KEYS: Final = frozenset({CALLER_SESSION_ID_METADATA_KEY, BAGGAGE_MAPPED_TAGS_METADATA_KEY}) +_PROXY_SET_METADATA_KEYS: Final = frozenset( + {CALLER_SESSION_ID_METADATA_KEY, BAGGAGE_MAPPED_TAGS_METADATA_KEY, COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY} +) _METADATA_TAGS_ADAPTER: Final = TypeAdapter(tuple[object, ...]) @@ -1855,7 +1889,7 @@ class LiteLLMProxyRequestSetup: metadata_from_headers: Final = LitellmMetadataFromRequestHeaders() request_body: Final = cast(Mapping[str, object], data) # cast-ok: request body is a str-keyed JSON object - _drop_client_caller_session(request_body) + _drop_client_session_metadata(request_body) spend_logs_metadata: Final = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers) if spend_logs_metadata is not None: metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata @@ -1870,6 +1904,9 @@ class LiteLLMProxyRequestSetup: agent_id_from_header: Final = headers.get("x-litellm-agent-id") # Explicit litellm headers take precedence; fall back to any x-*-session-id header. chain_id: Final = get_chain_id_from_headers(dict(headers)) + child_tier_metadata: Final = _child_tier_session_metadata( + normalized_headers, request_body, _metadata_variable_name + ) if agent_id_from_header: metadata_from_headers["agent_id"] = agent_id_from_header @@ -1939,9 +1976,15 @@ class LiteLLMProxyRequestSetup: if isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name].update( - {**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: caller_session_beside_trace} - if caller_session_beside_trace - else metadata_from_headers + { + **metadata_from_headers, + **child_tier_metadata, + **( + {CALLER_SESSION_ID_METADATA_KEY: caller_session_beside_trace} + if caller_session_beside_trace + else {} + ), + } ) return data diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 9c46e711cf6..7a843aecc0c 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -49,6 +49,8 @@ from litellm.proxy.management_endpoints.common_daily_activity import daily_activ from litellm.proxy.management_helpers.auto_router_permissions import ( authorize_member_auto_router_dependencies, authorize_member_auto_router_team, + auto_router_environment_reference_error, + raise_auto_router_environment_reference_error, validate_member_auto_router_config, ) from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository @@ -402,6 +404,10 @@ async def validate_complexity_router_config( """ member_team: Final = await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) + raise_auto_router_environment_reference_error( + auto_router_environment_reference_error(data.complexity_router_config, user_api_key_dict) + ) + from litellm.router_utils.auto_router_model_naming import ( validate_complexity_router_config_write, ) @@ -560,6 +566,12 @@ async def preview_auto_router_routing( detail={"error": CommonProxyErrors.no_llm_router.value}, ) resolved: Final = await _resolve_saved_routing_test(data, user_api_key_dict, llm_router) + if data.saved_model_id is None: + raise_auto_router_environment_reference_error( + auto_router_environment_reference_error( + data.complexity_router_config.model_dump(exclude_none=True), user_api_key_dict + ) + ) actor: Final = ( await _authorize_member_dry_run_config( config=resolved.complexity_router_config.model_dump(exclude_none=True), diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index e625ea6f9f2..45006b14df6 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -12,6 +12,7 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime +import itertools import json from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress @@ -42,6 +43,7 @@ from litellm.litellm_core_utils.ptu_pricing import ( ptu_config_error, ) from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload +from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import ( BlockModelRequest, CommonProxyErrors, @@ -100,8 +102,22 @@ from litellm.proxy.management_helpers.auto_router_permissions import ( authorize_member_auto_router_dependencies, authorize_member_auto_router_team, authorize_member_auto_router_write, + auto_router_environment_reference_error, + raise_auto_router_environment_reference_error, ) from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model +from litellm.proxy.management_helpers.model_quota_policy import ( + ENFORCE_PROJECT_MODEL_QUOTA_SETTING, + ModelQuotaPolicy, + QuotaLimit, + deployment_route, + exempts_any_route, + limits_added_by_team_write, + parse_model_quota_setting, + required_limits, + team_write_serves, + unmet_limits, +) from litellm.proxy.spend_tracking.ptu_feature_flag import ( PTU_COST_ATTRIBUTION_ENV_VAR, is_ptu_cost_attribution_enabled, @@ -111,6 +127,7 @@ from litellm.repositories.chunked_in import find_many_in from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ModelTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.router import Router @@ -152,6 +169,7 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, + LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams, @@ -762,26 +780,93 @@ async def _auto_router_capability_slot( ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" -def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLMParams, enforced: bool) -> None: - """Require both rpm and tpm (each a positive value) when the operator opts in via config.yaml. +_ROUTING_KEYS: Final = ("model", "custom_llm_provider", "api_base") +_RATE_LIMIT_KEYS: Final = ("rpm", "tpm") +_NOTHING_SENT: Final[Mapping[str, object]] = MappingProxyType({}) - Off by default, so deployments keep adding models without limits. When - ``enforce_rpm_tpm_on_model_add: true`` is set under general_settings, a model added - without both rpm and tpm set to a positive value is rejected rather than stored - unbounded (or effectively excluded from routing by a zero/negative limit). - """ - if not enforced: + +def _merged_litellm_params( + incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None, keys: tuple[str, ...] +) -> Mapping[str, object]: + sent: Final = _NOTHING_SENT if incoming is None else _dump_without_none(incoming) + stored: Final = _NOTHING_SENT if existing is None else _decrypted_litellm_params(existing) + return MappingProxyType({key: sent.get(key, stored.get(key)) for key in keys}) + + +def _written_routing_params( + incoming: GenericLiteLLMParams | None, existing: GenericLiteLLMParams | None +) -> Mapping[str, object]: + return _merged_litellm_params(incoming, existing, _ROUTING_KEYS) + + +def _quota_policy(general_settings: Mapping[str, object], setting: str) -> ModelQuotaPolicy | None: + try: + return parse_model_quota_setting(general_settings.get(setting)) + except ValidationError as e: + raise ProxyException( + message=f"general_settings.{setting} is invalid: {e}", + type=ProxyErrorTypes.internal_server_error.value, + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + param=None, + ) from e + + +def _raise_if_rate_limits_required_but_missing( + *, model_name: str, litellm_params: LiteLLM_Params, team_model: bool, policy: ModelQuotaPolicy | None +) -> None: + """Off by default, so deployments keep adding models without limits. Once the operator opts in under + general_settings, the policy decides per model name and provider which of rpm and tpm a new model needs, + and a model missing one, or carrying a zero or negative one, is rejected rather than stored unbounded + (or effectively excluded from routing).""" + if policy is None: return - missing: Final = tuple( - field - for field, value in (("rpm", litellm_params.rpm), ("tpm", litellm_params.tpm)) - if value is None or value <= 0 + _raise_on_missing_rate_limits( + model_name, + unmet_limits( + required_limits( + policy, model_name, deployment_route(_written_routing_params(litellm_params, None), team_model) + ), + MappingProxyType({"rpm": litellm_params.rpm, "tpm": litellm_params.tpm}), + ), ) + + +def _raise_if_model_update_adds_missing_rate_limits( + *, + existing: Deployment, + written_name: str, + written_team_model: bool, + incoming_params: GenericLiteLLMParams | None, + remaining_params: GenericLiteLLMParams, + policy: ModelQuotaPolicy | None, +) -> None: + if policy is None: + return + before: Final = required_limits( + policy, + existing.model_info.team_public_model_name or existing.model_name, + deployment_route( + _written_routing_params(None, existing.litellm_params), existing.model_info.team_id is not None + ), + ) + after: Final = required_limits( + policy, + written_name, + deployment_route(_written_routing_params(incoming_params, remaining_params), written_team_model), + ) + written_limits: Final = _merged_litellm_params(incoming_params, remaining_params, _RATE_LIMIT_KEYS) + checked: Final = after - before + _raise_on_missing_rate_limits( + written_name, unmet_limits(checked, MappingProxyType({limit: written_limits.get(limit) for limit in checked})) + ) + + +def _raise_on_missing_rate_limits(model_name: str, missing: tuple[QuotaLimit, ...]) -> None: if not missing: return raise ProxyException( message=( - f"{' and '.join(missing)} must be set to a positive value when " + f"{' and '.join(missing)} must be set to a positive value for '{model_name}' when " f"'{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING}' is enabled in general_settings" ), type=ProxyErrorTypes.validation_error.value, @@ -790,6 +875,113 @@ def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLM ) +def _project_limit(metadata: object, limits_key: str, model: str) -> object: + limits: Final = metadata.get(limits_key) if is_str_keyed_mapping(metadata) else None + return limits.get(model) if is_str_keyed_mapping(limits) else None + + +def _projects_short_on( + projects: Sequence["prisma_models.LiteLLM_ProjectTable"], model: str, added: frozenset[QuotaLimit] +) -> tuple[str, ...]: + return tuple( + project.project_alias or project.project_id + for project in projects + if model in project.models + and unmet_limits( + added, + MappingProxyType( + { + "rpm": _project_limit(project.metadata, "model_rpm_limit", model), + "tpm": _project_limit(project.metadata, "model_tpm_limit", model), + } + ), + ) + ) + + +async def _raise_if_team_write_leaves_projects_short( + *, + prisma_client: PrismaClient, + general_settings: Mapping[str, object], + llm_router: Router | None, + team_id: str | None, + public_name: str, + replaced_id: str | None, + written_params: Mapping[str, object] | None, +) -> None: + if team_id is None: + return + policy: Final = _quota_policy(general_settings, ENFORCE_PROJECT_MODEL_QUOTA_SETTING) + if policy is None or not exempts_any_route(policy): + return + projects: Final = await ProjectRepository(prisma_client).table.find_many(where={"team_id": team_id}) + listed: Final = frozenset(itertools.chain.from_iterable(project.models for project in projects)) + served: Final = sorted(model for model in listed if team_write_serves(llm_router, team_id, public_name, model)) + shortfall: Final = next( + ( + (model, added, short) + for model in served + if ( + added := limits_added_by_team_write( + policy, llm_router, team_id, public_name, replaced_id, written_params, model + ) + ) + and (short := _projects_short_on(projects, model, added)) + ), + None, + ) + if shortfall is None: + return + short_model, short_limits, short_projects = shortfall + raise ProxyException( + message=( + f"projects {', '.join(short_projects)} list '{short_model}' without a positive " + f"{' and '.join(sorted(short_limits))} " + f"quota, which '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' requires once this team model change applies. " + "Set those limits on the projects first" + ), + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param=None, + ) + + +async def _raise_if_team_update_leaves_projects_short( + *, + prisma_client: PrismaClient, + general_settings: Mapping[str, object], + llm_router: Router | None, + existing: Deployment, + written_team_id: str | None, + written_name: str, + incoming_params: GenericLiteLLMParams | None, + remaining_params: GenericLiteLLMParams, +) -> None: + if written_team_id is None and existing.model_info.team_id is None: + return + stored_name: Final = existing.model_info.team_public_model_name or existing.model_name + await _raise_if_team_write_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings, + llm_router=llm_router, + team_id=written_team_id, + public_name=written_name, + replaced_id=existing.model_info.id, + written_params=_written_routing_params(incoming_params, remaining_params), + ) + if (written_team_id, written_name) == (existing.model_info.team_id, stored_name): + return + await _raise_if_team_write_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings, + llm_router=llm_router, + team_id=existing.model_info.team_id, + public_name=stored_name, + replaced_id=existing.model_info.id, + written_params=None, + ) + + _PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"}) @@ -1308,6 +1500,7 @@ async def patch_model( ProxyException: For various error conditions including authentication and database errors """ from litellm.proxy.proxy_server import ( + general_settings_view, litellm_proxy_admin_name, llm_router, premium_user, @@ -1434,6 +1627,29 @@ async def patch_model( existing_params=remaining_params, ) + written_team_id: Final = ( + patch_data.model_info.team_id if patch_data.model_info else None + ) or db_model.model_info.team_id + written_name: Final = _get_public_model_name(effective_patch, db_model) + _raise_if_model_update_adds_missing_rate_limits( + existing=db_model, + written_name=written_name, + written_team_model=written_team_id is not None, + incoming_params=incoming_params, + remaining_params=remaining_params, + policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING), + ) + await _raise_if_team_update_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings_view(), + llm_router=llm_router, + existing=db_model, + written_team_id=written_team_id, + written_name=written_name, + incoming_params=incoming_params, + remaining_params=remaining_params, + ) + effective_params: Final = _effective_complexity_router_params(incoming_params, remaining_params) requested_model_name: Final = patch_data.model_name stored_model_name: str | None = None @@ -2413,6 +2629,14 @@ class ModelManagementAuthChecks: ModelManagementAuthChecks._reject_team_inline_api_key_secret_reference( model_params=model_params, incoming_params=incoming_params, incoming_model_params=incoming_model_params ) + raise_auto_router_environment_reference_error( + auto_router_environment_reference_error( + config=_OPTIONAL_JSON_OBJECT.validate_python( + incoming_params.get("complexity_router_config") if incoming_params is not None else None + ), + user_api_key_dict=user_api_key_dict, + ) + ) # Federation fields choose which server-side secret is read and where the org-scoped token # it buys is sent, so only a proxy admin may point a federated deployment somewhere else. # Evaluated on the RESULTING deployment: a patch attaching a credential by name inherits @@ -2509,6 +2733,7 @@ async def delete_model( from litellm.proxy.proxy_server import ( MODEL_RECONCILE_LOCK, + general_settings_view, llm_router, premium_user, prisma_client, @@ -2542,6 +2767,15 @@ async def delete_model( incoming_params=None, allow_missing_team=True, ) + await _raise_if_team_write_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings_view(), + llm_router=llm_router, + team_id=model_params.model_info.team_id, + public_name=model_params.model_info.team_public_model_name or model_params.model_name, + replaced_id=model_params.model_info.id, + written_params=None, + ) # update DB if store_model_in_db is True: @@ -2708,6 +2942,8 @@ async def add_new_model( """ from litellm.proxy.proxy_server import ( general_settings, + general_settings_view, + llm_router, premium_user, prisma_client, proxy_config, @@ -2766,8 +3002,19 @@ async def add_new_model( ) _raise_if_rate_limits_required_but_missing( + model_name=model_params.model_name, litellm_params=model_params.litellm_params, - enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)), + team_model=model_params.model_info.team_id is not None, + policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING), + ) + await _raise_if_team_write_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings_view(), + llm_router=llm_router, + team_id=model_params.model_info.team_id, + public_name=model_params.model_name, + replaced_id=None, + written_params=_written_routing_params(model_params.litellm_params, None), ) clean_model_info: Final = ModelInfo.model_validate( @@ -2906,6 +3153,7 @@ async def update_model( """ from litellm.proxy.proxy_server import ( LITELLM_PROXY_ADMIN_NAME, + general_settings_view, llm_router, premium_user, prisma_client, @@ -2970,6 +3218,30 @@ async def update_model( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, ) + renamed_to: Final = ( + model_params.model_name + if model_params.model_name not in (None, deployment.model_name) and deployment.model_info.team_id is None + else None + ) + written_name: Final = renamed_to or deployment.model_info.team_public_model_name or deployment.model_name + _raise_if_model_update_adds_missing_rate_limits( + existing=deployment, + written_name=written_name, + written_team_model=deployment.model_info.team_id is not None, + incoming_params=model_params.litellm_params, + remaining_params=deployment.litellm_params, + policy=_quota_policy(general_settings_view(), ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING), + ) + await _raise_if_team_update_leaves_projects_short( + prisma_client=prisma_client, + general_settings=general_settings_view(), + llm_router=llm_router, + existing=deployment, + written_team_id=deployment.model_info.team_id, + written_name=written_name, + incoming_params=model_params.litellm_params, + remaining_params=deployment.litellm_params, + ) effective_params: Final = _effective_complexity_router_params( model_params.litellm_params, deployment.litellm_params ) @@ -3009,12 +3281,6 @@ async def update_model( if value is not None or _existing_litellm_params_dict.get(key) is not None } - renamed_to: Final = ( - model_params.model_name - if model_params.model_name not in (None, deployment.model_name) - and deployment.model_info.team_id is None - else None - ) base_update: Final[PrismaCompatibleUpdateDBModel] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py index a4428eb3a85..995446db7a6 100644 --- a/litellm/proxy/management_helpers/auto_router_permissions.py +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -5,8 +5,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException -from pydantic import ConfigDict, Field, ValidationError -from typing_extensions import ReadOnly, TypedDict +from pydantic import ConfigDict, Field, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict, assert_never from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.project import LiteLLM_ProjectTable @@ -36,6 +36,7 @@ from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.router import Router from litellm.router_strategy.complexity_router.config import ( ComplexityRouterConfigWrite, + normalize_classifier_config_aliases, resolve_complexity_router_config_write, ) from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig @@ -47,6 +48,36 @@ if TYPE_CHECKING: from prisma import types as prisma_types +def auto_router_environment_reference_error( + config: Mapping[str, object] | None, user_api_key_dict: UserAPIKeyAuth +) -> Literal["environment_reference_denied"] | None: + if config is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return None + classifier: Final = normalize_classifier_config_aliases(config).get("opensource_classifier_config") + if not isinstance(classifier, Mapping): + return None + fields: Final = TypeAdapter(Mapping[str, object]).validate_python(classifier) + if any( + isinstance(value, str) and value.startswith("os.environ/") + for value in (fields.get("api_key"), fields.get("api_base")) + ): + return "environment_reference_denied" + return None + + +def raise_auto_router_environment_reference_error(error: Literal["environment_reference_denied"] | None) -> None: + match error: + case None: + return + case "environment_reference_denied": + raise HTTPException( + status_code=403, + detail="Only proxy admins can submit classifier connection environment references", + ) + case _: + assert_never(error) + + class _MemberRouterThinking(LiteLLMBaseModel): model_config = ConfigDict(extra="forbid", frozen=True) diff --git a/litellm/proxy/management_helpers/model_quota_policy.py b/litellm/proxy/management_helpers/model_quota_policy.py new file mode 100644 index 00000000000..752c78df862 --- /dev/null +++ b/litellm/proxy/management_helpers/model_quota_policy.py @@ -0,0 +1,327 @@ +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +import litellm +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider, is_openai_compatible_endpoint +from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping +from litellm.llms.openai_like.json_loader import JSONProviderRegistry +from litellm.router_utils.common_utils import provider_for_generic_call +from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm import Router + +QuotaLimit: TypeAlias = Literal["rpm", "tpm"] + +ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota" + +_LIMITS: Final[tuple[QuotaLimit, ...]] = ("rpm", "tpm") +_ROUTER_PATTERNS: Final = PatternMatchRouter() + + +def _is_known_provider(provider: str) -> bool: + return provider in litellm.provider_list or JSONProviderRegistry.exists(provider) + + +class ModelQuotaPolicy(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + default: frozenset[QuotaLimit] = frozenset(_LIMITS) + providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) + models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) + + @field_validator("providers") + @classmethod + def _known_providers(cls, providers: Mapping[str, frozenset[QuotaLimit]]) -> Mapping[str, frozenset[QuotaLimit]]: + unknown: Final = sorted(provider for provider in providers if not _is_known_provider(provider)) + if unknown: + raise ValueError(f"unknown provider ids {unknown}") + return providers + + +REQUIRE_BOTH_LIMITS: Final = ModelQuotaPolicy() +_POLICY_KEYS: Final = frozenset(ModelQuotaPolicy.model_fields) + + +@dataclass(frozen=True, slots=True) +class TeamDeployment: + provider: str | None + called_model: str | None + + +@dataclass(frozen=True, slots=True) +class ModelRoute: + providers: frozenset[str] = frozenset() + team_deployments: frozenset[TeamDeployment] = frozenset() + + +RouteOf: TypeAlias = Callable[[str], ModelRoute] +NO_ROUTE: Final = ModelRoute() + + +def parse_model_quota_setting(setting: object) -> ModelQuotaPolicy | None: + """Any falsy value is off and any other value requires both limits, as the boolean setting always + behaved, so only a mapping that names a policy key is read as a per-provider and per-model policy.""" + if not setting: + return None + if is_str_keyed_mapping(setting) and not _POLICY_KEYS.isdisjoint(setting): + return ModelQuotaPolicy.model_validate(setting) + return REQUIRE_BOTH_LIMITS + + +def _wildcard_matches(pattern: str, model: str) -> bool: + return re.match(_ROUTER_PATTERNS.pattern_to_regex(pattern), model) is not None + + +def _router_specificity(pattern: str) -> tuple[int, int]: + return PatternUtils.calculate_pattern_specificity(_ROUTER_PATTERNS.pattern_to_regex(pattern)) + + +def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[QuotaLimit]]) -> str | None: + matching: Final = tuple(pattern for pattern in entries if "*" in pattern and _wildcard_matches(pattern, model)) + return max(matching, key=_router_specificity, default=None) + + +def _model_entry(entries: Mapping[str, frozenset[QuotaLimit]], model: str) -> frozenset[QuotaLimit] | None: + exact: Final = entries.get(model) + if exact is not None: + return exact + pattern: Final = _most_specific_model_pattern(model, entries) + return None if pattern is None else entries[pattern] + + +def _provider_limits(policy: ModelQuotaPolicy, provider: str | None) -> frozenset[QuotaLimit]: + return policy.default if provider is None else policy.providers.get(provider, policy.default) + + +def _team_deployment_limits(policy: ModelQuotaPolicy, deployment: TeamDeployment) -> frozenset[QuotaLimit]: + entry: Final = None if deployment.called_model is None else _model_entry(policy.models, deployment.called_model) + return _provider_limits(policy, deployment.provider) if entry is None else entry + + +def required_limits(policy: ModelQuotaPolicy, model: str, route: ModelRoute) -> frozenset[QuotaLimit]: + """A model entry (exact name, then the most specific `*` pattern) wins over the providers of the + model's deployments, which win over the default. A model served by several providers needs every + limit any of them requires. A team picks its model's public name, so a team deployment matches + model entries by the model it calls instead.""" + entry: Final = None if route.team_deployments else _model_entry(policy.models, model) + if entry is not None: + return entry + if not route.providers and not route.team_deployments: + return policy.default + return frozenset[QuotaLimit]().union( + *(_team_deployment_limits(policy, deployment) for deployment in route.team_deployments), + *(_provider_limits(policy, provider) for provider in route.providers), + ) + + +def _is_positive(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + +def unmet_limits(required: frozenset[QuotaLimit], values: Mapping[QuotaLimit, object]) -> tuple[QuotaLimit, ...]: + return tuple( + limit + for limit in _LIMITS + if (limit in required or values.get(limit) is not None) and not _is_positive(values.get(limit)) + ) + + +def no_route(_model: str) -> ModelRoute: + return NO_ROUTE + + +def _resolved_api_base(api_base: object) -> object: + if isinstance(api_base, str) and api_base.startswith("os.environ/"): + return get_secret_str(api_base) + return api_base + + +def _provider_behind_api_base(litellm_params: Mapping[str, object]) -> str | None: + model: Final = litellm_params.get("model") + api_base: Final = _resolved_api_base(litellm_params.get("api_base")) + if litellm_params.get("custom_llm_provider") or not isinstance(model, str) or not isinstance(api_base, str): + return None + if model.split("/", 1)[0] in litellm.provider_list or not is_openai_compatible_endpoint(api_base): + return None + try: + _, provider, _, _ = get_llm_provider(model=model, api_base=api_base) + except BadRequestError: + return None + return provider + + +def deployment_provider(litellm_params: object) -> str | None: + if not is_str_keyed_mapping(litellm_params): + return None + return _provider_behind_api_base(litellm_params) or provider_for_generic_call(litellm_params) + + +def _called_model(litellm_params: Mapping[str, object]) -> str | None: + model: Final = litellm_params.get("model") + if not isinstance(model, str): + return None + prefix, _, rest = model.partition("/") + return rest if rest and _is_known_provider(prefix) else model + + +def team_deployment(litellm_params: object) -> TeamDeployment: + if not is_str_keyed_mapping(litellm_params): + return TeamDeployment(provider=None, called_model=None) + return TeamDeployment(provider=deployment_provider(litellm_params), called_model=_called_model(litellm_params)) + + +def deployment_route(litellm_params: Mapping[str, object], team_model: bool) -> ModelRoute: + if team_model: + return ModelRoute(team_deployments=frozenset((team_deployment(litellm_params),))) + provider: Final = deployment_provider(litellm_params) + return ModelRoute(providers=frozenset() if provider is None else frozenset((provider,))) + + +def _named_deployments(llm_router: "Router", model: str, team_id: str | None) -> tuple[Mapping[str, object], ...]: + return ( + *llm_router.get_all_deployments(model_name=model, team_id=team_id), + *llm_router.get_model_list_from_routing_groups(model_name=model), + ) + + +def _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: + target: Final = llm_router.routable_model_group(model) + return _named_deployments(llm_router, target, team_id) or llm_router.wildcard_route_deployments(target, team_id) + + +def _is_team_deployment(deployment: Mapping[str, object]) -> bool: + model_info: Final = deployment.get("model_info") + return is_str_keyed_mapping(model_info) and model_info.get("team_id") is not None + + +def router_model_routes(llm_router: "Router | None", team_id: str | None) -> RouteOf: + if llm_router is None: + return no_route + + def route_of(model: str) -> ModelRoute: + deployments: Final = _routed_deployments(llm_router, model, team_id) + shared: Final = tuple(deployment for deployment in deployments if not _is_team_deployment(deployment)) + return ModelRoute( + providers=frozenset( + provider + for deployment in shared + if (provider := deployment_provider(deployment.get("litellm_params"))) is not None + ), + team_deployments=frozenset( + team_deployment(deployment.get("litellm_params")) + for deployment in deployments + if _is_team_deployment(deployment) + ), + ) + + return route_of + + +def _deployment_id(deployment: Mapping[str, object]) -> object: + model_info: Final = deployment.get("model_info") + return model_info.get("id") if is_str_keyed_mapping(model_info) else None + + +def _routable_name(llm_router: "Router | None", model: str) -> str: + return model if llm_router is None else llm_router.routable_model_group(model) + + +def _patterns_tried_before(patterns: Sequence[str], regex: str) -> Sequence[str]: + if regex in patterns: + return patterns[: patterns.index(regex)] + specificity: Final = PatternUtils.calculate_pattern_specificity(regex) + return tuple(rival for rival in patterns if PatternUtils.calculate_pattern_specificity(rival) >= specificity) + + +def _team_pattern_serves(llm_router: "Router", team_id: str, pattern: str, model: str) -> bool: + regex: Final = _ROUTER_PATTERNS.pattern_to_regex(pattern) + if re.match(regex, model) is None or _named_deployments(llm_router, model, team_id): + return False + team_router: Final = llm_router.team_pattern_routers.get(team_id) + patterns: Final = () if team_router is None else tuple(team_router.patterns) + return not any(re.match(rival, model) is not None for rival in _patterns_tried_before(patterns, regex)) + + +def team_write_serves(llm_router: "Router | None", team_id: str, public_name: str, model: str) -> bool: + """Whether a request from the team for ``model`` reaches the team route named ``public_name``: the name + itself or an alias of it, or, for a `*` name, a model no named deployment shadows and no team pattern the + router tries first takes.""" + target: Final = _routable_name(llm_router, model) + if "*" not in public_name or llm_router is None: + return target == public_name + return _team_pattern_serves(llm_router, team_id, public_name, target) + + +def _served_params(litellm_params: object, public_name: str, model: str) -> object: + if "*" not in public_name or not is_str_keyed_mapping(litellm_params): + return litellm_params + matched: Final = re.match(_ROUTER_PATTERNS.pattern_to_regex(public_name), model) + called: Final = litellm_params.get("model") + if matched is None or not isinstance(called, str): + return litellm_params + return MappingProxyType( + { + **litellm_params, + "model": PatternMatchRouter.set_deployment_model_name( + matched_pattern=matched, litellm_deployment_litellm_model=called + ), + } + ) + + +def team_route_after_write( + llm_router: "Router | None", + team_id: str, + public_name: str, + replaced_id: str | None, + written_params: Mapping[str, object] | None, + model: str, +) -> ModelRoute | None: + deployments: Final = () if llm_router is None else llm_router.get_all_deployments(public_name, team_id=team_id) + kept: Final = tuple( + deployment.get("litellm_params") + for deployment in deployments + if _is_team_deployment(deployment) and _deployment_id(deployment) != replaced_id + ) + team_params: Final = kept if written_params is None else (*kept, written_params) + target: Final = _routable_name(llm_router, model) + if team_params: + return ModelRoute( + team_deployments=frozenset( + team_deployment(_served_params(litellm_params, public_name, target)) for litellm_params in team_params + ) + ) + if llm_router is None or not _routed_deployments(llm_router, model, None): + return None + return router_model_routes(llm_router, None)(model) + + +def limits_added_by_team_write( + policy: ModelQuotaPolicy, + llm_router: "Router | None", + team_id: str, + public_name: str, + replaced_id: str | None, + written_params: Mapping[str, object] | None, + model: str, +) -> frozenset[QuotaLimit]: + after: Final = team_route_after_write(llm_router, team_id, public_name, replaced_id, written_params, model) + if after is None: + return frozenset() + now: Final = router_model_routes(llm_router, team_id)(model) + return required_limits(policy, model, after) - required_limits(policy, model, now) + + +def exempts_any_route(policy: ModelQuotaPolicy) -> bool: + return any( + limits != REQUIRE_BOTH_LIMITS.default + for limits in (policy.default, *policy.providers.values(), *policy.models.values()) + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9eef4f0ef9d..e478096f7ad 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -87,6 +87,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, model_deployments_support_stream_options, open_sse_before_first_byte, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -99,10 +100,16 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_headers, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) +from litellm.proxy.common_utils.token_count_log import ( + deployment_model_id, + logged_provider_token_count, + rejected_token_count, +) from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, @@ -1259,6 +1266,7 @@ async def handle_bedrock_count_tokens( from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.proxy.proxy_server import llm_router + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) try: # Initialize the handler handler: Final = BedrockCountTokensHandler() @@ -1266,40 +1274,45 @@ async def handle_bedrock_count_tokens( # Extract model from request body model: Final = _optional_str(request_body.get("model")) if not model: - raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) + raise rejected_token_count( + HTTPException(status_code=400, detail={"error": "Model is required in request body"}), + litellm_call_id, + None, + ) # Get model parameters from router litellm_params: Final = {"user_api_key_dict": user_api_key_dict} resolved_model = model # Default fallback - if llm_router: - deployments: Final = llm_router.get_model_list(model_name=model) - if deployments and len(deployments) > 0: - # Get the first matching deployment - deployment: Final = deployments[0] - model_litellm_params: Final = deployment.get("litellm_params", {}) + deployments: Final = llm_router.get_model_list(model_name=model) if llm_router else None + deployment: Final = deployments[0] if deployments else None + if deployment is not None: + model_litellm_params: Final = deployment.get("litellm_params", {}) - # Get the resolved model ID from the configuration - if "model" in model_litellm_params: - resolved_model = model_litellm_params["model"] + # Get the resolved model ID from the configuration + if "model" in model_litellm_params: + resolved_model = model_litellm_params["model"] - # Copy all litellm_params - BaseAWSLLM will handle AWS credential discovery - for key, value in model_litellm_params.items(): - if key != "user_api_key_dict": # Don't overwrite user_api_key_dict - litellm_params[key] = value + # Copy all litellm_params - BaseAWSLLM will handle AWS credential discovery + for key, value in model_litellm_params.items(): + if key != "user_api_key_dict": # Don't overwrite user_api_key_dict + litellm_params[key] = value verbose_proxy_logger.debug("Count tokens litellm_params: %s", litellm_params) verbose_proxy_logger.debug("Resolved model: %s", resolved_model) - # Handle the count tokens request - result: Final = await handler.handle_count_tokens_request( - request_data=request_body, - litellm_params=litellm_params, - resolved_model=resolved_model, + return await logged_provider_token_count( + handler.handle_count_tokens_request( + request_data=request_body, + litellm_params=litellm_params, + resolved_model=resolved_model, + ), + litellm_call_id=litellm_call_id, + request_model=model, + model_id=deployment_model_id(deployment), + model_used=partial(handler.count_tokens_model_id, resolved_model), ) - return result - except BedrockError as e: # Convert BedrockError to HTTPException for FastAPI verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f2ea1752a86..650dcae3239 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -54,6 +54,7 @@ from litellm.integrations.custom_guardrail import ( without_server_streaming_classification, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, @@ -151,6 +152,7 @@ from .success_handler import PassThroughEndpointLogging from .upstream_usage_headers import ( UpstreamReportedUsage, apply_upstream_reported_model, + apply_upstream_reported_operation, apply_upstream_reported_usage, ) @@ -1134,6 +1136,7 @@ async def pass_through_request( guardrails_config: dict | None = None, timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, + operation: GenAIOperation | None = None, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -1155,6 +1158,7 @@ async def pass_through_request( timeout: Optional per-endpoint timeout in seconds. Falls back to general_settings.pass_through_request_timeout, then 600s. upstream_headers_filter: Optional last rewrite of the headers sent upstream, after x-pass- expansion + operation: Optional GenAI operation the route reports, unless the upstream names its own """ from litellm.exceptions import ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging @@ -1353,6 +1357,8 @@ async def pass_through_request( call_type="pass_through_endpoint", ) logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + if operation is not None: + logging_obj.model_call_details[PASSTHROUGH_OPERATION_KEY] = operation ## PASSTHROUGH MANAGED ID RESOLUTION (INPUT) ## # Resolve managed IDs in path, query params, and body back to raw @@ -1553,6 +1559,7 @@ async def pass_through_request( logging_obj=logging_obj, headers=response.headers, ) + apply_upstream_reported_operation(logging_obj, response.headers) relay_response: Final = await _log_passthrough_upstream_failure( response=response, @@ -1646,6 +1653,7 @@ async def pass_through_request( logging_obj=logging_obj, headers=response.headers, ) + apply_upstream_reported_operation(logging_obj, response.headers) if _is_streaming_response(response) is True: logging_obj.stream = True @@ -2161,6 +2169,7 @@ def create_pass_through_route( config_file_path: str | None = None, timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, + operation: GenAIOperation | None = None, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -2301,6 +2310,7 @@ def create_pass_through_route( guardrails_config=cast(dict | None, param_guardrails), timeout=cast(float | None, param_timeout), upstream_headers_filter=upstream_headers_filter, + operation=operation, ) finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): @@ -3232,6 +3242,7 @@ class InitPassThroughEndpointHelpers: guardrails=guardrails, config_file_path=config_file_path, timeout=timeout, + operation=GenAIOperation.LITELLM_PASSTHROUGH, ), methods=methods, dependencies=dependencies, @@ -3316,6 +3327,7 @@ class InitPassThroughEndpointHelpers: guardrails=guardrails, config_file_path=config_file_path, timeout=timeout, + operation=GenAIOperation.LITELLM_PASSTHROUGH, ), methods=methods, dependencies=dependencies, diff --git a/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py b/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py index e9d93fb306d..223f6a5e760 100644 --- a/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py +++ b/litellm/proxy/pass_through_endpoints/upstream_usage_headers.py @@ -5,18 +5,21 @@ price the request from the response body. Those targets report the totals for the whole HTTP request in ``x-litellm-*`` response headers instead; LiteLLM records the reported values without recomputing them. A target may also name the request in ``x-litellm-upstream-model`` so its spend can be told apart from -other requests that carry no ``model`` in the body. +other requests that carry no ``model`` in the body, and name the GenAI +operation it performed, which labels the request's trace. """ import math from collections.abc import Mapping from dataclasses import dataclass +from types import MappingProxyType from typing import Final import httpx from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_SPEND_LOG_MODEL_NAME_LENGTH +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import Usage @@ -25,6 +28,8 @@ UPSTREAM_TOTAL_TOKENS_HEADER: Final = "x-litellm-total-tokens" # Deliberately not x-litellm-model-name/-id/-group: the gateway already sends # those to clients, and pass-through relays upstream headers back to the client. UPSTREAM_MODEL_HEADER: Final = "x-litellm-upstream-model" +UPSTREAM_OPERATION_NAME_HEADER: Final = "x-litellm-operation-name" +_OPERATIONS_BY_NAME: Final = MappingProxyType({operation.value: operation for operation in GenAIOperation}) # model_call_details key holding what the upstream reported, so later stages of # the success path can tell an upstream-reported cost apart from one LiteLLM @@ -175,3 +180,19 @@ def has_upstream_reported_usage(logging_obj: LiteLLMLoggingObj) -> bool: means zero, never a fallback to someone else's estimate. """ return isinstance(logging_obj.model_call_details.get(UPSTREAM_REPORTED_USAGE_KEY), UpstreamReportedUsage) + + +def apply_upstream_reported_operation(logging_obj: LiteLLMLoggingObj, headers: Mapping[str, str]) -> None: + """Record the GenAI operation the upstream named, overriding the route's label. An unknown name is ignored.""" + raw_operation: Final = headers.get(UPSTREAM_OPERATION_NAME_HEADER) + if raw_operation is None: + return + operation: Final = _OPERATIONS_BY_NAME.get(raw_operation) + if operation is None: + verbose_proxy_logger.warning( + "pass_through_endpoint: upstream sent unknown %s=%r; ignoring it", + UPSTREAM_OPERATION_NAME_HEADER, + raw_operation, + ) + return + logging_obj.model_call_details[PASSTHROUGH_OPERATION_KEY] = operation # rebind-ok: annotates the request's record diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04f84f1f7b5..64fb2791913 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -100,7 +100,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, offload_token_count from litellm.llms.bedrock.count_tokens.kimi_k3 import ( HUGGING_FACE_SOURCE, KimiK3Encodings, @@ -552,6 +552,11 @@ from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_settings, get_budget_reset_time, ) +from litellm.proxy.common_utils.token_count_log import ( + deployment_model_id, + logged_token_count, + rejected_token_count, +) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, @@ -14848,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]: @@ -14870,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) @@ -15253,16 +15270,21 @@ def _system_message(system: object) -> ChatCompletionSystemMessage | None: dependencies=[Depends(user_api_key_auth)], response_model=TokenCountResponse, ) -async def token_counter(request: TokenCountRequest, call_endpoint: bool = False): +async def token_counter( + request: TokenCountRequest, + call_endpoint: bool = False, + litellm_call_id: Annotated[str | None, Header(alias=LITELLM_CALL_ID_HEADER)] = None, +): """ Args: request: TokenCountRequest call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs. + litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs Returns: TokenCountResponse """ - return await count_request_tokens(request=request, call_endpoint=call_endpoint) + return await count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=litellm_call_id) async def count_request_tokens( @@ -15273,7 +15295,69 @@ async def count_request_tokens( context_management: JsonValue = None, request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), responses_request: Mapping[str, object] | None = None, + litellm_call_id: str | None = None, client_oauth_token: str | None = None, +) -> TokenCountResponse: + call_id: Final = resolve_litellm_call_id(litellm_call_id) + if request.prompt is None and request.messages is None and request.contents is None: + raise rejected_token_count( + HTTPException(status_code=400, detail="prompt or messages or contents must be provided"), + call_id, + request.model, + ) + routed: Final = await _routed_count_deployment(request.model, responses_request) + count: Final = _count_request_tokens( + request=request, + routed=routed, + call_endpoint=call_endpoint, + messages_format=messages_format, + output_format=output_format, + context_management=context_management, + request_litellm_params=request_litellm_params, + responses_request=responses_request, + client_oauth_token=client_oauth_token, + ) + return await logged_token_count( + _with_anthropic_content_errors(count) if messages_format == "anthropic" else count, + litellm_call_id=call_id, + request_model=request.model, + model_id=deployment_model_id(routed), + ) + + +async def _with_anthropic_content_errors(count: Awaitable[TokenCountResponse]) -> TokenCountResponse: + try: + return await count + except UnsupportedContentTypeError as e: + raise ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400) from e + + +async def _routed_count_deployment( + model: str, responses_request: Mapping[str, object] | None +) -> Mapping[str, object] | None: + if llm_router is None: + return None + try: + return await llm_router.async_get_available_deployment( + model=model, request_kwargs=dict(_responses_routing_kwargs(responses_request)) + ) + except Exception: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment" + ) + return None + + +async def _count_request_tokens( + request: TokenCountRequest, + routed: Mapping[str, object] | None, + call_endpoint: bool, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + context_management: JsonValue, + request_litellm_params: Mapping[str, JsonValue], + responses_request: Mapping[str, object] | None, + client_oauth_token: str | None, ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -15281,26 +15365,9 @@ async def count_request_tokens( tools: Final = request.tools system: Final = request.system - ######################################################### - # Validate request - ######################################################### - if prompt is None and messages is None and contents is None: - raise HTTPException(status_code=400, detail="prompt or messages or contents must be provided") - - deployment: dict[str, Any] | None = None + deployment: dict[str, Any] | None = None if routed is None else dict(routed) litellm_model_name = None model_info: ModelMapInfo | None = None - if llm_router is not None: - # get 1 deployment corresponding to the model - try: - deployment = await llm_router.async_get_available_deployment( - model=request.model, - request_kwargs=dict(_responses_routing_kwargs(responses_request)), - ) - except Exception: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment" - ) if deployment is not None: litellm_model_name = deployment.get("litellm_params", {}).get("model") model_info = deployment.get("model_info", {}) @@ -15400,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, ) @@ -19899,12 +19967,14 @@ async def update_config_general_settings( ) try: - ConfigGeneralSettings.model_validate({data.field_name: requested_value}) + validated: Final = ConfigGeneralSettings.model_validate({data.field_name: requested_value}) except Exception: raise HTTPException( status_code=400, detail={"error": f"Invalid type of field value={type(requested_value)} passed in."}, ) + parsed_value: Final[object] = _UNTYPED_VALUE.validate_python(getattr(validated, data.field_name)) + stored_value: Final = parsed_value if isinstance(parsed_value, bool) else requested_value if data.field_name == "alerting_args": try: @@ -19933,10 +20003,10 @@ async def update_config_general_settings( proxy_config.reject_config_owned_writes( section_name="general_settings", - changed_keys={data.field_name: cast(JsonValue, requested_value)}, # cast-ok: validated above + changed_keys={data.field_name: cast(JsonValue, stored_value)}, # cast-ok: validated above ) - field_value = requested_value + field_value = stored_value if data.field_name == "plugins": field_value = _preserve_redacted_plugin_keys(field_value, general_settings.get("plugins")) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 0dd96c856f8..f2d88702470 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -4145,6 +4145,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "ZeroGPU", + "provider_display_name": "ZeroGPU", + "litellm_provider": "zerogpu", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.zerogpu.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "zerogpu/gpt-oss-120b" + }, { "provider": "CURSOR", "provider_display_name": "Cursor", diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 97bbf893e3e..4f8b1697f7e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -56,6 +56,7 @@ from litellm.proxy.common_request_processing import ( cost_headers, create_response, resolve_counted_tool_fields, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports @@ -69,6 +70,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, openai_error_type_for_status, ) +from litellm.proxy.common_utils.token_count_log import rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id @@ -1407,7 +1409,12 @@ class _TokenCountPayload(TypedDict): class _TokenCounter(Protocol): def __call__( - self, *, request: TokenCountRequest, call_endpoint: bool, responses_request: Mapping[str, object] | None + self, + *, + request: TokenCountRequest, + call_endpoint: bool, + responses_request: Mapping[str, object] | None, + litellm_call_id: str | None, ) -> Awaitable[TokenCountResponse]: ... @@ -1648,18 +1655,23 @@ async def responses_input_tokens( Returns: `{"object": "response.input_tokens", "input_tokens": }` """ + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) data: Final = await read_request_body(request=request) model_name: Final = data.get("model") input_value: Final = data.get("input") if not isinstance(model_name, str) or not model_name: - return _missing_responses_param_response("model") + return rejected_token_count(_missing_responses_param_response("model"), litellm_call_id, None) if input_value is None: - return _missing_responses_param_response("input") + return rejected_token_count(_missing_responses_param_response("input"), litellm_call_id, model_name) if isinstance(input_value, (str, list)) and not input_value: - return _responses_invalid_request_response( - message="""One of "input" or "previous_response_id" or 'prompt' or 'conversation' must be provided.""", - param=None, - code="missing_required_parameter", + return rejected_token_count( + _responses_invalid_request_response( + message="""One of "input" or "previous_response_id" or 'prompt' or 'conversation' must be provided.""", + param=None, + code="missing_required_parameter", + ), + litellm_call_id, + model_name, ) counted_fields: Final = { @@ -1667,7 +1679,11 @@ async def responses_input_tokens( } token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id) if isinstance(token_fields, ProxyException): - return _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None) + return rejected_token_count( + _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None), + litellm_call_id, + model_name, + ) try: payload: Final[_TokenCountPayload] = { @@ -1680,27 +1696,44 @@ async def responses_input_tokens( } token_request: Final = TokenCountRequest.model_validate(payload) except ValidationError as e: - return _responses_invalid_request_response( - message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + return rejected_token_count( + _responses_invalid_request_response( + message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + ), + litellm_call_id, + model_name, ) except Exception as e: - return _responses_invalid_request_response( - message=f"Invalid request for token counting: {e}", param=None, code=None + return rejected_token_count( + _responses_invalid_request_response( + message=f"Invalid request for token counting: {e}", param=None, code=None + ), + litellm_call_id, + model_name, ) request_fields: Final = _PROXY_SERVER_REQUEST_ADAPTER.validate_python(data) - responses_request: Final = await _with_managed_file_id_mapping( - _with_authorized_previous_response( - {**{name: value for name, value in request_fields.items() if name not in counted_fields}, **token_fields}, - response_id_security, + try: + responses_request: Final = await _with_managed_file_id_mapping( + _with_authorized_previous_response( + { + **{name: value for name, value in request_fields.items() if name not in counted_fields}, + **token_fields, + }, + response_id_security, + user_api_key_dict, + ), + managed_files, user_api_key_dict, - ), - managed_files, - user_api_key_dict, - ) + ) + except ProxyException as e: + raise rejected_token_count(e, litellm_call_id, model_name) try: token_response: Final = await token_counter( - request=token_request, call_endpoint=True, responses_request=responses_request + request=token_request, + call_endpoint=True, + responses_request=responses_request, + litellm_call_id=litellm_call_id, ) except ProxyException as e: if e.type != _PROVIDER_COUNT_ERROR_TYPE: 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/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index b5ef9e95741..219944de521 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -423,6 +423,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._is_tool_search_call(call_id, fn_name): self._queue_tool_search_call_added(call_id, output_index, fn_args_delta) continue + self._started_tool_call = True if call_id not in self._tool_args_by_call_id: self._tool_args_by_call_id[call_id] = "" 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 50ab3e4f75a..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, @@ -365,6 +369,7 @@ class BaseResponsesAPIStreamingIterator: """ _pending_logging_tasks: tuple[asyncio.Task[object], ...] = () + _started_tool_call: bool = False _earlier_rounds_spend: BilledRoundsSpend | None = None def __init__( @@ -488,12 +493,16 @@ class BaseResponsesAPIStreamingIterator: _args_delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_args_delta, str): self._generated_tool_arguments += _args_delta + if _event_type in _CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS: + self._started_tool_call = True _stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata) if _event_type in ( ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): _item: Final[object] = getattr(openai_responses_api_chunk, "item", None) + if getattr(_item, "type", None) in _CLIENT_RUN_TOOL_CALL_ITEM_TYPES: + self._started_tool_call = True if _item is not None: ResponsesAPIRequestUtils.encode_container_id_on_output_item( item=_item, @@ -573,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 @@ -726,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 "", @@ -976,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. @@ -985,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() @@ -1044,21 +1088,31 @@ class BaseResponsesAPIStreamingIterator: if pending: await asyncio.wait(pending) + @property + def started_tool_call(self) -> bool: + """Whether the stream began a tool call, which a text continuation could never finish.""" + return self._started_tool_call + def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: """ The response.incomplete event to bill when the client went away mid-stream, 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 @@ -1072,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: @@ -1208,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 @@ -1293,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 @@ -1690,13 +1753,26 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None: return None -_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset( +_CLIENT_RUN_TOOL_CALL_ITEM_TYPES: Final = frozenset( + { + "function_call", + "custom_tool_call", + "mcp_approval_request", + "computer_call", + "local_shell_call", + "shell_call", + "apply_patch_call", + } +) +_CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset( { ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA, - ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, } ) +_TOOL_ARGUMENTS_DELTA_EVENTS: Final = _CLIENT_RUN_TOOL_ARGUMENTS_DELTA_EVENTS | { + ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA +} def _estimate_usage_from_text( @@ -1704,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 @@ -1711,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 @@ -1728,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( @@ -1735,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) @@ -2253,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 e53bab4659e..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, @@ -305,14 +307,20 @@ from litellm.router_utils.kubernetes_pod_discovery import ( ) from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_MARKER, + RESPONSES_CONTINUATION_KWARG, ContinuationOutput, ContinuationPrefillDeploymentCheck, ContinuedPartial, + ResponsesContinuation, + ResponsesContinuationSeam, constrains_continuation, + constrains_responses_continuation, continuation_deployments, continuation_output_ceilings, continued_text, deployment_can_continue, + responses_continuation_input, + responses_continuation_input_for, ) from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, @@ -457,6 +465,7 @@ from litellm.utils import ( provider_rejectable_params, resolve_api_key_secret_reference, set_live_deployment_replay, + token_counter, ) from .router_utils.pattern_match_deployments import PatternMatchRouter @@ -488,7 +497,6 @@ if TYPE_CHECKING: from litellm.types.llms.openai import ( ResponseAPIUsage, ResponseIncompleteEvent, - ResponseInputParam, ResponsesAPIResponse, ) @@ -625,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]]: @@ -686,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 () @@ -826,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 @@ -956,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 @@ -3865,7 +3888,8 @@ class Router: ) -> Mapping[str, object]: """The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies.""" deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} - defaults: Final = {key: value for key, value in self.default_litellm_params.items() if value is not None} + default_params: Final = _REQUEST_KWARGS_ADAPTER.validate_python(self.default_litellm_params) + defaults: Final = {key: value for key, value in default_params.items() if value is not None} deployment_tool_choice: Final = ( {"tool_choice": deployment_params["tool_choice"]} if deployment_params.get("tool_choice") is not None @@ -3944,67 +3968,6 @@ class Router: extended: Final = _extend_trailing_text_block(blocks, new_text) return [*messages[:-1], {**last, "content": extended, "prefix": True}] - @staticmethod - def _build_responses_continuation_input( - input_val: Union[str, "ResponseInputParam"] | None, - generated_content: str, - ) -> "ResponseInputParam": - """ - Convert Responses-API input + partial assistant output into a - continuation input that asks the fallback model to pick up where the - prior assistant message stopped. - - Best effort across providers. The chat-completions path uses - Anthropic's `prefix: True` prefill trick on the assistant message; - the Responses-API input schema has no direct equivalent, so we - append an instruction (developer role) plus a prior assistant - message containing the partial output. Providers without prefill - semantics (OpenAI, Vertex) treat this as conversational context - and may regenerate — same trade-off as the chat-completions path - for non-Anthropic fallbacks. - """ - # base/continuation are List[Any] because ResponseInputParam items - # are a wide Union of TypedDicts (EasyInputMessageParam, Message, - # ResponseOutputMessageParam, ...) — annotating as List[Dict[str, Any]] - # rejects the list() spread of input_val. We cast the combined list to - # ResponseInputParam at the return. - base: list[object] - if isinstance(input_val, str): - base = [ - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": input_val}], - } - ] - elif isinstance(input_val, list): - base = list(input_val) - else: - base = [] - continuation: Final[list[object]] = [ - { - "type": "message", - "role": "developer", - "content": [ - { - "type": "input_text", - "text": ( - "The previous assistant response was interrupted " - "mid-stream. Continue exactly where it stopped — " - "do not repeat any of its content. Your response " - "must read as a seamless continuation." - ), - } - ], - }, - { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": generated_content}], - }, - ] - return cast("ResponseInputParam", base + continuation) - async def _aresponses_streaming_iterator( self, response: "BaseResponsesAPIStreamingIterator", @@ -4026,9 +3989,12 @@ class Router: Full parity with the chat-completions path: - Pre-first-chunk: retry with the original input unchanged. - - Partial content: inject a developer instruction + prior - assistant message carrying the generated text so the fallback - model continues rather than restarts. + - Partial content: continue from the generated text, as an + assistant prefill on a fallback that takes one and as assistant + history plus a user request to continue on any other. A request + whose output is constrained (a JSON schema, a forced tool), or a + stream that began a tool call, is not continued, and the primary's + error surfaces instead. - Usage combining: merge partial-stream usage onto the fallback's response.completed event so accounting reflects both attempts. - Stream cleanup: shielded aclose() on both source and fallback @@ -4178,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: @@ -4192,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 @@ -4217,6 +4190,90 @@ 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", + source_iterator: "BaseResponsesAPIStreamingIterator", + request_kwargs: Mapping[str, object], + ) -> ResponsesContinuation | None: + """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 _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) + cut_deployment_id: Final = (get_hidden_params(source_iterator) or {}).get("model_id") + cut_deployment: Final = ( + self.get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None + ) + if constrains_responses_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)): + Router._raise_original_mid_stream_error(e) + earlier: Final = request_kwargs.get(RESPONSES_CONTINUATION_KWARG) + if isinstance(earlier, ResponsesContinuation): + return earlier.extended(e.generated_content) + return ResponsesContinuation(request_kwargs.get("input"), e.generated_content) + + @staticmethod + def _responses_continuation_ceilings( + e: "MidStreamFallbackError", request_kwargs: Mapping[str, object] + ) -> Mapping[str, object]: + """The caller's output ceilings less what the cut stream already emitted, raising the primary's error once + nothing is left.""" + emitted_tokens: Final = token_counter(model=e.model or "", text=e.generated_content) + reduced: Final = continuation_output_ceilings(request_kwargs, emitted_tokens) + if reduced is None: + Router._raise_original_mid_stream_error(e) + return reduced + async def _aresponses_fallback_attempt( self, e: "MidStreamFallbackError", @@ -4235,6 +4292,11 @@ class Router: """ from litellm.exceptions import MidStreamFallbackError + continuation: Final = self._responses_continuation(e, source_iterator, initial_kwargs) + reduced_ceilings: Final[Mapping[str, object]] = ( + Router._responses_continuation_ceilings(e, initial_kwargs) if continuation is not None else {} + ) + seam: Final = ResponsesContinuationSeam(continuation.emitted_text if continuation is not None else "") partial_usage: Final = Router._extract_partial_responses_usage(source_iterator) fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted fallback_yielded = False # rebind-ok: flipped on the first fallback item so a fallback that dies before its first event still replays the primary's held announcement @@ -4252,11 +4314,10 @@ class Router: initial_kwargs["original_function"] = ( # rebind-ok: the fallback chain re-enters on the same kwargs self._ageneric_api_call_with_fallbacks_responses_attempt ) - if e.generated_content and not e.is_pre_first_chunk: - initial_kwargs["input"] = Router._build_responses_continuation_input( # rebind-ok: fallback hop input - initial_kwargs.get("input"), - e.generated_content, - ) + if continuation is not None: + initial_kwargs.update(reduced_ceilings) + initial_kwargs[RESPONSES_CONTINUATION_KWARG] = continuation # rebind-ok: shaped per fallback deployment + initial_kwargs["input"] = responses_continuation_input(continuation) # rebind-ok: fallback hop input # The Responses-API path stores observability metadata # under "litellm_metadata" (not the default "metadata") — # see _ageneric_api_call_with_fallbacks. Mirroring that @@ -4285,7 +4346,7 @@ class Router: ) prepared_fallback_hidden_params: Final = adopt_headers(fallback_response) if hasattr(fallback_response, "__aiter__"): - async for fallback_item in fallback_response: + async for fallback_item in seam.splice(fallback_response): Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if partial_usage is not None: Router._combine_responses_fallback_usage(fallback_item, partial_usage) @@ -4717,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 @@ -4808,11 +4875,12 @@ class Router: splits_stream_timeout: bool = False, ) -> None: """ - 4 jobs: + 5 jobs: - Adds selected deployment, model_info and api_base to kwargs["metadata"] (used for logging) - Adds default litellm params to kwargs, if set. - Merges tools from deployment with request (proxy-configured tools + request tools). - On a fallback hop, drops the encrypted reasoning this deployment cannot decrypt, keeping its summary. + - On a Responses continuation, sends the partial response as prefill where this deployment takes one. """ for key in self._forwarded_alias_marker_keys_the_deployment_sets( deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) @@ -5016,6 +5084,11 @@ class Router: (_FALLBACK_HOP_ADAPTER.validate_python(deployment),), unmarked_origin=deployment_that_just_failed, ) + if RESPONSES_CONTINUATION_KWARG in hop_kwargs: + del kwargs[RESPONSES_CONTINUATION_KWARG] # rebind-ok: this method prepares the call's kwargs in place + kwargs["input"] = responses_continuation_input_for( # rebind-ok: the deployment's continuation input + hop_kwargs, _FALLBACK_HOP_ADAPTER.validate_python(deployment) + ) def _get_async_openai_model_client(self, deployment: dict, kwargs: dict): """ @@ -6383,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 ) @@ -6492,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: @@ -6715,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 @@ -8730,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}" @@ -9053,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 @@ -9132,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, @@ -9228,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): @@ -9878,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. @@ -9896,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 @@ -9908,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] @@ -10281,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( @@ -13862,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, @@ -13875,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 () @@ -14444,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 """ @@ -14459,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 @@ -14484,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", {}) @@ -14513,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") @@ -14567,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 @@ -14594,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_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 3df1803d264..d06248a555f 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -477,6 +477,12 @@ unless `modality_routing` is also on. ### Session pin retention +With `classification_mode: user_turn`, a child request carrying `x-parent-session-id` classifies +its own new asks and stores the tier under its own session ID. Its continuations reuse that tier +without replacing the parent's tier. Existing trace grouping, per-tier model affinity and deployment +affinity retain their session identities. A child without a distinct, usable session ID classifies +without storing a tier pin. `session_affinity: true` retains precedence and the existing shared-pin behavior + `session_affinity_ttl_seconds` is the idle window for both the model pin selected by session affinity and the deployment pin. Every request that reuses a pin refreshes its TTL, so a session actively sending requests stays pinned. After the window passes with no pin reuse, the next request classifies again and creates a fresh pin. Omit the setting to track the default of 3600 seconds. ### Mid-task stall escalation diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 82aaa221d31..43af0485a2d 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -23,6 +23,7 @@ import random import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass from functools import partial from itertools import accumulate, chain, islice, takewhile from threading import Lock @@ -37,6 +38,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.affinity_cache import ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, read_affinity_pin from litellm.constants import ( + COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, OUTPUT_TOKEN_CEILING_PARAMS, @@ -1389,6 +1391,11 @@ def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) return {"model": model, "tier": tier_value} +@dataclass(frozen=True, slots=True) +class _JevConnectionError: + message: str + + class ComplexityRouter(CustomLogger): """ Complexity router that classifies requests and routes to appropriate models. @@ -1404,7 +1411,7 @@ class ComplexityRouter(CustomLogger): - Question complexity (multiple questions) """ - def _build_jev_client(self, config: OpenSourceClassifierConfig) -> JevClassifierClient: + def _build_jev_client(self, config: OpenSourceClassifierConfig) -> JevClassifierClient | _JevConnectionError: if config.deployment_name is not None: return RouterJevClassifierClient(self.litellm_router_instance, config.deployment_name) if config.provider in ("laya", "bespoke"): @@ -1425,12 +1432,22 @@ class ComplexityRouter(CustomLogger): http_client=get_async_httpx_client(httpxSpecialProvider.PassThroughEndpoint), provider="databricks", ) - api_key: Final = config.api_key or get_secret_str("TYPESAFE_API_KEY") - if not api_key: - raise ValueError( + api_key: Final = ( + get_secret_str(config.api_key) + if config.api_key is not None and config.api_key.startswith("os.environ/") + else config.api_key or get_secret_str("TYPESAFE_API_KEY") + ) + if not api_key or not api_key.strip(): + return _JevConnectionError( "opensource_classifier_config.api_key or TYPESAFE_API_KEY is required for classifier_type 'oss_classifier'" ) - api_base: Final = config.api_base or get_secret_str("TYPESAFE_API_BASE") or "https://api.typesafe.ai" + api_base: Final = ( + get_secret_str(config.api_base) + if config.api_base is not None and config.api_base.startswith("os.environ/") + else config.api_base or get_secret_str("TYPESAFE_API_BASE") or "https://api.typesafe.ai" + ) + if not api_base or not api_base.strip(): + return _JevConnectionError("opensource_classifier_config.api_base must resolve to a non-empty value") return HttpJevClassifierClient( api_key=api_key, api_base=api_base, @@ -1473,13 +1490,18 @@ class ComplexityRouter(CustomLogger): self.config.default_model = default_model jev_config: Final = self.config.opensource_classifier_config - self._jev_client: JevClassifierClient | None = ( + jev_client_result: Final = ( jev_client if jev_client is not None else self._build_jev_client(jev_config) if self.config.classifier_type == "oss_classifier" and jev_config is not None else None ) + match jev_client_result: + case _JevConnectionError(message=message): + raise ValueError(message) + case _: + self._jev_client: JevClassifierClient | None = jev_client_result self._tier_affinity_config = hashlib.sha256( self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode() @@ -4358,17 +4380,17 @@ class ComplexityRouter(CustomLogger): return newest_ask is not None and newest_ask.startswith(_CODEX_COMPACTION_SUMMARY_PREFIX) @staticmethod - def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: + def _iter_metadata_dicts(request_kwargs: Mapping[str, object]) -> list[Mapping[str, object]]: """Metadata may land on `metadata` or `litellm_metadata` depending on the endpoint, mirroring DeploymentAffinityCheck's precedence.""" return [ - metadata + cast(Mapping[str, object], metadata) # cast-ok: request metadata is a str-keyed JSON object for metadata_key in ("litellm_metadata", "metadata") if isinstance(metadata := request_kwargs.get(metadata_key), dict) ] @staticmethod - def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: + def _get_session_id_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: """Resolve a client-supplied session_id.""" for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): session_id = metadata.get("session_id") @@ -4376,6 +4398,21 @@ class ComplexityRouter(CustomLogger): return str(session_id) return None + def _get_tier_pin_session_id(self, request_kwargs: Mapping[str, object]) -> str | None: + if self.config.classification_mode == "user_turn" and not self.config.session_affinity: + child_metadata: Final = next( + ( + metadata + for metadata in self._iter_metadata_dicts(request_kwargs) + if COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY in metadata + ), + None, + ) + if child_metadata is not None: + child_session: Final = child_metadata[COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY] + return child_session if isinstance(child_session, str) and child_session else None + return self._get_session_id_from_request_kwargs(request_kwargs) + @staticmethod def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None: """Resolve the proxy-derived API key hash, the same trust boundary @@ -4460,7 +4497,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing: Final = _conversation_is_continuing(resolved_messages) use_session_affinity: Final = self._uses_tier_pin - session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None + session_id: Final = self._get_tier_pin_session_id(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None # In 'user_turn' mode a held pin is replayed only on continuation turns; a new human 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/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py index 8f211ef67db..d2a46fc316d 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -1,12 +1,13 @@ """Keeps a mid-stream continuation on deployments whose model supports assistant prefill and whose own params leave -the output unconstrained.""" +the output unconstrained, and shapes a Responses API continuation for the deployment that serves it.""" -from collections.abc import AsyncIterable, AsyncIterator, Mapping, Sequence +import json +from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping, Sequence from dataclasses import dataclass from itertools import accumulate from os.path import commonprefix from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar from pydantic import TypeAdapter, ValidationError @@ -14,7 +15,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.rules import Rules from litellm.router_utils.common_utils import MID_STREAM_CONTINUATION_KWARG -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, OutputTextDeltaEvent from litellm.types.utils import ModelResponse, ModelResponseStream, Usage from litellm.utils import supports_assistant_prefill @@ -28,7 +29,9 @@ MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker() _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) _UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"}) -_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens") +_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens", "max_output_tokens") +_QUOTED_ENDING_CHARS: Final = 100 +RESPONSES_CONTINUATION_KWARG: Final = "_responses_continuation" def _str_keyed(value: object) -> Mapping[str, object]: @@ -38,32 +41,55 @@ def _str_keyed(value: object) -> Mapping[str, object]: return _NO_FIELDS -def deployment_supports_prefill(deployment: object) -> bool: +def _underlying_model(deployment: object) -> str | None: deployment_map: Final = _str_keyed(deployment) model_info: Final = _str_keyed(deployment_map.get("model_info")) - declared: Final = model_info.get("supports_assistant_prefill") - if isinstance(declared, bool): - return declared litellm_params: Final = _str_keyed(deployment_map.get("litellm_params")) candidates: Final = (model_info.get("base_model"), litellm_params.get("base_model"), litellm_params.get("model")) - model: Final = next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None) + return next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None) + + +def deployment_supports_prefill(deployment: object) -> bool: + declared: Final = _str_keyed(_str_keyed(deployment).get("model_info")).get("supports_assistant_prefill") + if isinstance(declared, bool): + return declared + model: Final = _underlying_model(deployment) return model is not None and supports_assistant_prefill(model=model) -def constrains_continuation(params: Mapping[str, object]) -> bool: +def _forces_tool_call(tool_choice: object) -> bool: + return tool_choice == "required" or isinstance(tool_choice, Mapping) + + +def _forces_responses_tool_call(tool_choice: object) -> bool: + choice: Final = _str_keyed(tool_choice) + leaves_text_open: Final = choice.get("type") == "allowed_tools" and choice.get("mode") == "auto" + return _forces_tool_call(tool_choice) and not leaves_text_open + + +def _constrains(params: Mapping[str, object], forces_tool_call: Callable[[object], bool]) -> bool: sent: Final = {**params, **_str_keyed(params.get("extra_body"))} choice_count: Final = sent.get("n") response_format: Final = sent.get("response_format") - tool_choice: Final = sent.get("tool_choice") + text_format: Final = _str_keyed(sent.get("text")).get("format") return ( (isinstance(choice_count, int) and choice_count > 1) or (response_format is not None and response_format != _UNCONSTRAINED_RESPONSE_FORMAT) - or tool_choice == "required" - or isinstance(tool_choice, Mapping) + or (text_format is not None and text_format != _UNCONSTRAINED_RESPONSE_FORMAT) + or sent.get("text_format") is not None + or forces_tool_call(sent.get("tool_choice")) or sent.get("merge_reasoning_content_in_choices") is True ) +def constrains_continuation(params: Mapping[str, object]) -> bool: + return _constrains(params, _forces_tool_call) + + +def constrains_responses_continuation(params: Mapping[str, object]) -> bool: + return _constrains(params, _forces_responses_tool_call) + + def deployment_can_continue(deployment: object) -> bool: litellm_params: Final = _str_keyed(_str_keyed(deployment).get("litellm_params")) return deployment_supports_prefill(deployment) and not constrains_continuation(litellm_params) @@ -99,6 +125,7 @@ def _delta_text(chunk: ModelResponseStream) -> str: _CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) +_INPUT_ITEMS_ADAPTER: Final = TypeAdapter(list[object]) def _prefill_text(prefill_message: Mapping[str, object]) -> str: @@ -114,7 +141,7 @@ def _prefill_text(prefill_message: Mapping[str, object]) -> str: ) -def _resendable_whitespace(prefill: str) -> str: +def resendable_whitespace(prefill: str) -> str: return prefill[len(prefill.rstrip()) :] if prefill.strip() else "" @@ -124,7 +151,7 @@ def continued_text(prefill_message: Mapping[str, object], generated_content: str return None echo: Final = len(prefill) if prefill and generated_content.startswith(prefill) else 0 after_echo: Final = generated_content[echo:] - return after_echo[len(commonprefix((after_echo, _resendable_whitespace(prefill)))) :] + return after_echo[len(commonprefix((after_echo, resendable_whitespace(prefill)))) :] def _reduced_ceilings(params: Mapping[str, object], emitted_tokens: int) -> Mapping[str, int]: @@ -211,7 +238,7 @@ class ContinuationOutput: self._rules: Final = Rules() self._held: tuple[ModelResponseStream, ...] = () self._echo_settled = not self._prefill_text - self._seam = _resendable_whitespace(self._prefill_text) + self._seam = resendable_whitespace(self._prefill_text) async def splice(self, stream: AsyncIterable[object]) -> AsyncIterator[object]: async for item in stream: @@ -262,3 +289,91 @@ class ContinuationOutput: self._text += _delta_text(chunk) self._rules.post_call_rules(input=self._text, model=self._model) return chunks + + +@dataclass(frozen=True, slots=True) +class ResponsesContinuation: + prior_input: object + emitted_text: str + + def extended(self, generated_content: str) -> "ResponsesContinuation": + return ResponsesContinuation(self.prior_input, f"{self.emitted_text}{generated_content}") + + +def deployment_takes_responses_prefill(deployment: object) -> bool: + """Claude takes a trailing assistant turn as prefill with no `prefix` flag, which a Responses input item cannot + carry.""" + model: Final = _underlying_model(deployment) + return model is not None and "claude" in model.lower() and deployment_supports_prefill(deployment) + + +def _continuation_request(emitted_text: str) -> Mapping[str, object]: + seam: Final = ( + "It already ends with whitespace, so do not add another space before your first word." + if emitted_text[-1:].isspace() + else "It ends without trailing whitespace, so begin with a space if your first word starts a new word." + ) + ending: Final = json.dumps(emitted_text[-_QUOTED_ENDING_CHARS:], ensure_ascii=False) + instruction: Final = ( + f"The previous assistant response was interrupted mid-stream. It ended with {ending}. Continue exactly where " + "it stopped, so your first characters are the ones that would have come right after it, and do not repeat " + f"any of its content. {seam} Your response must read as a seamless continuation." + ) + return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": instruction}]} + + +def _input_items(request_input: object) -> tuple[object, ...]: + if isinstance(request_input, str): + return ({"type": "message", "role": "user", "content": [{"type": "input_text", "text": request_input}]},) + try: + return tuple(_INPUT_ITEMS_ADAPTER.validate_python(request_input)) + except ValidationError: + return () + + +def responses_continuation_input(continuation: ResponsesContinuation) -> Sequence[object]: + """The prior input, the partial response as assistant history, and a user request to continue it.""" + partial: Final = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": continuation.emitted_text}], + } + return [*_input_items(continuation.prior_input), partial, _continuation_request(continuation.emitted_text)] + + +def responses_continuation_input_for(request_kwargs: Mapping[str, object], deployment: Mapping[str, object]) -> object: + """The continuation input as the deployment receives it, without the request to continue where the deployment + takes the partial response as prefill. It reads the input the router already stripped of reasoning this + deployment cannot decrypt, never the continuation's prior input.""" + request_input: Final = request_kwargs.get("input") + continuation: Final = request_kwargs.get(RESPONSES_CONTINUATION_KWARG) + if not isinstance(continuation, ResponsesContinuation): + return request_input + items: Final = _input_items(request_input) + if items[-1:] != (_continuation_request(continuation.emitted_text),) or not deployment_takes_responses_prefill( + deployment + ): + return request_input + return list(items[:-1]) + + +_StreamItem = TypeVar("_StreamItem") + + +class ResponsesContinuationSeam: + """Drops whitespace a fallback sends again at the start of its text when the client already has it, since + Anthropic strips a prefill's trailing whitespace and the model then writes it back.""" + + def __init__(self, emitted_text: str) -> None: + self._seam = resendable_whitespace(emitted_text) + + async def splice(self, stream: AsyncIterable[_StreamItem]) -> AsyncIterator[_StreamItem]: + async for item in stream: + yield self._past(item) + + def _past(self, item: _StreamItem) -> _StreamItem: + if not self._seam or not isinstance(item, OutputTextDeltaEvent): + return item + matched: Final = len(commonprefix((item.delta, self._seam))) + self._seam = self._seam[matched:] if matched == len(item.delta) else "" + return item.model_copy(update={"delta": item.delta[matched:]}) if matched else item 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/llms/mistral.py b/litellm/types/llms/mistral.py index e64aa9c13ab..cd61f32f430 100644 --- a/litellm/types/llms/mistral.py +++ b/litellm/types/llms/mistral.py @@ -1,6 +1,6 @@ from typing import Literal -from typing_extensions import TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict class FunctionCall(TypedDict): @@ -22,3 +22,4 @@ class MistralTextBlock(TypedDict): class MistralThinkingBlock(TypedDict): type: Literal["thinking"] thinking: list[MistralTextBlock] + closed: NotRequired[ReadOnly[bool]] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 4912f9f1375..8e1db87c16f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1342,6 +1342,10 @@ class ResponsesAPIStreamOptions(TypedDict, total=False): include_obfuscation: bool +class ResponsesAPIModeration(TypedDict, total=False): + model: Required[ReadOnly[str]] + + class ResponsesAPIOptionalRequestParams(TypedDict, total=False): """TypedDict for Optional parameters supported by the responses API.""" @@ -1369,6 +1373,7 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False): prompt_cache_key: str | None prompt_cache_retention: str | None prompt_cache_options: ReadOnly[PromptCacheOptions | None] + moderation: ReadOnly[ResponsesAPIModeration | None] stream_options: ResponsesAPIStreamOptions | None top_logprobs: int | None partial_images: int | None # Number of partial images to generate (1-3) for streaming image generation diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 055c18f2e36..b21bf0e14d4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -225,6 +225,8 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_xhigh_reasoning_effort: bool | None supports_max_reasoning_effort: bool | None reasoning_effort_levels: ReadOnly[Sequence[str] | None] + reasoning_modes: ReadOnly[Sequence[str] | None] + reasoning_contexts: ReadOnly[Sequence[str] | None] default_reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None] supports_output_config: bool | None supports_image_size: bool | None @@ -4586,6 +4588,7 @@ class LlmProviders(str, Enum): PRISM = "prism" OPENINFER = "openinfer" INFERENCEAPIS = "inferenceapis" + ZEROGPU = "zerogpu" DARKBLOOM = "darkbloom" META = "meta" SAIL = "sail" diff --git a/litellm/utils.py b/litellm/utils.py index 1e095966a43..a231644addf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7000,6 +7000,8 @@ def get_model_info_helper( supports_xhigh_reasoning_effort=_model_info.get("supports_xhigh_reasoning_effort", None), supports_max_reasoning_effort=_model_info.get("supports_max_reasoning_effort", None), reasoning_effort_levels=_model_info.get("reasoning_effort_levels", None), + reasoning_modes=_model_info.get("reasoning_modes", None), + reasoning_contexts=_model_info.get("reasoning_contexts", None), default_reasoning_effort=_model_info.get("default_reasoning_effort", None), bedrock_output_config_effort_ceiling=_model_info.get("bedrock_output_config_effort_ceiling", None), bedrock_converse_supports_strict_tools=_model_info.get("bedrock_converse_supports_strict_tools", None), @@ -9612,18 +9614,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 8207c3122f9..3e0e86012a4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4039,6 +4039,7 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure_ai/gpt-oss-120b": { + "deprecation_date": "2027-09-15", "input_cost_per_token": 1.5e-07, "output_cost_per_token": 6e-07, "litellm_provider": "azure_ai", @@ -4243,6 +4244,54 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, @@ -4904,6 +4953,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -4922,6 +4976,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -5009,6 +5068,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5048,6 +5112,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5113,6 +5182,11 @@ "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.1-codex": { @@ -5188,6 +5262,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5344,6 +5423,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, @@ -5428,6 +5512,11 @@ "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/global/gpt-5.1-codex": { @@ -5680,6 +5769,11 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5717,6 +5811,11 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5754,6 +5853,11 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -5791,6 +5895,11 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6203,6 +6312,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -6221,6 +6335,11 @@ "mode": "chat", "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -6710,6 +6829,11 @@ "default_reasoning_effort": "none", "input_cost_per_token_batches": 6.25e-07, "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_minimal_reasoning_effort": true, "cache_read_input_token_cost_batches": 6.25e-08 @@ -6797,6 +6921,11 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6835,6 +6964,11 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6906,6 +7040,11 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_batches": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6944,6 +7083,11 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_batches": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -6981,6 +7125,11 @@ "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -7016,6 +7165,11 @@ "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -7051,6 +7205,11 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/responses" @@ -7112,6 +7271,11 @@ "input_cost_per_token_priority": 2.5e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/gpt-5-chat": { @@ -7459,6 +7623,11 @@ "output_cost_per_token": 1.4e-05, "output_cost_per_token_batches": 7e-06, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -7498,6 +7667,11 @@ "output_cost_per_token": 1.4e-05, "output_cost_per_token_batches": 7e-06, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -7567,6 +7741,11 @@ "mode": "responses", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/responses" @@ -7599,9 +7778,9 @@ "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -7637,9 +7816,9 @@ "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -7996,9 +8175,9 @@ "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8043,9 +8222,9 @@ "output_cost_per_token_batches": 9e-05, "output_cost_per_token_flex": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8095,9 +8274,9 @@ "output_cost_per_token_priority": 4e-05, "output_cost_per_token_above_272k_tokens_priority": 6e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -8157,9 +8336,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8220,9 +8399,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_flex": 1e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8283,9 +8462,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_flex": 6e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8346,9 +8525,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_flex": 6e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8409,9 +8588,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8472,9 +8651,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -8519,9 +8698,9 @@ "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8567,9 +8746,9 @@ "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -8615,9 +8794,9 @@ "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -8663,9 +8842,9 @@ "output_cost_per_token": 5e-07, "output_cost_per_token_above_272k_tokens": 7.5e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -8718,7 +8897,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8770,7 +8949,7 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8800,6 +8979,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -8815,9 +9041,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8863,9 +9089,9 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8896,6 +9122,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/us/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, + "cache_read_input_token_cost": 5.5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-07, + "input_cost_per_token": 1.1e-06, + "input_cost_per_token_above_272k_tokens": 2.2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "output_cost_per_token_above_272k_tokens": 8.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/us/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, @@ -8911,9 +9184,9 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8944,6 +9217,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/eu/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_above_272k_tokens": 2.4e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_272k_tokens": 9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/eu/gpt-6.1-sol": { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, @@ -8959,9 +9279,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -8992,6 +9312,53 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/apac/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.2e-07, + "input_cost_per_token": 1.2e-06, + "input_cost_per_token_above_272k_tokens": 2.4e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_272k_tokens": 9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/gpt-6-1-terra-now-generally-available-in-microsoft-foundry-build-agents-that-see/4567769", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/apac/gpt-6.1-sol": { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, @@ -9007,9 +9374,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", "supported_endpoints": [ @@ -9134,9 +9501,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9188,9 +9555,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9243,9 +9610,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, "output_cost_per_token_priority": 2.64e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9298,9 +9665,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-06, "output_cost_per_token_priority": 2.64e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9345,9 +9712,9 @@ "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9393,9 +9760,9 @@ "output_cost_per_token": 5.5e-07, "output_cost_per_token_above_272k_tokens": 8.25e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -9448,7 +9815,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_priority": 3.3e-05, "output_cost_per_token_priority": 2.2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": {"search_context_size_high":0.014,"search_context_size_low":0.014,"search_context_size_medium":0.014}, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -9536,9 +9903,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9590,9 +9957,9 @@ "output_cost_per_token_above_272k_tokens_priority": 6.6e-05, "output_cost_per_token_priority": 4.4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9645,9 +10012,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, "output_cost_per_token_priority": 2.64e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9700,9 +10067,9 @@ "output_cost_per_token_above_272k_tokens_priority": 3.96e-06, "output_cost_per_token_priority": 2.64e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9760,9 +10127,9 @@ "output_cost_per_token_above_272k_tokens_priority": 9e-05, "output_cost_per_token_batches": 1.5e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9814,9 +10181,9 @@ "output_cost_per_token_above_272k_tokens_batches": 2.475e-05, "output_cost_per_token_batches": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9868,9 +10235,9 @@ "output_cost_per_token_above_272k_tokens_batches": 2.475e-05, "output_cost_per_token_batches": 1.65e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -9921,9 +10288,9 @@ "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -9981,9 +10348,9 @@ "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10034,9 +10401,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10086,9 +10453,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10139,9 +10506,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10191,9 +10558,9 @@ "max_output_tokens": 128000, "mode": "chat", "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10239,9 +10606,9 @@ "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/batch", @@ -10282,9 +10649,9 @@ "output_cost_per_token": 0.00018, "output_cost_per_token_above_272k_tokens": 0.00027, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/batch", @@ -10322,9 +10689,9 @@ "mode": "chat", "output_cost_per_token": 4.5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10372,9 +10739,9 @@ "mode": "chat", "output_cost_per_token": 4.5e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10423,9 +10790,9 @@ "mode": "chat", "output_cost_per_token": 1.25e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10470,9 +10837,9 @@ "mode": "chat", "output_cost_per_token": 1.25e-06, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "supported_endpoints": [ "/v1/chat/completions", @@ -10934,6 +11301,11 @@ "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -10965,6 +11337,11 @@ "mode": "chat", "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -10996,9 +11373,9 @@ "mode": "responses", "output_cost_per_token": 4e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ @@ -11065,6 +11442,11 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/batch", @@ -11095,6 +11477,11 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/batch", @@ -11124,6 +11511,11 @@ "max_output_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -11155,6 +11547,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "output_cost_per_token_batches": 2.2e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": false, @@ -11268,6 +11665,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11305,6 +11707,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11428,6 +11835,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -11446,6 +11858,11 @@ "max_output_tokens": 16384, "mode": "chat", "output_cost_per_token": 6.6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -11533,6 +11950,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11572,6 +11994,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11608,6 +12035,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11673,6 +12105,11 @@ "input_cost_per_token_priority": 2.75e-06, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.1-codex": { @@ -11800,6 +12237,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11832,6 +12274,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ "/v1/chat/completions", @@ -11898,6 +12345,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": false, @@ -11918,6 +12370,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": false, @@ -11935,6 +12392,7 @@ "output_cost_per_second": 0.0001 }, "azure_ai/Cohere-embed-v3-english": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 512, @@ -11945,6 +12403,7 @@ "supports_embedding_image_input": true }, "azure_ai/Cohere-embed-v3-multilingual": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 512, @@ -11955,6 +12414,7 @@ "supports_embedding_image_input": true }, "azure_ai/Codestral-2501": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 3e-07, "litellm_provider": "azure_ai", "max_input_tokens": 256000, @@ -11974,6 +12434,7 @@ "supports_native_streaming": true }, "azure_ai/FLUX-1.1-pro": { + "deprecation_date": "2027-01-05", "litellm_provider": "azure_ai", "max_input_tokens": 5000, "mode": "image_generation", @@ -11984,6 +12445,7 @@ ] }, "azure_ai/FLUX.1-Kontext-pro": { + "deprecation_date": "2027-01-05", "litellm_provider": "azure_ai", "max_input_tokens": 5000, "mode": "image_generation", @@ -11994,6 +12456,7 @@ ] }, "azure_ai/flux.2-pro": { + "deprecation_date": "2027-01-20", "litellm_provider": "azure_ai", "max_input_tokens": 32000, "mode": "image_generation", @@ -12004,6 +12467,7 @@ ] }, "azure_ai/FLUX.2-flex": { + "deprecation_date": "2027-01-20", "input_cost_per_pixel": 5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 32000, @@ -12368,6 +12832,7 @@ "supports_tool_choice": true }, "azure_ai/Llama-3.3-70B-Instruct": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 7.1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12379,6 +12844,7 @@ "supports_tool_choice": true }, "azure_ai/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "deprecation_date": "2026-11-19", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -12391,6 +12857,7 @@ "supports_vision": true }, "azure_ai/Llama-4-Scout-17B-16E-Instruct": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 10000000, @@ -12624,6 +13091,7 @@ ] }, "azure_ai/cohere-command-a": { + "deprecation_date": "2027-02-15", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12696,6 +13164,7 @@ "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models" }, "azure_ai/cohere-rerank-v4.0-pro": { + "deprecation_date": "2027-09-13", "input_cost_per_query": 0.0025, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -12706,6 +13175,7 @@ "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-cohere-rerank-4-0-in-microsoft-foundry/4477076" }, "azure_ai/cohere-rerank-v4.0-fast": { + "deprecation_date": "2027-09-13", "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -12716,6 +13186,7 @@ "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-cohere-rerank-4-0-in-microsoft-foundry/4477076" }, "azure_ai/deepseek-v3.2": { + "deprecation_date": "2027-11-01", "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12730,6 +13201,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3.2-speciale": { + "deprecation_date": "2027-11-01", "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12800,6 +13272,7 @@ "supports_tool_choice": true }, "azure_ai/embed-v-4-0": { + "deprecation_date": "2027-06-01", "input_cost_per_token": 1.2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12818,6 +13291,7 @@ "input_cost_per_image_token": 4.7e-07 }, "azure_ai/grok-4": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 262000, @@ -12906,6 +13380,7 @@ "supports_web_search": true }, "azure_ai/grok-4-1-fast-non-reasoning": { + "deprecation_date": "2027-05-19", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -12919,6 +13394,7 @@ "supports_web_search": true }, "azure_ai/grok-4-1-fast-reasoning": { + "deprecation_date": "2027-05-19", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -12933,6 +13409,7 @@ "supports_web_search": true }, "azure_ai/grok-code-fast-1": { + "deprecation_date": "2027-01-31", "input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "max_input_tokens": 256000, @@ -13100,7 +13577,7 @@ "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { - "deprecation_date": "2026-12-09", + "deprecation_date": "2027-02-15", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -14619,11 +15096,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14638,6 +15116,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -14647,11 +15126,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14665,7 +15145,8 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -14787,11 +15268,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14806,6 +15288,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -14815,11 +15298,12 @@ "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, "cache_read_input_token_cost": 3.6e-07, + "deprecation_date": "2027-04-08", "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, - "max_output_tokens": 8192, + "max_output_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.8e-05, "supports_assistant_prefill": true, @@ -14833,7 +15317,8 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -15197,6 +15682,7 @@ "supports_tool_choice": true }, "cerebras/llama-3.3-70b": { + "deprecation_date": "2026-02-16", "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15207,6 +15693,7 @@ "supports_tool_choice": true }, "cerebras/llama3.1-70b": { + "deprecation_date": "2025-01-17", "input_cost_per_token": 6e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15217,6 +15704,7 @@ "supports_tool_choice": true }, "cerebras/llama3.1-8b": { + "deprecation_date": "2026-05-27", "input_cost_per_token": 1e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15241,6 +15729,7 @@ "supports_tool_choice": true }, "cerebras/qwen-3-32b": { + "deprecation_date": "2026-02-16", "input_cost_per_token": 4e-07, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -15552,6 +16041,7 @@ "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "deprecation_date": "2027-04-08", "litellm_provider": "bedrock", "supports_tool_search": true, "max_input_tokens": 200000, @@ -15570,7 +16060,8 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, "input_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_batches": 7.5e-06 + "output_cost_per_token_batches": 7.5e-06, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "claude-opus-4-5-20251101": { "supports_web_search": true, @@ -15764,7 +16255,6 @@ "us": 1.1 }, "supports_output_config": true, - "supports_speed": true, "prompt_cache_min_tokens": 2048, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, @@ -15807,7 +16297,6 @@ "us": 1.1 }, "supports_output_config": true, - "supports_speed": true, "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { @@ -20520,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, @@ -20532,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, @@ -20548,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, @@ -20560,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, @@ -20578,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, @@ -20601,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, @@ -20611,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, @@ -20624,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, @@ -20634,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, @@ -20647,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, @@ -20657,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, @@ -20671,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, @@ -20681,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, @@ -20695,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, @@ -20706,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, @@ -20721,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, @@ -20732,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, @@ -20748,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, @@ -20759,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, @@ -20855,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, @@ -20865,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, @@ -20878,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, @@ -20888,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, @@ -20901,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, @@ -21057,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", @@ -21087,13 +21592,19 @@ "supports_video_input": true }, "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", @@ -21107,13 +21618,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", @@ -21129,7 +21646,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, @@ -21139,12 +21658,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", @@ -21154,13 +21670,20 @@ "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 }, "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", @@ -21209,13 +21732,17 @@ "supports_tool_choice": true }, "databricks/databricks-gemini-3-8-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": "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_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21233,13 +21760,17 @@ "supports_vision": true }, "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": "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_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21258,9 +21789,9 @@ }, "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_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, @@ -21268,8 +21799,8 @@ "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, + "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21530,7 +22061,9 @@ "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, @@ -21540,6 +22073,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": [ @@ -21613,8 +22147,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, @@ -21624,6 +22161,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": [ @@ -21640,10 +22178,13 @@ "supports_vision": true }, "databricks/databricks-gpt-5-6-terra": { - "cache_creation_input_token_cost": 3.12501e-06, - "cache_read_input_token_cost": 2.4997e-07, - "input_cost_per_token": 2.49998e-06, - "input_dbu_cost_per_token": 3.5714e-05, + "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, @@ -21651,8 +22192,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": "chat", - "output_cost_per_token": 1.500002e-05, - "output_dbu_cost_per_token": 0.000214286, + "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": [ "text", @@ -21668,10 +22210,13 @@ "supports_vision": true }, "databricks/databricks-gpt-5-6-luna": { - "cache_creation_input_token_cost": 1.24999e-06, - "cache_read_input_token_cost": 1.0003e-07, - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, + "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, "max_output_tokens": 128000, @@ -21679,8 +22224,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": "chat", - "output_cost_per_token": 5.99998e-06, - "output_dbu_cost_per_token": 8.5714e-05, + "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": [ "text", @@ -21698,7 +22244,9 @@ "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, @@ -21708,6 +22256,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": [ @@ -21727,7 +22276,8 @@ "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, @@ -21735,8 +22285,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", @@ -22055,6 +22606,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, @@ -34478,6 +35048,7 @@ "source": "https://developers.openai.com/api/docs/pricing" }, "gpt-4o-mini-tts": { + "deprecation_date": "2027-01-06", "input_cost_per_token": 6e-07, "litellm_provider": "openai", "mode": "audio_speech", @@ -35095,6 +35666,7 @@ "cache_read_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.25e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, @@ -35107,6 +35679,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, @@ -35184,6 +35763,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, @@ -35351,6 +35937,7 @@ "cache_read_input_token_cost_batches": 6.25e-08, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2027-04-01", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -35359,6 +35946,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, @@ -35410,6 +36004,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, @@ -35525,6 +36126,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, @@ -35574,6 +36182,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, @@ -35615,6 +36230,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, @@ -35650,15 +36272,15 @@ "gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_12hr": 2.5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, - "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, "cache_creation_input_token_cost_above_12hr_priority": 5e-05, - "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -35702,6 +36324,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": { @@ -35786,6 +36417,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": { @@ -35828,11 +36468,11 @@ "gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_12hr": 2.5e-07, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, - "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, @@ -35878,6 +36518,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": { @@ -35961,6 +36610,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": { @@ -36042,6 +36700,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": { @@ -36125,6 +36792,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": { @@ -36207,6 +36883,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": { @@ -36444,6 +37129,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", @@ -36491,6 +37184,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": { @@ -36553,6 +37254,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": { @@ -36607,6 +37316,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": { @@ -36658,6 +37375,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": { @@ -36717,6 +37442,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": [ @@ -36774,6 +37507,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": [ @@ -36823,6 +37564,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": { @@ -36875,6 +37624,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": { @@ -36929,6 +37686,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": { @@ -36983,6 +37748,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": { @@ -37034,6 +37807,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": { @@ -37075,6 +37856,7 @@ "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, + "deprecation_date": "2027-04-01", "input_cost_per_token": 2e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, @@ -37085,6 +37867,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": { @@ -37122,6 +37912,7 @@ "supports_minimal_reasoning_effort": false }, "gpt-5-pro": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -37130,6 +37921,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, @@ -37172,6 +37970,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, @@ -37222,6 +38027,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, @@ -37332,6 +38144,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, @@ -37401,6 +38220,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -37413,6 +38233,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, @@ -37464,6 +38291,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, @@ -37501,6 +38335,7 @@ "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_batches": 2.5e-09, "cache_read_input_token_cost_flex": 2.5e-09, + "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, "input_cost_per_token_batches": 2.5e-08, "input_cost_per_token_flex": 2.5e-08, @@ -37511,6 +38346,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, @@ -37559,6 +38401,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, @@ -38210,6 +39059,7 @@ }, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -38230,6 +39080,7 @@ }, "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -38249,6 +39100,7 @@ "search_context_size_medium": 0.005 }, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -40996,9 +41848,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 3e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/devstral-small-2-25-12", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41010,9 +41867,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41024,9 +41886,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41331,6 +42198,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, @@ -41547,13 +42430,13 @@ "mode": "embedding" }, "mistral/mistral-large-4": { - "cache_read_input_token_cost": 6.8e-08, - "input_cost_per_token": 6.8e-07, + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.36e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 2.09e-06, + "output_cost_per_token": 4.18e-06, "reasoning_effort_levels": [ "none", "high" @@ -41568,13 +42451,13 @@ "supports_vision": true }, "mistral/mistral-large-4-0": { - "cache_read_input_token_cost": 6.8e-08, - "input_cost_per_token": 6.8e-07, + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.36e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 2.09e-06, + "output_cost_per_token": 4.18e-06, "reasoning_effort_levels": [ "none", "high" @@ -41721,8 +42604,13 @@ "max_output_tokens": 8191, "mode": "chat", "output_cost_per_token": 3e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -41921,6 +42809,7 @@ }, "moonshot/kimi-k2.5": { "cache_read_input_token_cost": 1e-07, + "deprecation_date": "2026-08-31", "input_cost_per_token": 6e-07, "litellm_provider": "moonshot", "max_input_tokens": 262144, @@ -41951,7 +42840,31 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3.1-pro": { + "cache_creation_input_token_cost": 2e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://platform.kimi.ai/docs/pricing/chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, "moonshot/kimi-k3": { + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "moonshot", @@ -41966,6 +42879,7 @@ ], "source": "https://platform.kimi.ai/docs/pricing/chat-k3", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -41973,6 +42887,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-128k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -41985,6 +42900,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-128k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -41998,6 +42914,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-32k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 1e-06, "litellm_provider": "moonshot", "max_input_tokens": 32768, @@ -42010,6 +42927,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-32k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 1e-06, "litellm_provider": "moonshot", "max_input_tokens": 32768, @@ -42023,6 +42941,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-8k": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-07, "litellm_provider": "moonshot", "max_input_tokens": 8192, @@ -42035,6 +42954,7 @@ "supports_tool_choice": true }, "moonshot/moonshot-v1-8k-vision-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-07, "litellm_provider": "moonshot", "max_input_tokens": 8192, @@ -42048,6 +42968,7 @@ "supports_vision": true }, "moonshot/moonshot-v1-auto": { + "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", "max_input_tokens": 131072, @@ -43240,6 +44161,7 @@ "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -43252,6 +44174,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, @@ -43360,6 +44289,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, @@ -43430,6 +44366,7 @@ "supports_vision": false }, "o3-pro": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -43438,6 +44375,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, @@ -43475,6 +44419,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, @@ -47290,13 +48241,13 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 4.5e-06, + "output_cost_per_token": 7.5e-06, "source": "https://www.together.ai/pricing", "supports_prompt_caching": true }, @@ -47806,6 +48757,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 9e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, + "deprecation_date": "2027-04-08", "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, @@ -47824,7 +48776,7 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md", "token_count_model": { "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_provider": "bedrock" @@ -52901,15 +53853,21 @@ "xai/grok-4.7": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "cache_read_input_token_cost_above_200k_tokens_batches": 8e-07, + "cache_read_input_token_cost_batches": 4e-07, "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 3.2e-06, + "input_cost_per_token_batches": 1.6e-06, "litellm_provider": "xai", "max_input_tokens": 500000, "max_output_tokens": 500000, "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, + "output_cost_per_token_above_200k_tokens_batches": 9.6e-06, + "output_cost_per_token_batches": 4.8e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -56990,6 +57948,7 @@ ] }, "gpt-4o-mini-transcribe-2025-12-15": { + "deprecation_date": "2027-02-26", "input_cost_per_audio_token": 1.25e-06, "input_cost_per_second": 5e-05, "input_cost_per_token": 1.25e-06, @@ -64527,6 +65486,7 @@ ] }, "xai/grok-imagine-image-pro": { + "deprecation_date": "2026-05-15", "input_cost_per_image": 0.05, "litellm_provider": "xai", "mode": "image_generation", @@ -64933,9 +65893,14 @@ "max_output_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://mistral.ai/news/devstral-2-vibe-cli", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, @@ -65279,6 +66244,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", @@ -65364,6 +66503,7 @@ ] }, "scaleway/glm-5.2": { + "deprecation_date": "2026-11-16", "input_cost_per_token": 1.8e-06, "litellm_provider": "scaleway", "max_input_tokens": 256000, @@ -67922,6 +69062,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-mini": { @@ -67936,6 +69081,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { @@ -67974,6 +69124,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5-codex": { @@ -67998,6 +69153,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5-nano": { @@ -68010,6 +69170,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5-pro": { @@ -68020,6 +69185,11 @@ "mode": "responses", "output_cost_per_token": 0.000132, "output_cost_per_token_batches": 6.6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.1-codex-max": { @@ -68044,6 +69214,11 @@ "output_cost_per_token": 1.54e-05, "output_cost_per_token_batches": 7.7e-06, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.2-codex": { @@ -68062,6 +69237,11 @@ "mode": "responses", "output_cost_per_token": 0.0001848, "output_cost_per_token_batches": 9.24e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.3-codex": { @@ -68074,6 +69254,11 @@ "mode": "responses", "output_cost_per_token": 1.54e-05, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-mini": { @@ -68089,6 +69274,11 @@ "output_cost_per_token": 4.95e-06, "output_cost_per_token_batches": 2.475e-06, "output_cost_per_token_priority": 9.9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-nano": { @@ -68101,6 +69291,11 @@ "mode": "chat", "output_cost_per_token": 1.375e-06, "output_cost_per_token_batches": 6.875e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-pro": { @@ -68115,6 +69310,11 @@ "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-astra": { @@ -68131,6 +69331,11 @@ "mode": "chat", "output_cost_per_token": 6e-05, "output_cost_per_token_above_272k_tokens": 9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supported_endpoints": [ "/v1/chat/completions", @@ -68175,9 +69380,9 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_above_272k_tokens": 9e-07, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -68223,9 +69428,9 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ @@ -68277,6 +69482,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/o3-2025-04-16": { @@ -68290,6 +69500,11 @@ "mode": "chat", "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/o3-deep-research": { @@ -68312,6 +69527,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/o4-mini-2025-04-16": { @@ -68325,6 +69545,11 @@ "mode": "chat", "output_cost_per_token": 4.84e-06, "output_cost_per_token_batches": 2.42e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/text-embedding-3-large": { @@ -68418,6 +69643,11 @@ "output_cost_per_token": 8.8e-06, "output_cost_per_token_batches": 4.4e-06, "output_cost_per_token_priority": 1.54e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-mini": { @@ -68432,6 +69662,11 @@ "output_cost_per_token": 1.76e-06, "output_cost_per_token_batches": 8.8e-07, "output_cost_per_token_priority": 3.08e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { @@ -68470,6 +69705,11 @@ "output_cost_per_token": 1.1e-05, "output_cost_per_token_batches": 5.5e-06, "output_cost_per_token_priority": 2.2e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5-codex": { @@ -68494,6 +69734,11 @@ "output_cost_per_token": 2.2e-06, "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_priority": 3.96e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5-nano": { @@ -68506,6 +69751,11 @@ "mode": "chat", "output_cost_per_token": 4.4e-07, "output_cost_per_token_batches": 2.2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5-pro": { @@ -68516,6 +69766,11 @@ "mode": "responses", "output_cost_per_token": 0.000132, "output_cost_per_token_batches": 6.6e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.1-codex-max": { @@ -68540,6 +69795,11 @@ "output_cost_per_token": 1.54e-05, "output_cost_per_token_batches": 7.7e-06, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.2-codex": { @@ -68558,6 +69818,11 @@ "mode": "responses", "output_cost_per_token": 0.0001848, "output_cost_per_token_batches": 9.24e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.3-codex": { @@ -68570,6 +69835,11 @@ "mode": "responses", "output_cost_per_token": 1.54e-05, "output_cost_per_token_priority": 3.08e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-mini": { @@ -68585,6 +69855,11 @@ "output_cost_per_token": 4.95e-06, "output_cost_per_token_batches": 2.475e-06, "output_cost_per_token_priority": 9.9e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-nano": { @@ -68597,6 +69872,11 @@ "mode": "chat", "output_cost_per_token": 1.375e-06, "output_cost_per_token_batches": 6.875e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-pro": { @@ -68611,6 +69891,11 @@ "output_cost_per_token_above_272k_tokens": 0.000297, "output_cost_per_token_above_272k_tokens_batches": 0.0001485, "output_cost_per_token_batches": 9.9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/o1-mini": { @@ -68657,7 +69942,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2026-10-31", + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -71002,6 +72287,7 @@ "xai/grok-code-fast-1": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "deprecation_date": "2026-05-15", "input_cost_per_image_token": 1e-06, "input_cost_per_token": 1e-06, "input_cost_per_token_above_200k_tokens": 2e-06, @@ -71898,7 +73184,7 @@ "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, - "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost": 1e-07, "deprecation_date": "2027-09-29", "input_cost_per_token": 2e-06, "litellm_provider": "azure_ai", @@ -72505,15 +73791,15 @@ "gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, - "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, - "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, @@ -72557,6 +73843,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": { @@ -74474,14 +75769,14 @@ "supports_vision": true }, "wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2.5e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 943718, "mode": "chat", - "output_cost_per_token": 7.5e-07, - "source": "https://openrouter.ai/api/v1/models/nvidia/nemotron-3.5-super-vl/endpoints", + "output_cost_per_token": 1.5e-06, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -76056,6 +77351,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}, @@ -76188,7 +77492,7 @@ "supports_xhigh_reasoning_effort": true }, "baseten/decision-models/baseten-elephant-9b-latest": { - "input_cost_per_token": 0.0, + "input_cost_per_token": 3.5e-08, "litellm_provider": "baseten", "max_input_tokens": 16384, "mode": "decisions", @@ -76197,7 +77501,7 @@ "supports_vision": true }, "baseten/decision-models/baseten-mammoth-27b-latest": { - "input_cost_per_token": 0.0, + "input_cost_per_token": 1e-07, "litellm_provider": "baseten", "max_input_tokens": 16384, "mode": "decisions", @@ -76229,13 +77533,13 @@ "supports_vision": true }, "refuseless/glm-5.3-abliterated": { - "cache_read_input_token_cost": 2.5e-07, - "input_cost_per_token": 1.25e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.25e-06, + "output_cost_per_token": 2e-06, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -76250,13 +77554,13 @@ ] }, "refuseless/glm-5.3-flash-abliterated": { - "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_token": 2.5e-07, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 1.2e-07, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8e-07, + "output_cost_per_token": 5e-07, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -76271,13 +77575,13 @@ ] }, "refuseless/deepseek-v4.1-flash-abliterated": { - "cache_read_input_token_cost": 4.5e-08, - "input_cost_per_token": 2.2e-07, + "cache_read_input_token_cost": 3.5e-08, + "input_cost_per_token": 1.4e-07, "litellm_provider": "refuseless", "max_input_tokens": 1048576, "max_output_tokens": 130000, "mode": "chat", - "output_cost_per_token": 7.5e-07, + "output_cost_per_token": 5e-07, "source": "https://refuseless.com/pricing", "supports_function_calling": true, "supports_prompt_caching": true, @@ -76333,6 +77637,15 @@ "max" ] }, + "fireworks_ai/accounts/fireworks/models/rune": { + "cache_read_input_token_cost": 8.4e-09, + "input_cost_per_token": 4.2e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://api.fireworks.ai/v1/serverless/models?format=nested" + }, "bedrock_mantle/openai.gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, @@ -76370,5 +77683,614 @@ "supports_vision": true, "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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 + }, + "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", + "input_cost_per_token": 1.05e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 2.1e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-spot": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.05e-05, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-spot": { + "cache_read_input_token_cost": 1.82e-07, + "input_cost_per_token": 9.8e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.08e-06, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-1017-spot": { + "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 2.1e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 8.4e-07, + "source": "https://api.fireworks.ai/v1/serverless/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 1ab03fb3fdc..d73ac19bae1 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 0b3480cab23..00cdb1587aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -296,7 +296,7 @@ ci = [ "anthropic==0.84.0", "blockbuster==1.5.26", "beautifulsoup4==4.14.3", - "pylint==4.0.5", + "pylint==4.0.9", "langchain-openai==1.1.14", "langgraph>=1.2.4,<1.3.0", "langgraph-prebuilt>=1.1.0,<1.3.0", @@ -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/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 853b10ad97b..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 @@ -171,7 +172,13 @@ ignored_function_names = [ "_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py "_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py "_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py + "_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/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index bb75569ad99..a0f185c4077 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -183,6 +183,48 @@ class TestResponses: ) assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" + @pytest.mark.covers("llm.responses.openai.basic.nonstream.works") + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(Provider.OPENAI,), + models=(OPENAI_RESPONSES_BACKEND,), + mode=Mode.NONSTREAM, + ) + ) + def test_responses_forwards_moderation( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register(proxy, resources, _openai_responses_params()) + client: Final = sdk.openai(resources.key()) + + response: Final = client.responses.create( + model=model, + input="hi", + extra_body={**NO_PROXY_CACHE, "moderation": {"model": "omni-moderation-latest"}}, + ) + response_body: Final = TypeAdapter(dict[str, object]).validate_python(cast(object, response.model_dump())) + moderation_value: Final = response_body.get("moderation") + assert moderation_value is not None, f"response had no moderation results: {response_body!r}" + moderation: Final = TypeAdapter(dict[str, object]).validate_python(moderation_value) + moderation_input: Final = TypeAdapter(dict[str, object]).validate_python(moderation["input"]) + moderation_output: Final = TypeAdapter(dict[str, object]).validate_python(moderation["output"]) + assert moderation_input.get("type") == "moderation_result", ( + f"unexpected input moderation: {moderation_input!r}" + ) + assert ( + moderation_output.get("type") == "moderation_result" + ), f"unexpected output moderation: {moderation_output!r}" + + with pytest.raises(openai.BadRequestError) as exc_info: + client.responses.create( + model=model, + input="hi", + extra_body={**NO_PROXY_CACHE, "moderation": {"model": "does-not-exist"}}, + ) + assert exc_info.value.status_code == 400 + @pytest.mark.covers("llm.responses.openai.basic.stream.works") @meta( Subject( diff --git a/tests/integration/_support/riva.py b/tests/integration/_support/riva.py new file mode 100644 index 00000000000..8552bc2197d --- /dev/null +++ b/tests/integration/_support/riva.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import io +import threading +import wave +from array import array +from collections.abc import Iterable, Iterator +from concurrent import futures +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from queue import SimpleQueue +from typing import Final + +import grpc +from riva.client.proto import riva_asr_pb2, riva_asr_pb2_grpc + +SAMPLE_RATE_HZ: Final = 16000 +OFFLINE_ONLY_REFUSAL: Final = ( + "Error: Unavailable model requested given these parameters: language_code=en-US; sample_rate=16000; type=online;" +) +_STALL_LIMIT_SECONDS: Final = 120 + + +class Rpc(Enum): + RECOGNIZE = "Recognize" + STREAMING = "StreamingRecognize" + + +@dataclass(frozen=True, slots=True) +class Word: + word: str + start_ms: int + end_ms: int + + +@dataclass(frozen=True, slots=True) +class RivaCall: + rpc: Rpc + audio: bytes + cancelled: bool + + +@dataclass(frozen=True, slots=True) +class RivaServer: + address: str + calls: SimpleQueue[RivaCall] + + def drain(self) -> tuple[RivaCall, ...]: + return tuple(self.calls.get_nowait() for _ in range(self.calls.qsize())) + + +def _clip_samples(seconds: float) -> array[int]: + return array("h", (((index * 37) % 2001 - 1000) * 16 for index in range(int(seconds * SAMPLE_RATE_HZ)))) + + +def pcm_wav(seconds: float) -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(SAMPLE_RATE_HZ) + writer.writeframes(_clip_samples(seconds).tobytes()) + return buffer.getvalue() + + +def pcm_bytes(seconds: float) -> int: + return int(seconds * SAMPLE_RATE_HZ) * 2 + + +def carries_clip(audio: bytes, seconds: float) -> bool: + sent: Final = _clip_samples(seconds) + if len(audio) != len(sent) * 2: + return False + received: Final = array("h", audio) + return all(abs(left - right) <= 1 for left, right in zip(sent, received, strict=True)) + + +class _Asr(riva_asr_pb2_grpc.RivaSpeechRecognitionServicer): + def __init__(self, words: tuple[Word, ...], offline_only: bool, stall: bool, calls: SimpleQueue[RivaCall]) -> None: + self._words: Final = words + self._offline_only: Final = offline_only + self._stall: Final = stall + self._calls: Final = calls + + def _alternative(self, config: riva_asr_pb2.RecognitionConfig) -> riva_asr_pb2.SpeechRecognitionAlternative: + timed: Final = self._words if config.enable_word_time_offsets else () + return riva_asr_pb2.SpeechRecognitionAlternative( + transcript=" ".join(word.word for word in self._words), + confidence=0.9, + words=[ + riva_asr_pb2.WordInfo(word=word.word, start_time=word.start_ms, end_time=word.end_ms) for word in timed + ], + ) + + def Recognize( + self, request: riva_asr_pb2.RecognizeRequest, context: grpc.ServicerContext + ) -> riva_asr_pb2.RecognizeResponse: + if self._stall: + released: Final = threading.Event() + context.add_callback(released.set) + released.wait(timeout=_STALL_LIMIT_SECONDS) + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=not context.is_active())) + return riva_asr_pb2.RecognizeResponse() + self._calls.put(RivaCall(Rpc.RECOGNIZE, request.audio, cancelled=False)) + return riva_asr_pb2.RecognizeResponse( + results=[riva_asr_pb2.SpeechRecognitionResult(alternatives=[self._alternative(request.config)])] + ) + + def StreamingRecognize( + self, request_iterator: Iterable[riva_asr_pb2.StreamingRecognizeRequest], context: grpc.ServicerContext + ) -> Iterator[riva_asr_pb2.StreamingRecognizeResponse]: + if self._offline_only: + self._calls.put(RivaCall(Rpc.STREAMING, b"", cancelled=False)) + context.abort(grpc.StatusCode.INVALID_ARGUMENT, OFFLINE_ONLY_REFUSAL) + requests: Final = tuple(request_iterator) + config: Final = requests[0].streaming_config.config if requests else riva_asr_pb2.RecognitionConfig() + self._calls.put( + RivaCall(Rpc.STREAMING, b"".join(request.audio_content for request in requests), cancelled=False) + ) + yield riva_asr_pb2.StreamingRecognizeResponse( + results=[riva_asr_pb2.StreamingRecognitionResult(alternatives=[self._alternative(config)], is_final=True)] + ) + + +@contextmanager +def riva_server( + words: tuple[Word, ...], + *, + offline_only: bool = False, + stall: bool = False, + max_receive_bytes: int | None = None, +) -> Iterator[RivaServer]: + calls: Final[SimpleQueue[RivaCall]] = SimpleQueue() + options: Final = () if max_receive_bytes is None else (("grpc.max_receive_message_length", max_receive_bytes),) + server: Final = grpc.server(futures.ThreadPoolExecutor(max_workers=8), options=options) + riva_asr_pb2_grpc.add_RivaSpeechRecognitionServicer_to_server(_Asr(words, offline_only, stall, calls), server) + port: Final = server.add_insecure_port("127.0.0.1:0") + server.start() + try: + yield RivaServer(f"127.0.0.1:{port}", calls) + finally: + server.stop(grace=None).wait() diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index df6036ff8d7..e278a796146 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -19,6 +19,7 @@ class Request: target: str headers: Mapping[str, str] body: bytes + header_lines: tuple[tuple[str, str], ...] = () class AbortedBody(Exception): @@ -136,7 +137,7 @@ def wire_server( self.close_connection = True disconnected.put(self.path) return - request: Final = Request(self.command, self.path, headers, body) + request: Final = Request(self.command, self.path, headers, body, tuple(self.headers.items())) received.put(request) try: reply = respond(request) diff --git a/tests/integration/authorization/test_agent_least_privilege.py b/tests/integration/authorization/test_agent_least_privilege.py index 498a53c3ca9..d993920a32e 100644 --- a/tests/integration/authorization/test_agent_least_privilege.py +++ b/tests/integration/authorization/test_agent_least_privilege.py @@ -313,7 +313,7 @@ def test_agent_access_default_deny_toggles_at_runtime_through_config_api( for field in config_fields.json() if object_value(field)["field_name"] == "agent_access_default_deny" ) - assert setting["field_type"] == "Boolean" and setting["field_value"] == "true", setting + assert setting["field_type"] == "Boolean" and setting["field_value"] is True, setting eventually( lambda: tuple(_card(candidate, agent_id, ungranted_key) for _ in range(8)), lambda responses: all(response.status_code == 403 for response in responses), 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/configuration/test_openapi_docs_route.py b/tests/integration/configuration/test_openapi_docs_route.py new file mode 100644 index 00000000000..ccb30c2f784 --- /dev/null +++ b/tests/integration/configuration/test_openapi_docs_route.py @@ -0,0 +1,165 @@ +import time +from collections.abc import Iterator +from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, list_value, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from pydantic import JsonValue, TypeAdapter + +from litellm.proxy._lazy_features import LAZY_FEATURES + +DOCS_SECONDS: Final = 120 +COLD_DOCS: Final = 4 +CHATS: Final = 8 +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +OWNED_PROXY_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +TWO_OWNED_PROXIES_TIMEOUT: Final = pytest.mark.timeout(2 * (2 * graceful_stop_seconds() + 120)) +LAZY_ROUTES_FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES" +WARMED_FEATURE: Final = next(feature for feature in LAZY_FEATURES if feature.name == "policies") + + +def _timed_get(candidate: Gateway, path: str) -> tuple[httpx.Response, float]: + started: Final = time.perf_counter() + response: Final = candidate.request("GET", path) + return response, time.perf_counter() - started + + +def _liveliness_until(done: Future[tuple[httpx.Response, float]], candidate: Gateway) -> Iterator[tuple[int, float]]: + while not done.done(): + response, seconds = _timed_get(candidate, "/health/liveliness") + yield response.status_code, seconds + + +def _schema(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text[:500] + return JSON_OBJECT.validate_json(response.content) + + +def _servers(response: httpx.Response) -> list[JsonValue]: + return list(list_value(_schema(response).get("servers", []))) + + +def _chat(candidate: Gateway, model: str, key: str, content: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": content}], "cache": {"no-cache": True}}, + key=key, + ) + + +def _warmed_feature_operations(response: httpx.Response) -> dict[str, JsonValue]: + paths: Final = object_value(_schema(response)["paths"]) + return {path: operations for path, operations in paths.items() if WARMED_FEATURE.matches(path)} + + +@OWNED_PROXY_TIMEOUT +def test_liveliness_answers_while_a_cold_openapi_schema_builds(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + ThreadPoolExecutor(1) as pool, + ): + docs: Final = pool.submit(_timed_get, owned.gateway, "/openapi.json") + probes: Final = tuple(_liveliness_until(docs, owned.gateway)) + docs_response, docs_seconds = docs.result() + assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"]) + assert probes, "no liveliness probe ran while the docs request was pending" + assert {status for status, _ in probes} == {200}, probes + slowest: Final = max(seconds for _, seconds in probes) + assert slowest < docs_seconds / 3, ( + f"a liveliness probe waited {slowest:.2f}s of the {docs_seconds:.2f}s cold docs request" + ) + + +@OWNED_PROXY_TIMEOUT +def test_concurrent_cold_openapi_requests_share_the_schema_later_requests_get(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + ThreadPoolExecutor(COLD_DOCS) as pool, + ): + cold: Final = tuple(pool.map(lambda _: owned.gateway.request("GET", "/openapi.json"), range(COLD_DOCS))) + warm: Final = owned.gateway.request("GET", "/openapi.json") + assert [response.status_code for response in (*cold, warm)] == [200] * (COLD_DOCS + 1) + assert {response.content for response in cold} == {warm.content} + assert "/v1/chat/completions" in object_value(_schema(warm)["paths"]) + + +@OWNED_PROXY_TIMEOUT +def test_chat_completions_overlapping_a_cold_openapi_build_all_succeed(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, client_timeout=DOCS_SECONDS) as owned, + owned.gateway.scenario() as scenario, + ThreadPoolExecutor(CHATS + 1) as pool, + ): + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + first_chat: Final = _chat(owned.gateway, model, key, "before the docs build") + docs: Final = pool.submit(owned.gateway.request, "GET", "/openapi.json") + chats: Final = tuple( + pool.map(lambda index: _chat(owned.gateway, model, key, f"overlapping docs build {index}"), range(CHATS)) + ) + chats_returned_during_build: Final = not docs.done() + docs_response: Final = docs.result() + assert first_chat.status_code == 200, first_chat.text[:300] + assert chats_returned_during_build, "the chat calls returned only after the cold docs build finished" + assert [(response.status_code, response.text[:300]) for response in chats if response.status_code != 200] == [] + replies: Final = tuple(JSON_OBJECT.validate_json(response.content) for response in chats) + assert {reply["object"] for reply in replies} == {"chat.completion"} + assert all(list_value(reply["choices"]) for reply in replies), replies + assert "/v1/chat/completions" in object_value(_schema(docs_response)["paths"]) + + +@OWNED_PROXY_TIMEOUT +def test_docs_filtered_mode_serves_only_llm_routes(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy_process(gateway, tmp_path, {"DOCS_FILTERED": "True"}, client_timeout=DOCS_SECONDS) as owned: + paths: Final = object_value(_schema(owned.gateway.request("GET", "/openapi.json"))["paths"]) + assert "/v1/chat/completions" in paths + assert {"/key/generate", "/user/new", "/team/new"}.isdisjoint(paths), sorted(paths) + + +@OWNED_PROXY_TIMEOUT +def test_each_root_path_prefix_lists_itself_first_in_servers(gateway: Gateway, tmp_path: Path) -> None: + prefixes: Final = ("/tenant-a", "/tenant-b", "/tenant-a") + with owned_proxy_process( + gateway, tmp_path, {"SERVER_ROOT_PATHS": "/tenant-a,/tenant-b"}, client_timeout=DOCS_SECONDS + ) as owned: + served: Final = tuple( + (prefix, _servers(owned.gateway.request("GET", f"{prefix}/openapi.json"))) for prefix in prefixes + ) + assert served == tuple((prefix, [{"url": prefix}]) for prefix in prefixes) + + +@OWNED_PROXY_TIMEOUT +def test_a_scalar_server_root_path_is_listed_once_in_servers(gateway: Gateway, tmp_path: Path) -> None: + overrides: Final = {"SERVER_ROOT_PATH": "/gateway", "LITELLM_UI_PATH": str(tmp_path / "ui")} + with owned_proxy_process(gateway, tmp_path, overrides, client_timeout=DOCS_SECONDS) as owned: + servers: Final = _servers(owned.gateway.request("GET", "/gateway/openapi.json")) + assert servers == [{"url": "/gateway"}] + + +@TWO_OWNED_PROXIES_TIMEOUT +def test_a_lazy_feature_loaded_during_a_cold_openapi_build_shows_its_live_routes_in_the_next_schema( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy_process(gateway, tmp_path, {LAZY_ROUTES_FLAG: "true"}, client_timeout=DOCS_SECONDS) as eager: + live: Final = _warmed_feature_operations(eager.gateway.request("GET", "/openapi.json")) + with ( + owned_proxy_process( + gateway, tmp_path, {}, remove_environment=(LAZY_ROUTES_FLAG,), client_timeout=DOCS_SECONDS + ) as lazy, + ThreadPoolExecutor(1) as pool, + ): + docs: Final = pool.submit(lazy.gateway.request, "GET", "/openapi.json") + warmed: Final = lazy.gateway.request("POST", f"/lazy/warm/{WARMED_FEATURE.name}") + warmed_during_build: Final = not docs.done() + during: Final = _warmed_feature_operations(docs.result()) + after: Final = _warmed_feature_operations(lazy.gateway.request("GET", "/openapi.json")) + assert warmed.status_code == 200, warmed.text + assert warmed_during_build, "the feature finished loading only after the cold docs build finished" + assert live, "the eager proxy published no route of the warmed feature" + assert during != live, "the cold docs build started after the feature had already loaded" + assert after == live diff --git a/tests/integration/cost_calculation/cases/decisions.py b/tests/integration/cost_calculation/cases/decisions.py index 1cc57afe8ed..6594510b232 100644 --- a/tests/integration/cost_calculation/cases/decisions.py +++ b/tests/integration/cost_calculation/cases/decisions.py @@ -162,6 +162,37 @@ TYPESAFE_JEV_LATEST_SYSTEMONE: Final = CostTrackingTestCase( completion_tokens=22, ), ) +TYPESAFE_JEV_PREVIEW_ANSWERED_AS_PRICED_VERSION_SYSTEMONE: Final = CostTrackingTestCase( + name="typesafe/jev-preview-answered-as-priced-version-systemone", + covers="quota_management.spend_tracking.systemone.logs_cost", + model="typesafe/jev-preview", + endpoint="/v1/systemone", + request={ + "model": "$MODEL", + "state": "The replacement part arrived but the wrong size.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + upstream_body={ + "model": "jev-preview", + "state": "The replacement part arrived but the wrong size.", + "questions": {"is_refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}}, + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "jev-latest", + "answers": {"is_refund": {"type": "noul", "noul": 0.31}}, + "usage": {"input_tokens": 311, "output_tokens": 9}, + }, + ), + expected=ExactExpected( + spend=311 * 4.2e-08, + input_cost=311 * 4.2e-08, + output_cost=0.0, + prompt_tokens=311, + completion_tokens=9, + ), +) # perplexity/pplx-decider-v1-27b PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS: Final = CostTrackingTestCase( @@ -514,6 +545,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( OPENAI_GPT_6_LUNA_SYSTEMONE, TYPESAFE_JEV_LATEST_DECISIONS, TYPESAFE_JEV_LATEST_SYSTEMONE, + TYPESAFE_JEV_PREVIEW_ANSWERED_AS_PRICED_VERSION_SYSTEMONE, PERPLEXITY_PPLX_DECIDER_V1_27B_DECISIONS, PERPLEXITY_PPLX_DECIDER_V1_27B_SYSTEMONE, CLOUDFLARE_CLEF_FLASH_DECISIONS, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 16107def2ad..32722b37c9b 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -121,6 +121,12 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token": 4.2e-08, "output_cost_per_token": 0.0, }, + "typesafe/jev-preview": { + "litellm_provider": "typesafe", + "mode": "evaluation", + "input_cost_per_token": 8.4e-08, + "output_cost_per_token": 0.0, + }, "fireworks_ai/accounts/fireworks/routers/rune": { "litellm_provider": "fireworks_ai", "mode": "evaluation", diff --git a/tests/integration/management/test_scim_admin_group.py b/tests/integration/management/test_scim_admin_group.py index 3aeed45d186..ffbec93d4a8 100644 --- a/tests/integration/management/test_scim_admin_group.py +++ b/tests/integration/management/test_scim_admin_group.py @@ -691,9 +691,6 @@ def test_scim_group_put_returns_the_updated_members(gateway: Gateway, tmp_path: @pytest.mark.timeout(300) def test_scim_admin_group_delete_revokes_the_members_admin_key(gateway: Gateway, tmp_path: Path) -> None: - pytest.skip( - "BUG: DELETE /scim/v2/Groups/{id} demotes the admin group's members in the DB but their existing keys keep proxy_admin access on /key/generate" - ) group_name: Final = f"scim-admin-delete-{uuid.uuid4().hex}" config: Final = _owned_scim_config(tmp_path / "scim-admin-group-delete.yaml", {"scim_admin_group": group_name}) with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned, owned.gateway.scenario() as scenario: diff --git a/tests/integration/management/test_scim_group_delete.py b/tests/integration/management/test_scim_group_delete.py index 4a4a93f8bd2..53af34c12e0 100644 --- a/tests/integration/management/test_scim_group_delete.py +++ b/tests/integration/management/test_scim_group_delete.py @@ -50,7 +50,7 @@ from integration._support.scim import ( ) from pydantic import JsonValue, TypeAdapter -from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS pytestmark: Final = pytest.mark.shares_worker_state @@ -70,6 +70,17 @@ FROM "LiteLLM_AuditLog" WHERE object_id IN (%s, %s) AND action = 'deleted' ORDER BY table_name """ +AUDIT_ROWS_NAMING_SQL: Final = """ +SELECT table_name, object_id, changed_by_api_key +FROM "LiteLLM_AuditLog" +WHERE object_id = ANY(%s::text[]) AND action = 'deleted' +ORDER BY table_name, object_id +""" +OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """ +SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age +FROM pg_stat_activity +WHERE %s::int = ANY(pg_blocking_pids(pid)) +""" _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _PIDS: Final = TypeAdapter(tuple[int, ...]) @@ -115,15 +126,28 @@ def _deleted_audit_rows(team: str, hashed_key: str) -> list[dict[str, JsonValue] return read_rows(DELETED_AUDIT_ROWS_SQL, (team, hashed_key)) +def _audit_rows_naming(object_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]: + return read_rows(AUDIT_ROWS_NAMING_SQL, ("{" + ",".join(object_ids) + "}",)) + + +def _oldest_transaction_blocked_by(holder: int) -> float: + age: Final = read_rows(OLDEST_TRANSACTION_BLOCKED_BY_SQL, (str(holder),))[0]["age"] + assert isinstance(age, float), age + return age + + def _worker_startups(log: Path) -> tuple[tuple[int, ...], int]: text: Final = log.read_text() return _PIDS.validate_python(_STARTED_WORKER.findall(text)), text.count("Application startup complete.") @contextmanager -def _fresh_connection(candidate: Gateway) -> Iterator[Gateway]: +def _fresh_connection(candidate: Gateway, timeout: httpx.Timeout | float | None = None) -> Iterator[Gateway]: with httpx.Client( - base_url=candidate.client.base_url, timeout=candidate.client.timeout, trust_env=False, limits=GATEWAY_LIMITS + base_url=candidate.client.base_url, + timeout=candidate.client.timeout if timeout is None else timeout, + trust_env=False, + limits=GATEWAY_LIMITS, ) as client: yield Gateway(client, candidate.key, candidate.upstream_url) @@ -294,6 +318,82 @@ def test_a_group_delete_revokes_member_keys_and_a_re_created_group_gets_its_serv assert _models_status(gateway, member_key) == 401 +@pytest.mark.timeout(OWNED_PROXY_TIMEOUT) +def test_a_group_delete_sent_with_the_idps_scim_token_revokes_member_keys_and_audits_them_under_that_token( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + scim_token: Final = scenario.key(allowed_routes=["/scim/*"], models=[]) + users: Final = seed_users(scenario, 3) + team: Final = created_team(scenario, create_group(candidate, users)) + member_keys: Final = (team_key(candidate, team, users[0]), team_key(candidate, team, users[1])) + service_key: Final = service_account_key(candidate, scenario, team) + + deleted: Final = delete_group(candidate, team, key=scim_token) + assert deleted.status_code == 204, deleted.text + + assert keys_of(team) == frozenset({_hashed(service_key)}) + for key in member_keys: + _assert_revoked_everywhere((candidate,), key) + audited: Final = (team, *(_hashed(key) for key in member_keys), _hashed(service_key)) + rows: Final = eventually(lambda: _audit_rows_naming(audited), lambda found: len(found) >= 3, seconds=30) + token: Final = _hashed(scim_token) + assert rows == [ + {"table_name": "LiteLLM_TeamTable", "object_id": team, "changed_by_api_key": token}, + *( + {"table_name": "LiteLLM_VerificationToken", "object_id": hashed, "changed_by_api_key": token} + for hashed in sorted(_hashed(key) for key in member_keys) + ), + ] + + +@pytest.mark.timeout(TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 120) +def test_a_group_delete_that_outlives_its_transaction_on_a_locked_member_row_changes_nothing_and_a_retry_lands( + gateway: Gateway, +) -> None: + patient_timeout: Final = TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 90 + with ( + gateway.scenario() as scenario, + _fresh_connection(gateway, timeout=patient_timeout) as patient, + ThreadPoolExecutor(max_workers=1) as pool, + ): + users: Final = seed_users(scenario, 3) + team: Final = created_team(scenario, create_group(gateway, users)) + key: Final = team_key(gateway, team, users[0]) + + with held_user_row(users[0]) as holder: + delete: Final = pool.submit(delete_group, patient, team) + eventually( + lambda: ( + delete.done() + or _oldest_transaction_blocked_by(holder) > TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 1 + ), + lambda outlived: outlived, + seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + 30, + ) + failed: Final = delete.result(timeout=patient_timeout) + + assert failed.status_code >= 500, failed.text + assert_landed(gateway, team, users) + assert keys_of(team) == frozenset({_hashed(key)}) + assert read_rows(DELETED_TEAM_RECORDS_SQL, (team,)) == [] + assert read_rows(DELETED_KEY_RECORDS_SQL, (_hashed(key),)) == [] + assert _models_status(gateway, key) == 200 + + retried: Final = delete_group(gateway, team) + assert retried.status_code == 204, retried.text + assert_gone(gateway, team, users, key) + + def test_deleting_an_unknown_group_is_404(gateway: Gateway) -> None: missing: Final = f"missing-{uuid.uuid4().hex}" response: Final = delete_group(gateway, missing) diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index e752c835e4f..1a8140373b0 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -33,7 +33,7 @@ from tests.integration._support.client import ( string_value, ) from tests.integration._support.database import read_rows, write_rows -from tests.integration._support.process import owned_proxy_process +from tests.integration._support.process import graceful_stop_seconds, owned_proxy_process LARGE_ROSTER: Final = 250 POOL_LIMIT: Final = 5 @@ -131,6 +131,13 @@ def _roster_user_ids(roster: JsonValue) -> list[str]: return sorted(string_value(object_value(member)["user_id"]) for member in roster) +def _deleted_audit_rows(object_id: str, table_name: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT before_value FROM "LiteLLM_AuditLog" WHERE object_id = %s AND table_name = %s AND action = %s', + (object_id, table_name, "deleted"), + ) + + def _waiters_on_lock_held_by(backend_pid: int) -> int: rows: Final = read_rows(WAITERS_ON_HELD_LOCK_SQL, ("%pg_advisory_xact_lock%", str(backend_pid))) waiting: Final = rows[0]["waiting"] @@ -551,6 +558,41 @@ def test_delete_writes_one_audit_row_for_the_team_and_one_per_key(gateway: Gatew ] +@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120)) +def test_two_deletes_racing_on_one_team_write_one_tombstone_and_one_audit_row(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned, + owned.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=2) as pool, + ): + user: Final = scenario.user() + team: Final = _create_team(scenario) + _bulk_member_add(owned.gateway, team, [user]) + + with psycopg.connect(os.environ["DATABASE_URL"]) as holder: + holder.execute(TAKE_TEAM_LOCK_SQL, (team,)) + deletes: Final = tuple( + _post_in_background(pool, owned.gateway, "/team/delete", {"team_ids": [team]}) for _ in range(2) + ) + eventually( + lambda: _waiters_on_lock_held_by(holder.info.backend_pid), lambda waiting: waiting == 2, seconds=30 + ) + + assert [delete.result(timeout=60).status_code for delete in deletes] == [200, 200] + tombstones: Final = _tombstones(team) + assert [row["team_id"] for row in tombstones] == [team] + assert _roster_user_ids(tombstones[0]["members_with_roles"]) == sorted([PROXY_ADMIN, user]) + + team_audit_rows: Final = _deleted_audit_rows(team, "LiteLLM_TeamTable") + assert len(team_audit_rows) == 1, team_audit_rows + assert object_value(team_audit_rows[0]["before_value"])["team_id"] == team + + def test_second_delete_of_the_same_team_is_404_with_one_tombstone(gateway: Gateway) -> None: with gateway.scenario() as scenario: user: Final = scenario.user() diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 490dd3cfa0a..d0fdc6935c7 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -1041,6 +1041,45 @@ def test_streaming_chat_auto_execute_answered_in_one_round_ends_with_the_usage_i ), usage +NO_CACHE: Final = {"no-cache": True} + + +def _upstream_rounds(rig: Rig) -> tuple[int, int]: + answered: Final = tuple( + _has_tool_result(JSON_OBJECT.validate_json(request.body)) + for request in rig.wire.drain() + if request.method == "POST" + ) + return answered.count(False), answered.count(True) + + +def _billed_rows(call_id: str, call_type: str, billed: float) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s', + (call_id, call_type), + ), + lambda found: sum(float(str(row["spend"])) for row in found) >= billed * (1 - 1e-6), + seconds=70, + ) + + +@pytest.mark.parametrize(("surface", "call_type"), (("responses", "aresponses"), ("chat", "acompletion"))) +def test_nonstream_follow_up_failing_on_every_retry_bills_each_attempts_first_round_on_one_failure_row( + gateway: Gateway, surface: Surface, call_type: str +) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, surface, billing=billing, prices={**PRICES, "num_retries": 2}) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO], cache=NO_CACHE) + assert response.status_code == 500 and FOLLOW_UP_OUTAGE in response.text, response.text + first_rounds, follow_ups = _upstream_rounds(rig) + assert first_rounds == follow_ups > 1, (first_rounds, follow_ups) + paid: Final = _price(*(FIRST_ROUND,) * first_rounds) + rows: Final = _billed_rows(response.headers["x-litellm-call-id"], call_type, paid) + assert [row["status"] for row in rows] == ["failure"], rows + assert float(str(rows[0]["spend"])) == pytest.approx(paid), rows + + def _tool_rows(call_id: str, expected: int) -> list[dict[str, JsonValue]]: return eventually( lambda: read_rows( diff --git a/tests/integration/observability/test_passthrough_operation_name.py b/tests/integration/observability/test_passthrough_operation_name.py new file mode 100644 index 00000000000..7b01389a993 --- /dev/null +++ b/tests/integration/observability/test_passthrough_operation_name.py @@ -0,0 +1,163 @@ +import json +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.otlp_sink import Span, SpanSinks, recorded_spans +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +pytestmark: Final = pytest.mark.timeout(180) + +_EXPORTERS: Final = { + "otel-v1-semconv": {"OTEL_SEMCONV_STABILITY_OPT_IN": "gen_ai_latest_experimental"}, + "otel-v2": {"LITELLM_OTEL_V2": "1"}, +} +_ANTHROPIC_MESSAGE: Final = { + "id": "msg_passthrough_operation", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, +} +_STREAM_REQUESTS: Final = {"stream-requested": {"query": "q", "stream": True}, "stream-detected": {"query": "q"}} + + +@dataclass(frozen=True, slots=True) +class _Rig: + proxy: Gateway + sink: str + + +def _upstream(request: Request) -> Reply: + if request.target.endswith("/v1/messages"): + return Reply(body=json.dumps(_ANTHROPIC_MESSAGE).encode()) + _, _, reported = request.target.partition("/report/") + headers: Final = {"x-litellm-operation-name": reported} if reported else {} + if request.target.startswith("/app/stream/cut"): + return Reply( + content_type="text/event-stream", + chunks=(b'data: {"chunk": 1}\n\n', b'data: {"chunk": 2}\n\n'), + abort_after=1, + ) + if request.target.startswith("/app/stream"): + return Reply(content_type="text/event-stream", chunks=(b'data: {"chunk": 1}\n\n',), headers=headers) + return Reply(body=b'{"chunks": []}', headers=headers) + + +@pytest.fixture(scope="module", params=sorted(_EXPORTERS)) +def rig( + request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks +) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("passthrough-operation") + exporter: Final = _EXPORTERS[request.param] + with gateway_from_environment() as base, wire_server(_upstream) as upstream: + shared: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **shared, + "litellm_settings": {**shared["litellm_settings"], "callbacks": ["otel"]}, + "callback_settings": { + "otel": {"exporter": "http/json", "endpoint": audit_sinks.operator, "use_simple_processor": True} + }, + "general_settings": { + **shared["general_settings"], + "pass_through_endpoints": [ + {"path": "/doc-app", "target": upstream.url + "/app", "include_subpath": True} + ], + }, + "environment_variables": { + **shared["environment_variables"], + "ANTHROPIC_API_BASE": upstream.url, + "ANTHROPIC_API_KEY": "scripted", + }, + } + path: Final = directory / "passthrough-operation.yaml" + path.write_text(yaml.safe_dump(config)) + overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, **exporter} + unset: Final = tuple( + name for name in ("LITELLM_OTEL_V2", "OTEL_SEMCONV_STABILITY_OPT_IN") if name not in exporter + ) + with owned_proxy(base, directory, overrides, config=path, remove_environment=unset) as proxy: + yield _Rig(proxy, audit_sinks.operator) + + +def _span_of(rig: _Rig, since: int, call_id: str) -> Span: + spans: Final = eventually( + lambda: tuple( + span + for span in recorded_spans(rig.sink, since)[1] + if span["attributes"].get("litellm.call_id") == call_id and "gen_ai.operation.name" in span["attributes"] + ), + lambda found: len(found) == 1, + seconds=60, + return_last_on_timeout=True, + ) + assert len(spans) == 1, spans + return spans[0] + + +def _operation_span(rig: _Rig, method: str, path: str, body: dict[str, object] | None = None) -> Span: + since, _ = recorded_spans(rig.sink) + response: Final = rig.proxy.request(method, path, body) + assert response.status_code == 200, response.text + return _span_of(rig, since, response.headers["x-litellm-call-id"]) + + +def _cut_stream_span(rig: _Rig, path: str) -> Span: + since, _ = recorded_spans(rig.sink) + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + with rig.proxy.client.stream("POST", path, json={"query": "q", "stream": True}, headers=headers) as response: + assert response.status_code == 200 + with pytest.raises(httpx.RemoteProtocolError): + response.read() + return _span_of(rig, since, response.headers["x-litellm-call-id"]) + + +def _assert_operation(span: Span, operation: str) -> None: + assert span["attributes"]["gen_ai.operation.name"] == operation, span + assert span["name"].split(" ")[0] == operation, span + + +def test_configured_route_without_the_header_is_labelled_passthrough(rig: _Rig) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/search", {"query": "q"}), "litellm.passthrough") + + +def test_operation_reported_by_the_target_labels_the_span(rig: _Rig) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/report/retrieval", {"query": "q"}), "retrieval") + + +def test_unknown_operation_reported_by_the_target_is_ignored(rig: _Rig) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/report/summarize", {"query": "q"}), "litellm.passthrough") + + +@pytest.mark.parametrize("body", _STREAM_REQUESTS.values(), ids=_STREAM_REQUESTS.keys()) +def test_streamed_configured_route_without_the_header_is_labelled_passthrough( + rig: _Rig, body: dict[str, object] +) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/stream", body), "litellm.passthrough") + + +@pytest.mark.parametrize("body", _STREAM_REQUESTS.values(), ids=_STREAM_REQUESTS.keys()) +def test_operation_reported_by_a_streamed_target_labels_the_span(rig: _Rig, body: dict[str, object]) -> None: + _assert_operation(_operation_span(rig, "POST", "/doc-app/stream/report/retrieval", body), "retrieval") + + +def test_configured_route_stream_cut_by_the_target_is_labelled_passthrough(rig: _Rig) -> None: + _assert_operation(_cut_stream_span(rig, "/doc-app/stream/cut"), "litellm.passthrough") + + +def test_provider_route_stays_chat(rig: _Rig) -> None: + body: Final = {"model": "claude-sonnet-5-5", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]} + _assert_operation(_operation_span(rig, "POST", "/anthropic/v1/messages", body), "chat") + + +def test_provider_route_litellm_does_not_parse_stays_chat(rig: _Rig) -> None: + _assert_operation(_operation_span(rig, "GET", "/anthropic/v1/models"), "chat") 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_api_base_trailing_slash_wire.py b/tests/integration/providers/test_api_base_trailing_slash_wire.py index e914199e825..01b68f52032 100644 --- a/tests/integration/providers/test_api_base_trailing_slash_wire.py +++ b/tests/integration/providers/test_api_base_trailing_slash_wire.py @@ -42,7 +42,7 @@ _MESSAGES_ROWS: Final = ( Row("ollama_chat/llama3.1", "/ol", "/api/chat"), ) _EMBEDDING_ROWS: Final = ( - Row("ollama/nomic-embed-text", "/ol", "/api/embed", keyed=False), + Row("ollama/nomic-embed-text", "/ol", "/api/embed"), Row("jina_ai/jina-embeddings-v3", "/jn/v1", "/embeddings"), Row("voyage/voyage-3", "/vy/v1", "/embeddings"), Row("voyage/voyage-multimodal-3", "/vy/v1", "/multimodalembeddings"), @@ -65,6 +65,7 @@ _FULL_ROUTE_EMBEDDING_ROWS: Final = ( Row("jina_ai/jina-embeddings-v3", "/jn/v1/embeddings", ""), Row("voyage/voyage-3", "/vy/v1/embeddings", ""), Row("perplexity/pplx-embed-v1-0.6b", "/px/v1/embeddings", ""), + Row("ollama/nomic-embed-text", "/ol/api/embed", ""), ) 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_converse_redacted_thinking_wire.py b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py index 79ec0545c44..efb0609d560 100644 --- a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py +++ b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py @@ -6,7 +6,7 @@ from typing import Final from urllib.parse import unquote import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, Scenario from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, ConfigDict, JsonValue @@ -414,3 +414,92 @@ def test_converse_redacted_thinking_replays_to_the_anthropic_api_as_claude_nativ {"type": "redacted_thinking", "data": NATIVE_REDACTED}, {"type": "text", "text": FIRST_ANSWER}, ) + + +_SIGNED_THINKING: Final[dict[str, JsonValue]] = { + "type": "thinking", + "thinking": "Checking the forecast.", + "signature": "scripted-signature-" + "s" * 32, +} +_CONVERSE_REDACTED_PART: Final[dict[str, JsonValue]] = {"reasoningContent": {"redactedContent": CONVERSE_REDACTED}} +_CONVERSE_SIGNED_PART: Final[dict[str, JsonValue]] = { + "reasoningContent": { + "reasoningText": {"text": _SIGNED_THINKING["thinking"], "signature": _SIGNED_THINKING["signature"]} + } +} + + +def _replayed_to_converse( + gateway: Gateway, wire: Wire, model: str, assistants: list[JsonValue], **fields: JsonValue +) -> tuple[dict[str, JsonValue], ...]: + reply: Final = _chat( + gateway, + { + "model": model, + "thinking": THINKING, + "max_tokens": 2048, + "messages": [{"role": "user", "content": QUESTION}, *assistants, {"role": "user", "content": FOLLOW_UP}], + **fields, + }, + ) + assert reply.content == FINAL_ANSWER + return _replayed_assistant_turn(wire).content + + +def _claude_on_converse(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"bedrock/converse/{CLAUDE_MODEL_ID}", + api_key=TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + + +@pytest.mark.parametrize( + "thinking_blocks", [None, [_redacted(NATIVE_REDACTED)]], ids=["content", "content-and-thinking-blocks"] +) +def test_claude_redacted_thinking_replayed_in_content_reaches_converse_once( + gateway: Gateway, thinking_blocks: list[JsonValue] | None +) -> None: + assistant: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": [_redacted(NATIVE_REDACTED), {"type": "text", "text": FIRST_ANSWER}], + **({} if thinking_blocks is None else {"thinking_blocks": thinking_blocks}), + } + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse(gateway, wire, _claude_on_converse(scenario, wire), [assistant]) + assert sent == (_CONVERSE_REDACTED_PART, {"text": FIRST_ANSWER}) + + +@pytest.mark.parametrize( + "block, converse_part", + [(_redacted(NATIVE_REDACTED), _CONVERSE_REDACTED_PART), (_SIGNED_THINKING, _CONVERSE_SIGNED_PART)], + ids=["redacted", "signed"], +) +def test_block_split_across_merged_assistant_messages_reaches_converse_once( + gateway: Gateway, block: dict[str, JsonValue], converse_part: dict[str, JsonValue] +) -> None: + assistants: Final[list[JsonValue]] = [ + {"role": "assistant", "content": [block, {"type": "text", "text": FIRST_ANSWER}]}, + {"role": "assistant", "content": "Checking again.", "thinking_blocks": [block]}, + ] + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse(gateway, wire, _claude_on_converse(scenario, wire), assistants) + assert sent == (converse_part, {"text": FIRST_ANSWER}, {"text": "Checking again."}) + + +@pytest.mark.parametrize("drop", [True, False], ids=["drop-params", "no-drop-params"]) +def test_converse_cache_point_after_replayed_thinking_follows_drop_params(gateway: Gateway, drop: bool) -> None: + assistant: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": [ + {**_redacted(NATIVE_REDACTED), "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": FIRST_ANSWER}, + ], + } + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + sent: Final = _replayed_to_converse( + gateway, wire, _claude_on_converse(scenario, wire), [assistant], drop_params=drop + ) + replayed: Final = (_CONVERSE_REDACTED_PART, {"text": FIRST_ANSWER}) + assert sent == (replayed if drop else (*replayed, {"cachePoint": {"type": "default"}})) 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_claude_thinking_content_replay_wire.py b/tests/integration/providers/test_claude_thinking_content_replay_wire.py new file mode 100644 index 00000000000..4f2f3cc66c4 --- /dev/null +++ b/tests/integration/providers/test_claude_thinking_content_replay_wire.py @@ -0,0 +1,195 @@ +import uuid +from collections.abc import Mapping +from typing import Final, Literal + +import pytest +from integration._support.anthropic_thinking import ( + BEDROCK_MODEL, + MODEL, + NO_CACHE, + SIGNATURE, + THINKING, + answer, + marker_of, + prompt, + standard_peer, + thinking_block, +) +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue + +Route = Literal["anthropic", "bedrock_invoke", "claude_platform", "vertex_ai", "azure_ai"] +Endpoint = Literal["chat", "messages"] + +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_ENABLED: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024} +_EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"} +_SIGNED: Final = thinking_block(THINKING, SIGNATURE) +_REDACTED: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": "scripted-redacted-" + "r" * 32} +_BLANK_REDACTED: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": " "} +_ROUTES: Final[tuple[Route, ...]] = ("anthropic", "bedrock_invoke", "claude_platform", "vertex_ai", "azure_ai") + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _SentMessage(_Frozen): + role: str + content: str | tuple[dict[str, JsonValue], ...] + + +class _SentRequest(_Frozen): + messages: tuple[_SentMessage, ...] + + +def _deployment(scenario: Scenario, route: Route, wire_url: str, upstream_url: str) -> str: + match route: + case "anthropic": + return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key="scripted-key") + case "bedrock_invoke": + return scenario.model( + model=f"bedrock/invoke/{BEDROCK_MODEL}", + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + ) + case "claude_platform": + return scenario.model( + model=f"bedrock/claude_platform/{MODEL}", + api_base=wire_url, + api_key="scripted-platform-key", + aws_region_name="us-east-1", + workspace_id="scripted-workspace", + ) + case "vertex_ai": + return scenario.model( + model=f"vertex_ai/{MODEL}", + api_base=( + f"{wire_url}/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}" + f"/publishers/anthropic/models/{MODEL}" + ), + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=service_account_json(_VERTEX_PROJECT, upstream_url.rstrip("/")), + ) + case "azure_ai": + return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key") + + +def _text(text: str) -> dict[str, JsonValue]: + return {"type": "text", "text": text} + + +def _replayed_turn( + gateway: Gateway, + wire: Wire, + *, + endpoint: Endpoint = "chat", + model: str, + marker: str, + assistant: Mapping[str, JsonValue], + **fields: JsonValue, +) -> tuple[dict[str, JsonValue], ...]: + follow_up: Final = uuid.uuid4().hex + body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 2048, + "thinking": _ENABLED, + "messages": [ + {"role": "user", "content": prompt(marker)}, + {"role": "assistant", **assistant}, + {"role": "user", "content": prompt(follow_up)}, + ], + **NO_CACHE, + **fields, + } + path: Final = "/v1/chat/completions" if endpoint == "chat" else "/v1/messages" + response: Final = gateway.request("POST", path, body) + assert response.status_code == 200, response.text + assert answer(follow_up) in response.text, response.text + received: Final = wire.drain() + assert len(received) == 1, [request.body for request in received] + assert marker_of(received[0]) == follow_up, received[0].body + sent: Final = _SentRequest.model_validate_json(received[0].body).messages[1] + assert sent.role == "assistant", received[0].body + assert not isinstance(sent.content, str), received[0].body + return sent.content + + +@pytest.mark.parametrize("route", _ROUTES) +def test_redacted_thinking_replayed_in_assistant_content_reaches_claude_once(gateway: Gateway, route: Route) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, route, wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, wire, model=model, marker=marker, assistant={"content": [_REDACTED, _text(answer(marker))]} + ) + assert sent == (_REDACTED, _text(answer(marker))) + + +def test_block_in_content_and_thinking_blocks_reaches_claude_once_in_order(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, + wire, + model=model, + marker=marker, + assistant={"content": [_REDACTED, _text(answer(marker))], "thinking_blocks": [_SIGNED, _REDACTED]}, + ) + assert sent == (_SIGNED, _REDACTED, _text(answer(marker))) + + +def test_streamed_thinking_replayed_in_content_reaches_claude_without_its_index(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, + wire, + model=model, + marker=marker, + assistant={"content": [{**_SIGNED, "index": 0}, _text(answer(marker))]}, + ) + assert sent == (_SIGNED, _text(answer(marker))) + + +@pytest.mark.parametrize("drop", [True, False], ids=["drop-params", "no-drop-params"]) +def test_replayed_thinking_cache_control_and_blank_data_follow_drop_params(gateway: Gateway, drop: bool) -> None: + marker: Final = uuid.uuid4().hex + content: Final = ({**_SIGNED, "cache_control": _EPHEMERAL}, _BLANK_REDACTED, _REDACTED, _text(answer(marker))) + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", wire.url, gateway.upstream_url) + sent: Final = _replayed_turn( + gateway, wire, model=model, marker=marker, assistant={"content": list(content)}, drop_params=drop + ) + assert sent == ((_SIGNED, _REDACTED, _text(answer(marker))) if drop else content) + + +@pytest.mark.parametrize("endpoint", ["chat", "messages"]) +def test_assistant_cache_injection_point_never_marks_a_thinking_part(gateway: Gateway, endpoint: Endpoint) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(standard_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{MODEL}", + api_base=wire.url, + api_key="scripted-key", + cache_control_injection_points=[{"location": "message", "role": "assistant"}], + ) + sent: Final = _replayed_turn( + gateway, + wire, + endpoint=endpoint, + model=model, + marker=marker, + assistant={"content": [_text(answer(marker)), _SIGNED]}, + ) + assert sent == ({**_text(answer(marker)), "cache_control": _EPHEMERAL}, _SIGNED) 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 89ca674ef31..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 @@ -137,9 +137,41 @@ _PROVIDERS: Final = ( "databricks-openjev-qwen35-4b", _API_KEY, False, - None, + "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: @@ -576,13 +796,26 @@ def test_upstream_errors_keep_their_status_and_log_an_unbilled_failure(gateway: ) -def test_upstream_success_without_answers_is_a_gateway_side_server_error(gateway: Gateway) -> None: +@pytest.mark.parametrize( + ("reply", "field"), + ( + pytest.param({"model": _PERPLEXITY.body_model, "usage": _USAGE}, "answers", id="missing_answers"), + pytest.param( + {"model": _PERPLEXITY.body_model, "answers": {"defect": {"type": "Request timed out"}}, "usage": _USAGE}, + "answers.defect", + id="answer_type_timeout_text", + ), + ), +) +def test_a_malformed_upstream_success_is_a_server_error_naming_the_field( + gateway: Gateway, reply: dict[str, JsonValue], field: str +) -> None: with gateway.scenario() as scenario: - handle: Final = _register(scenario, {"model": _PERPLEXITY.body_model, "usage": _USAGE}) + handle: Final = _register(scenario, reply) model: Final = _deployment(scenario, handle, _PERPLEXITY) response: Final = _decide(gateway, model, num_retries=0) - assert 500 <= response.status_code < 600, response.text - assert "answers" in response.text + assert response.status_code == 500, response.text + assert field in response.text assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" 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_huggingface_rerank_truncate_wire.py b/tests/integration/providers/test_huggingface_rerank_truncate_wire.py new file mode 100644 index 00000000000..0d29dbf0ded --- /dev/null +++ b/tests/integration/providers/test_huggingface_rerank_truncate_wire.py @@ -0,0 +1,74 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway +from integration._support.wire import Reply, Request, wire_server + +MODEL: Final = "huggingface/BAAI/bge-reranker-base" +QUERY: Final = "what is a fox" +TEXT: Final = "a short doc about foxes" +RESPONSE: Final = json.dumps([{"index": 0, "score": 0.82}]).encode() +TRUNCATION_CASES: Final = ( + pytest.param( + MappingProxyType({}), + MappingProxyType({"truncation_direction": "Right"}), + id="omitted-truncate-uses-server-default", + ), + pytest.param( + MappingProxyType({"truncate": True, "truncation_direction": "Left"}), + MappingProxyType({"truncate": True, "truncation_direction": "Left"}), + id="explicit-true-and-left-are-forwarded", + ), + pytest.param( + MappingProxyType({"truncate": False}), + MappingProxyType({"truncate": False, "truncation_direction": "Right"}), + id="explicit-false-is-forwarded", + ), +) + + +def rerank_peer(request: Request, expected_body: Mapping[str, object]) -> Reply: + assert request.method == "POST" and request.target == "/rerank", request.target + assert request.headers["authorization"] == "Bearer integration-provider-key" + body: Final = JSON_OBJECT.validate_json(request.body) + assert body == expected_body, body + return Reply(body=RESPONSE) + + +@pytest.mark.parametrize(("request_params", "expected_truncation_params"), TRUNCATION_CASES) +def test_huggingface_rerank_forwards_truncation_options( + gateway: Gateway, + request_params: Mapping[str, bool | str], + expected_truncation_params: Mapping[str, bool | str], +) -> None: + expected_body: Final = { + "raw_scores": False, + "truncation_direction": "Right", + "query": QUERY, + "texts": [TEXT], + "return_text": False, + **expected_truncation_params, + } + with wire_server(lambda request: rerank_peer(request, expected_body)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=MODEL, + api_base=wire.url, + model_info={"mode": "rerank"}, + ) + response: Final = gateway.request( + "POST", + "/v1/rerank", + { + "model": model, + "query": QUERY, + "documents": [TEXT], + "return_documents": False, + **request_params, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["results"] == [{"index": 0, "relevance_score": 0.82}], response.text + assert len(wire.drain()) == 1, "Expected exactly one provider rerank call" diff --git a/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py b/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py new file mode 100644 index 00000000000..ca1bfb5f78a --- /dev/null +++ b/tests/integration/providers/test_huggingface_tokenizer_downloads_wire.py @@ -0,0 +1,201 @@ +import json +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from tokenizers import Regex, Tokenizer, models, pre_tokenizers + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_LLAMA_3_BACKEND: Final = "databricks-meta-llama-3-3-70b-instruct" +_DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731" +_API_KEY: Final = "synthetic-databricks-key" +_PROMPT: Final = "Count this message before the tokenizer comes back" +_COUNTS_AFTER_FAILURE: Final = 4 +_TOKENIZER_FILE: Final = "/resolve/main/tokenizer.json" +_TOKENIZER_HEADERS: Final = MappingProxyType({"X-Repo-Commit": "0" * 40, "ETag": '"scripted-char-tokenizer"'}) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +class _Usage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + prompt_tokens: int + + +class _Chunk(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + id: str + usage: _Usage | None = None + + +def _hf_environment(directory: Path) -> dict[str, str]: + return { + "HF_HOME": str(directory / "hf-home"), + "HF_HUB_CACHE": str(directory / "hf-cache"), + "HF_HUB_OFFLINE": "0", + } + + +def _missing(request: Request) -> Reply: + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + +@pytest.fixture(scope="module") +def unreachable_hub_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[OwnedProxy, Wire]]: + directory: Final = tmp_path_factory.mktemp("hf-tokenizer-unreachable") + with ( + gateway_from_environment() as gateway, + wire_server(_missing) as hub, + owned_proxy_process(gateway, directory, {**_hf_environment(directory), "HF_ENDPOINT": hub.url}) as owned, + ): + yield owned, hub + + +def _llama_3(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"databricks/{_LLAMA_3_BACKEND}", api_base=api_base, api_key=_API_KEY) + + +def _local_count(gateway: Gateway, model: str) -> tuple[JsonValue, JsonValue]: + response: Final = gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "messages": [{"role": "user", "content": _PROMPT}]}, + params={"call_endpoint": "false"}, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + return payload["tokenizer_type"], payload["total_tokens"] + + +def _frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None) -> bytes: + chunk: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _LLAMA_3_BACKEND, + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + } + return b"data: " + json.dumps(chunk).encode() + b"\n\n" + + +def _stream_without_usage(identity: str) -> Callable[[Request], Reply]: + frames: Final = ( + _frame(identity, {"role": "assistant", "content": "Counted "}, None), + _frame(identity, {"content": "locally."}, None), + _frame(identity, {}, "stop"), + b"data: [DONE]\n\n", + ) + + def respond(_request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=frames) + + return respond + + +def test_counts_after_a_failed_llama_3_tokenizer_download_never_reach_the_hub_again( + unreachable_hub_proxy: tuple[OwnedProxy, Wire], +) -> None: + owned, hub = unreachable_hub_proxy + with owned.gateway.scenario() as scenario: + model: Final = _llama_3(scenario, owned.gateway.upstream_url) + first: Final = _local_count(owned.gateway, model) + attempted: Final = hub.received.qsize() + later: Final = tuple(_local_count(owned.gateway, model) for _ in range(_COUNTS_AFTER_FAILURE)) + retried: Final = hub.received.qsize() - attempted + assert attempted > 0, "the proxy never asked the scripted hub for the llama-3 tokenizer" + assert first[0] == "openai_tokenizer", first + assert later == (first,) * _COUNTS_AFTER_FAILURE, later + assert retried == 0, f"{retried} hub requests after the download already failed" + + +def test_streamed_llama_3_chat_after_a_failed_download_counts_usage_without_reaching_the_hub( + unreachable_hub_proxy: tuple[OwnedProxy, Wire], +) -> None: + owned, hub = unreachable_hub_proxy + identity: Final = f"chatcmpl-{uuid.uuid4().hex}" + with wire_server(_stream_without_usage(identity)) as upstream, owned.gateway.scenario() as scenario: + model: Final = _llama_3(scenario, upstream.url) + counted: Final = _local_count(owned.gateway, model) + attempted: Final = hub.received.qsize() + with owned.gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + headers={"Authorization": f"Bearer {owned.gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) + retried: Final = hub.received.qsize() - attempted + sent: Final = tuple((request.method, request.target) for request in upstream.drain()) + assert counted[0] == "openai_tokenizer", counted + assert lines[-1] == "data: [DONE]", lines + chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) + assert {chunk.id for chunk in chunks} == {identity}, lines + usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) + assert tuple(usage.prompt_tokens for usage in usages) == (counted[1],), (counted, lines) + assert sent == (("POST", "/chat/completions"),), sent + assert retried == 0, f"{retried} hub requests while streaming after the download already failed" + + +def test_a_failed_llama_3_download_leaves_the_deepseek_tokenizer_counting( + unreachable_hub_proxy: tuple[OwnedProxy, Wire], +) -> None: + owned, _hub = unreachable_hub_proxy + with owned.gateway.scenario() as scenario: + llama: Final = _llama_3(scenario, owned.gateway.upstream_url) + deepseek: Final = scenario.model(model=_DEEPSEEK_V4, api_base=owned.gateway.upstream_url, api_key=_API_KEY) + assert _local_count(owned.gateway, llama)[0] == "openai_tokenizer" + counted: Final = _local_count(owned.gateway, deepseek) + assert counted[0] == "huggingface_tokenizer", counted + + +def _char_tokenizer() -> bytes: + tokenizer: Final = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated") + return tokenizer.to_str().encode() + + +def _recovering_hub(recovered: threading.Event, tokenizer: bytes) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if recovered.is_set() and request.target.endswith(_TOKENIZER_FILE): + return Reply(body=tokenizer, headers=_TOKENIZER_HEADERS) + return _missing(request) + + return respond + + +def test_a_failed_llama_3_download_is_retried_once_the_retry_window_passes(gateway: Gateway, tmp_path: Path) -> None: + recovered: Final = threading.Event() + with ( + wire_server(_recovering_hub(recovered, _char_tokenizer())) as hub, + owned_proxy_process( + gateway, + tmp_path, + {**_hf_environment(tmp_path), "HF_ENDPOINT": hub.url, "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS": "1"}, + ) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = _llama_3(scenario, owned.gateway.upstream_url) + failed: Final = _local_count(owned.gateway, model) + hub.drain() + recovered.set() + counted: Final = eventually( + lambda: _local_count(owned.gateway, model), lambda count: count[0] == "huggingface_tokenizer", seconds=30 + ) + downloads: Final = tuple(request.target for request in hub.drain() if request.target.endswith(_TOKENIZER_FILE)) + assert failed[0] == "openai_tokenizer", failed + assert counted[1] != failed[1], (failed, counted) + assert downloads, "the tokenizer came back without a new download" diff --git a/tests/integration/providers/test_image_gen_drop_params_wire.py b/tests/integration/providers/test_image_gen_drop_params_wire.py index 4307d20c598..f61d7d6dc17 100644 --- a/tests/integration/providers/test_image_gen_drop_params_wire.py +++ b/tests/integration/providers/test_image_gen_drop_params_wire.py @@ -1,16 +1,101 @@ import json -from typing import Final +import threading +from types import MappingProxyType +from typing import Final, Literal +import httpx import openai import pytest from integration._support.client import Gateway from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter +import litellm + _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _PROMPT: Final = "a transparent lighthouse at dawn" +@pytest.mark.asyncio +@pytest.mark.parametrize("router_retries", (0, 1)) +@pytest.mark.parametrize("sdk_retries", (0, 5)) +@pytest.mark.parametrize("outcome", ("503", "timeout", "success")) +async def test_proxy_image_retry_attempts_are_owned_only_by_router( + monkeypatch: pytest.MonkeyPatch, + router_retries: int, + sdk_retries: int, + outcome: Literal["503", "timeout", "success"], +) -> None: + from litellm.proxy import proxy_server + + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/images/generations" + assert request.headers["authorization"] == "Bearer synthetic-image-key" + assert _JSON_OBJECT.validate_json(request.body) == {"model": "gpt-image-1-mini", "prompt": "a red circle"} + if outcome == "timeout": + return Reply(chunks=(b"{}",), gate_before_first=release) + if outcome == "503": + return Reply( + status=503, + headers=MappingProxyType({"retry-after-ms": "1"}), + body=b'{"error":{"message":"scripted unavailable","type":"server_error"}}', + ) + return Reply(body=b'{"created":1700000000,"data":[{"b64_json":"aW1n"}]}') + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "num_retries", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + with wire_server(respond) as wire: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "image-retries", + "litellm_params": { + "model": "openai/gpt-image-1-mini", + "api_base": f"{wire.url}/v1", + "api_key": "synthetic-image-key", + "max_retries": sdk_retries, + "timeout": 0.1, + }, + } + ], + num_retries=router_retries, + disable_cooldowns=True, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "master_key", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + try: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://proxy.test" + ) as client: + response: Final = await client.post( + "/v1/images/generations", + json={"model": "image-retries", "prompt": "a red circle", "max_retries": sdk_retries}, + ) + finally: + release.set() + router.discard() + litellm.in_memory_llm_clients_cache.flush_cache() + requests: Final = wire.drain() + status: Final = {"503": 503, "timeout": 408, "success": 200}[outcome] + attempts: Final = 1 if outcome == "success" else router_retries + 1 + assert response.status_code == status, response.text + assert len(requests) == attempts, tuple(request.headers.get("x-stainless-retry-count") for request in requests) + if outcome == "success": + data: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python( + _JSON_OBJECT.validate_json(response.content)["data"] + ) + assert data[0]["b64_json"] == "aW1n" + + def test_image_generation_additional_drop_params_reaches_provider_body(gateway: Gateway) -> None: def respond(request: Request) -> Reply: assert request.method == "POST" diff --git a/tests/integration/providers/test_nvidia_riva_wire.py b/tests/integration/providers/test_nvidia_riva_wire.py new file mode 100644 index 00000000000..ef36c449761 --- /dev/null +++ b/tests/integration/providers/test_nvidia_riva_wire.py @@ -0,0 +1,137 @@ +from collections.abc import Mapping +from typing import Final + +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.openai_wire import chat_reply +from integration._support.riva import RivaServer, Rpc, Word, carries_clip, pcm_bytes, pcm_wav, riva_server +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL: Final = "nvidia_riva/nvidia/parakeet-0.6b-tdt" +_WORDS: Final = (Word("riva", 0, 400), Word("offline", 400, 900), Word("works", 900, 1400)) +_TRANSCRIPT: Final = "riva offline works" +_CLIP_SECONDS: Final = 0.5 +_REFUSED_FIELD: Final = "riva_offline is not allowed in request body" + + +def _riva_model(scenario: Scenario, server: RivaServer, **parameters: JsonValue) -> str: + return scenario.model( + model=_MODEL, + api_base=server.address, + model_info={"mode": "audio_transcription"}, + **parameters, + ) + + +def _transcribe( + gateway: Gateway, key: str, model: str, seconds: float = _CLIP_SECONDS, **fields: str +) -> httpx.Response: + return gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model, **fields}, + {"file": ("clip.wav", pcm_wav(seconds), "audio/wav")}, + key=key, + ) + + +def _sdk(gateway: Gateway, key: str) -> openai.OpenAI: + return openai.OpenAI(base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=key, max_retries=0) + + +def test_a_form_riva_offline_on_transcriptions_is_refused_before_riva_is_called(gateway: Gateway) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server) + refused: Final = _transcribe(gateway, scenario.key(), model, riva_offline="true") + assert refused.status_code == 401, refused.text + assert _REFUSED_FIELD in refused.json()["error"]["message"], refused.text + assert server.drain() == () + + +def test_a_json_riva_offline_on_chat_is_refused_before_the_provider_is_called(gateway: Gateway) -> None: + def respond(_: Request) -> Reply: + return chat_reply("chatcmpl-riva-audit", "gpt-4o-mini", "served", stream=False) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + refused: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "riva flag on chat"}], "riva_offline": True}, + key=scenario.key(), + ) + assert refused.status_code == 401, refused.text + assert _REFUSED_FIELD in refused.json()["error"]["message"], refused.text + assert wire.drain() == () + + +def test_an_offline_deployment_sends_the_whole_clip_in_one_recognize_call(gateway: Gateway) -> None: + with riva_server(_WORDS, offline_only=True) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, riva_offline=True) + transcription: Final = _sdk(gateway, scenario.key()).audio.transcriptions.create( + model=model, + file=("clip.wav", pcm_wav(_CLIP_SECONDS), "audio/wav"), + response_format="verbose_json", + timestamp_granularities=["word"], + ) + assert transcription.text == _TRANSCRIPT + assert [(word.word, word.start, word.end) for word in transcription.words or ()] == [ + (word.word, word.start_ms / 1000, word.end_ms / 1000) for word in _WORDS + ] + assert [(call.rpc, carries_clip(call.audio, _CLIP_SECONDS)) for call in server.drain()] == [ + (Rpc.RECOGNIZE, True) + ] + + +def test_a_deployment_without_riva_offline_still_streams_the_clip(gateway: Gateway) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, timeout=30) + served: Final = _transcribe(gateway, scenario.key(), model) + assert served.status_code == 200, served.text + assert served.json()["text"] == _TRANSCRIPT, served.text + assert [(call.rpc, carries_clip(call.audio, _CLIP_SECONDS)) for call in server.drain()] == [ + (Rpc.STREAMING, True) + ] + + +@pytest.mark.parametrize( + ("sent", "deployment", "status", "rpcs"), + ( + pytest.param("true", {}, 200, (Rpc.RECOGNIZE,), id="true-on-a-streaming-deployment"), + pytest.param("false", {"riva_offline": True}, 200, (Rpc.STREAMING,), id="false-on-an-offline-deployment"), + pytest.param("yes", {}, 400, (), id="malformed"), + pytest.param("yes", {"drop_params": True}, 200, (Rpc.STREAMING,), id="malformed-under-drop-params"), + ), +) +def test_an_opted_in_caller_riva_offline_picks_the_rpc( + gateway: Gateway, sent: str, deployment: Mapping[str, JsonValue], status: int, rpcs: tuple[Rpc, ...] +) -> None: + with riva_server(_WORDS) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, configurable_clientside_auth_params=["riva_offline"], **deployment) + answered: Final = _transcribe(gateway, scenario.key(), model, riva_offline=sent) + assert answered.status_code == status, answered.text + assert tuple(call.rpc for call in server.drain()) == rpcs + + +def test_offline_audio_over_the_servers_message_limit_answers_413(gateway: Gateway) -> None: + seconds: Final = 2.0 + limit: Final = pcm_bytes(seconds) // 2 + with ( + riva_server(_WORDS, offline_only=True, max_receive_bytes=limit) as server, + gateway.scenario() as scenario, + ): + model: Final = _riva_model(scenario, server, riva_offline=True) + refused: Final = _transcribe(gateway, scenario.key(), model, seconds) + assert refused.status_code == 413, refused.text + assert server.drain() == () + + +def test_a_stalled_offline_recognize_answers_408_at_the_deployment_timeout(gateway: Gateway) -> None: + with riva_server(_WORDS, offline_only=True, stall=True) as server, gateway.scenario() as scenario: + model: Final = _riva_model(scenario, server, riva_offline=True, timeout=2) + timed_out: Final = _transcribe(gateway, scenario.key(), model) + assert timed_out.status_code == 408, timed_out.text + calls: Final = eventually(server.drain, lambda observed: len(observed) > 0, seconds=30) + assert all(call.rpc is Rpc.RECOGNIZE and call.cancelled for call in calls), calls diff --git a/tests/integration/providers/test_ollama_embedding_wire.py b/tests/integration/providers/test_ollama_embedding_wire.py new file mode 100644 index 00000000000..fd38260b21c --- /dev/null +++ b/tests/integration/providers/test_ollama_embedding_wire.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import json +import threading +import time +import uuid +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value +from integration._support.timing import assert_cut_short +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +_MODEL: Final = "nomic-embed-text" +_DEPLOYMENT_KEY: Final = "synthetic-ollama-deployment-key" +_CALLER_TOKEN: Final = "Bearer synthetic-deployment-header-token" +_DELAY_SECONDS: Final = 5.0 + + +def _vector(index: int) -> list[float]: + return [0.5, float(index)] + + +def _embeddings(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + count: Final = len(body["input"]) if isinstance(body["input"], list) else 1 + return Reply( + body=json.dumps({"embeddings": [_vector(index) for index in range(count)], "prompt_eval_count": 3}).encode() + ) + + +def _vectors(response: httpx.Response) -> list[JsonValue]: + data: Final = list_value(JSON_OBJECT.validate_json(response.content)["data"]) + return [object_value(row)["embedding"] for row in data] + + +def _only(wire: Wire) -> Request: + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [("POST", "/api/embed")], [ + (request.method, request.target) for request in received + ] + return received[0] + + +def _values_of(request: Request, name: str) -> list[str]: + return [value for header, value in request.header_lines if header.lower() == name.lower()] + + +def _texts(count: int) -> list[str]: + return [f"ollama proxy embedding {uuid.uuid4().hex}" for _ in range(count)] + + +@pytest.mark.parametrize( + ("route", "batched"), (("/v1/embeddings", False), ("/embeddings", True)), ids=("v1-string", "alias-batch") +) +@pytest.mark.parametrize( + ("credentials", "authorization"), + ( + ({"api_key": _DEPLOYMENT_KEY}, f"Bearer {_DEPLOYMENT_KEY}"), + ( + {"api_key": _DEPLOYMENT_KEY, "extra_headers": {"authorization": _CALLER_TOKEN, "X-Embed-Tenant": "t-1"}}, + _CALLER_TOKEN, + ), + ), + ids=("api-key", "lowercase-header-overrides-key"), +) +def test_regression_proxy_ollama_embedding_sends_deployment_credentials( + gateway: Gateway, route: str, batched: bool, credentials: Mapping[str, JsonValue], authorization: str +) -> None: + texts: Final = _texts(2 if batched else 1) + with wire_server(_embeddings) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_MODEL}", api_base=wire.url, **credentials) + response: Final = gateway.request("POST", route, {"model": model, "input": texts if batched else texts[0]}) + request: Final = _only(wire) + assert response.status_code == 200, response.text + assert _vectors(response) == [_vector(i) for i in range(len(texts))] + assert _values_of(request, "authorization") == [authorization], request.header_lines + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": texts}, request.body + expected_tenant: Final = ["t-1"] if "extra_headers" in credentials else [] + assert _values_of(request, "x-embed-tenant") == expected_tenant, request.header_lines + + +@pytest.mark.parametrize("caller_timeout", ({}, {"timeout": 0}), ids=("omitted", "zero")) +def test_regression_proxy_ollama_deployment_timeout_cuts_a_held_read( + gateway: Gateway, caller_timeout: Mapping[str, JsonValue] +) -> None: + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + release.wait(timeout=_DELAY_SECONDS) + return _embeddings(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_MODEL}", api_base=wire.url, timeout=0.5) + started: Final = time.monotonic() + try: + response: Final = gateway.request( + "POST", "/v1/embeddings", {"model": model, "input": _texts(1)[0], **caller_timeout} + ) + finally: + elapsed: Final = time.monotonic() - started + release.set() + _only(wire) + assert response.status_code == 408, response.text + assert_cut_short(elapsed, injected_delay=_DELAY_SECONDS) + + +def test_guard_proxy_keyless_ollama_embedding_sends_no_authorization(gateway: Gateway) -> None: + texts: Final = _texts(2) + with wire_server(_embeddings) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_MODEL}", api_base=f"{wire.url}/", api_key=None) + response: Final = gateway.request("POST", "/v1/embeddings", {"model": model, "input": texts}) + request: Final = _only(wire) + assert response.status_code == 200, response.text + assert _vectors(response) == [_vector(0), _vector(1)] + assert _values_of(request, "authorization") == [], request.header_lines + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": texts}, request.body + + +def test_guard_proxy_ollama_upstream_error_reaches_the_caller(gateway: Gateway) -> None: + marker: Final = f"scripted-bad-request-{uuid.uuid4().hex}" + + def respond(_: Request) -> Reply: + return Reply(status=400, body=json.dumps({"error": marker}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_MODEL}", api_base=wire.url) + response: Final = gateway.request("POST", "/v1/embeddings", {"model": model, "input": _texts(1)[0]}) + _only(wire) + assert response.status_code == 400, response.text + assert marker in response.text, response.text 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/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 2918f364c56..4e31122440d 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -60,6 +60,8 @@ _OWNED_MODEL: Final = "partner-claude" _DEPRECATED_OUTPUT_FORMAT: Final = "output_format: this field is deprecated, use output_config.format" _BAD_OUTPUT_CONFIG: Final = "output_config: Input should be a valid dictionary" _BAD_FORMAT: Final = "output_config.format: Input should be a json_schema object" +_BAD_OUTPUT_FORMAT: Final = "output_format: Input should be a valid dictionary" +_MISTYPED_OUTPUT_FORMAT: Final = "json_schema" _SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, @@ -94,6 +96,11 @@ _VALUES_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { "type": "json_schema", "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, } +_REQUIRED_ONLY_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer"}}, + "required": ["payload"], +} _ADDITIONAL_OBJECT_SCHEMA: Final[dict[str, JsonValue]] = { **_OPTIONAL_SCHEMA, "additionalProperties": {"type": "object"}, @@ -146,6 +153,7 @@ _STREAM_SHAPES: Final = ( "number_first", "bool_first", "null_first", + "named_wrapper", ) _HELD_STREAM_SHAPES: Final = frozenset( { @@ -166,12 +174,15 @@ _HELD_STREAM_SHAPES: Final = frozenset( "named_encoded", "undeclared_encoded", "encoded_deep", + "named_wrapper", } ) _AS_WRITTEN_SHAPES: Final = ("label_non_ascii", "label_two_objects_non_ascii") _TEXT_ANSWER_SHAPES: Final = frozenset({"undeclared_wrapper"}) _LABEL: Final = "RiskAssessment" _NON_ASCII_LABEL: Final = "Évaluation des risques 风险" +_NON_ASCII_SUMMARY: Final = "Hôtel 東京" +_LONE_SURROGATE: Final = "\ud800" _ENCODED_DEPTH: Final = 200_000 _FAMILY: Final = ( "anthropic_unmapped", @@ -181,6 +192,9 @@ _FAMILY: Final = ( "claude_platform_unmapped", "vertex_haiku", ) +_CUT_FAMILY: Final = ("anthropic_unmapped", "bedrock_invoke") +_CUT_EVENTS: Final = 3 +_OUTPUT_FORMAT_ROUTES: Final = ("anthropic_chat", "anthropic_messages", "vertex_chat", "vertex_messages") _FRAGMENT_SIZE: Final = 20 _TIMING_BYTES: Final = 512 * 1024 _TIMING_FRAGMENT: Final = 4096 @@ -228,9 +242,9 @@ def _nest(payload: JsonValue, depth: int) -> JsonValue: def _tool_input(shape: str, marker_hex: str) -> JsonValue: payload: Final = _payload(marker_hex) match shape: - case "plain": + case "plain" | "cut_plain": return payload - case "wrapped": + case "wrapped" | "cut_wrapped": return {"tool_name": _JSON_TOOL, "parameters": payload} case "named": return {"name": _JSON_TOOL, "parameters": payload} @@ -256,6 +270,14 @@ def _tool_input(shape: str, marker_hex: str) -> JsonValue: return {"tool_name": "other_tool", "parameters": payload} case "named_input": return {"name": _JSON_TOOL, "input": payload} + case "named_wrapper": + return {"name": _JSON_TOOL, "arguments": {"input": payload}} + case "required_only": + return {"payload": payload} + case "wrapped_non_ascii": + return {"tool_name": _JSON_TOOL, "parameters": _payload(f"{_NON_ASCII_SUMMARY} {marker_hex}")} + case "wrapped_lone_surrogate": + return {"tool_name": _JSON_TOOL, "parameters": _payload(marker_hex + _LONE_SURROGATE)} case "named_two_objects": return {"tool_name": _JSON_TOOL, "parameters": payload, "extra": {"note": marker_hex}} case "named_string": @@ -348,8 +370,13 @@ def _expected(shape: str, marker_hex: str) -> JsonValue: | "parameter_label" | "function_label" | "named_encoded" + | "named_wrapper" + | "cut_plain" + | "cut_wrapped" ): return _payload(marker_hex) + case "wrapped_non_ascii" | "wrapped_lone_surrogate": + return object_value(_tool_input(shape, marker_hex))["parameters"] case "large": return {**_payload(marker_hex), "summary": marker_hex + "x" * 5000} case "named_partial" | "inner_incomplete": @@ -423,9 +450,14 @@ def _aws_chunk(data: Mapping[str, JsonValue]) -> bytes: def _stream_reply( - message: dict[str, JsonValue], block: dict[str, JsonValue], deltas: tuple[dict[str, JsonValue], ...], framing: str + message: dict[str, JsonValue], + block: dict[str, JsonValue], + deltas: tuple[dict[str, JsonValue], ...], + framing: str, + cut: bool = False, ) -> Reply: - events: Final = _events(message, block, deltas) + complete: Final = _events(message, block, deltas) + events: Final = complete[:-_CUT_EVENTS] if cut else complete if framing == "aws": return Reply(content_type=_EVENT_STREAM, chunks=tuple(_aws_chunk(data) for _, data in events)) return Reply( @@ -445,6 +477,8 @@ def _error(status: int, message: str) -> Reply: def _format_fault(body: Mapping[str, JsonValue], vertex: bool) -> str | None: if vertex and "output_format" in body: return _DEPRECATED_OUTPUT_FORMAT + if body.get("output_format") is not None and not isinstance(body["output_format"], dict): + return _BAD_OUTPUT_FORMAT output_config: Final = body.get("output_config") if output_config is None: return None @@ -486,7 +520,7 @@ def _tool_reply(message_id: str, shape: str, marker_hex: str, user_tool: bool, f deltas: Final = tuple( {"type": "input_json_delta", "partial_json": piece} for piece in _fragments(partial, _fragment_size(shape)) ) - return _stream_reply(message, {**json_block, "input": {}}, deltas, framing) + return _stream_reply(message, {**json_block, "input": {}}, deltas, framing, cut=shape.startswith("cut_")) def _text_reply(message_id: str, text: str, framing: str | None) -> Reply: @@ -748,7 +782,9 @@ def _messages_stream_text(response: httpx.Response) -> JsonValue: return json.loads("".join(string_value(delta["text"]) for delta in deltas if delta["type"] == "text_delta")) -def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, backend: str) -> str: +def _vertex_deployment( + gateway: Gateway, scenario: Scenario, api_base: str, backend: str, **litellm_params: JsonValue +) -> str: return scenario.model( model=f"vertex_ai/{backend}", api_base=api_base, @@ -756,6 +792,7 @@ def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, back vertex_project=_PROJECT, vertex_location=_LOCATION, vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + **litellm_params, ) @@ -814,6 +851,33 @@ def _family_deployment(gateway: Gateway, scenario: Scenario, api_base: str, memb raise AssertionError(member) +def _output_format_deployment( + gateway: Gateway, scenario: Scenario, api_base: str, route: str, **litellm_params: JsonValue +) -> str: + if route.startswith("vertex_"): + return _vertex_deployment(gateway, scenario, api_base, _NATIVE_MODEL, **litellm_params) + return scenario.model( + model=f"anthropic/{_NATIVE_MODEL}", api_base=api_base, api_key=_ANTHROPIC_KEY, **litellm_params + ) + + +def _output_format_call( + gateway: Gateway, model: str, route: str, marker: str, output_format: JsonValue +) -> httpx.Response: + if route.endswith("_chat"): + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}], "output_format": output_format}, + ) + return gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, output_format=output_format), + headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA}, + ) + + def _owned_config(path: Path, gateway: Gateway, api_base: str) -> Path: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) path.write_text( @@ -1188,13 +1252,18 @@ def test_vertex_chat_json_schema_without_required_keys_unwraps_only_the_tool_nam @pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize( + ("shape", "schema"), + [("values", _VALUES_SCHEMA), ("required_only", _REQUIRED_ONLY_SCHEMA)], + ids=["declared_property", "required_only"], +) def test_vertex_chat_json_schema_that_declares_the_wrapper_key_keeps_the_answer_whole( - gateway: Gateway, stream: bool + gateway: Gateway, shape: str, schema: dict[str, JsonValue], stream: bool ) -> None: - marker, marker_hex = _marker("values") + marker, marker_hex = _marker(shape) response_format: Final[dict[str, JsonValue]] = { "type": "json_schema", - "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, + "json_schema": {"name": "valued", "schema": schema}, } with wire_server(_peer) as wire, gateway.scenario() as scenario: model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) @@ -1202,9 +1271,9 @@ def test_vertex_chat_json_schema_that_declares_the_wrapper_key_keeps_the_answer_ "POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format) ) content: Final = _stream_content(response) if stream else _chat_content(response) - assert content == _tool_input("values", marker_hex), response.text + assert content == _tool_input(shape, marker_hex), response.text if stream: - assert len(_content_frames(response)) >= _fragment_count("values", marker_hex), response.text + assert len(_content_frames(response)) >= _fragment_count(shape, marker_hex), response.text _one_delivery(wire, marker_hex) @@ -1265,6 +1334,22 @@ def test_vertex_chat_json_schema_unwraps_a_tool_named_envelope_only_around_one_o _one_delivery(wire, marker_hex) +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize("shape", ["wrapped_non_ascii", "wrapped_lone_surrogate"]) +def test_vertex_chat_json_schema_answer_keeps_non_ascii_text_and_escapes_a_lone_surrogate( + gateway: Gateway, shape: str, stream: bool +) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream)) + text: Final = _stream_text(response) if stream else string_value(_chat_message(response)["content"]) + assert json.loads(text) == _expected(shape, marker_hex), response.text + if shape == "wrapped_non_ascii": + assert _NON_ASCII_SUMMARY in text, text + _one_delivery(wire, marker_hex) + + @pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) @pytest.mark.parametrize("shape", ["name_property", "name_value_plain"]) def test_vertex_chat_json_schema_with_a_name_property_unwraps_only_the_tool_named_wrapper( @@ -1350,6 +1435,19 @@ def test_tool_path_family_stream_unwraps_the_envelope(gateway: Gateway, member: assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body +@pytest.mark.parametrize("shape", ["plain", "cut_plain", "cut_wrapped"]) +@pytest.mark.parametrize("member", _CUT_FAMILY) +def test_tool_path_stream_cut_before_the_tool_block_closes_still_delivers_the_answer( + gateway: Gateway, member: str, shape: str +) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _family_deployment(gateway, scenario, wire.url, member) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) + assert _stream_content(response) == _payload(marker_hex), response.text + _one_delivery(wire, marker_hex) + + @pytest.mark.parametrize("accepts_forced_tools", [True, False], ids=["accepts_forced_tools", "rejects_forced_tools"]) def test_vertex_chat_json_schema_forces_the_json_tool_only_when_the_model_accepts_forced_tools( gateway: Gateway, accepts_forced_tools: bool @@ -1865,6 +1963,74 @@ def test_vertex_messages_duplicate_output_format_keys_fold_the_last_value(gatewa assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT +@pytest.mark.parametrize("route", _OUTPUT_FORMAT_ROUTES) +def test_string_output_format_under_deployment_drop_params_is_dropped_before_the_peer( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker("plain_text") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _output_format_deployment(gateway, scenario, wire.url, route, drop_params=True) + response: Final = _output_format_call(gateway, model, route, marker, _MISTYPED_OUTPUT_FORMAT) + assert response.status_code == 200, response.text + assert f"plain text {marker_hex}" in response.text, response.text + sent: Final = _one_delivery(wire, marker_hex) + assert not {"output_format", "output_config"} & sent.body.keys(), sent.body + + +@pytest.mark.parametrize("route", ["anthropic_chat", "anthropic_messages"]) +def test_object_output_format_under_deployment_drop_params_still_reaches_anthropic( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _output_format_deployment(gateway, scenario, wire.url, route, drop_params=True) + response: Final = _output_format_call(gateway, model, route, marker, _OUTPUT_FORMAT) + assert response.status_code == 200, response.text + assert _one_delivery(wire, marker_hex).body["output_format"] == _OUTPUT_FORMAT + + +def test_vertex_messages_string_output_format_under_drop_params_keeps_the_callers_output_config_format( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL, drop_params=True) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body( + model, marker, output_format=_MISTYPED_OUTPUT_FORMAT, output_config={"format": _OUTPUT_FORMAT} + ), + ) + assert _messages_text(response) == _payload(marker_hex), response.text + assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT + + +@pytest.mark.parametrize("route", ["anthropic_chat", "anthropic_messages"]) +def test_string_output_format_without_drop_params_reaches_anthropic_and_returns_its_400( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker("plain_text") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _output_format_deployment(gateway, scenario, wire.url, route) + response: Final = _output_format_call(gateway, model, route, marker, _MISTYPED_OUTPUT_FORMAT) + assert response.status_code == 400, response.text + assert _BAD_OUTPUT_FORMAT in response.text, response.text + assert _one_delivery(wire, marker_hex).body["output_format"] == _MISTYPED_OUTPUT_FORMAT + + +def test_vertex_chat_string_output_format_without_drop_params_is_refused_without_reaching_the_peer( + gateway: Gateway, +) -> None: + marker, _ = _marker("plain_text") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _output_format_deployment(gateway, scenario, wire.url, "vertex_chat") + response: Final = _output_format_call(gateway, model, "vertex_chat", marker, _MISTYPED_OUTPUT_FORMAT) + assert response.status_code == 400, response.text + assert "needs `output_format` to be a JSON object" in response.text, response.text + assert _delivered(wire) == () + + @pytest.mark.parametrize( "schema", [ diff --git a/tests/integration/providers/test_vertex_partner_labels_wire.py b/tests/integration/providers/test_vertex_partner_labels_wire.py new file mode 100644 index 00000000000..ebdf3f2022e --- /dev/null +++ b/tests/integration/providers/test_vertex_partner_labels_wire.py @@ -0,0 +1,120 @@ +import base64 +import json +from typing import Final + +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-5-5" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_MODEL_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/{_BACKEND}" +_TARGET: Final = f"{_MODEL_PATH}:rawPredict" +_STREAM_TARGET: Final = f"{_MODEL_PATH}:streamRawPredict?alt=sse" +_LABELS_HEADER: Final = "x-vertex-ai-labels" +_REPLY_TEXT: Final = "labels control" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_LABELS: Final[dict[str, JsonValue]] = {"team": "triage", "env": "prod"} + +_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_vertex_labels", + "type": "message", + "role": "assistant", + "model": _BACKEND, + "content": [{"type": "text", "text": _REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, +} +_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "message_start", "message": {**_REPLY, "content": [], "stop_reason": None}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _REPLY_TEXT}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, +) +_SSE: Final = tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _EVENTS) + + +def _peer(request: Request) -> Reply: + if request.target == "/_oauth/token": + return Reply( + body=json.dumps({"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}).encode() + ) + if request.target == _STREAM_TARGET: + return Reply(content_type="text/event-stream", chunks=_SSE) + if request.target == _TARGET: + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + +def _deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, api_base), + num_retries=0, + ) + + +def _chat(gateway: Gateway, wire: Wire, scenario: Scenario, stream: bool, **extra: JsonValue) -> Request: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "Reply ok"}], + "max_tokens": 16, + "stream": stream, + **extra, + }, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + (sent,) = tuple(request for request in wire.drain() if request.target in (_TARGET, _STREAM_TARGET)) + assert sent.target == (_STREAM_TARGET if stream else _TARGET), sent.target + return sent + + +def _decoded_labels(sent: Request) -> JsonValue: + return json.loads(base64.b64decode(sent.headers[_LABELS_HEADER])) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_regression_vertex_claude_labels_move_from_body_to_header(gateway: Gateway, stream: bool) -> None: + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat(gateway, wire, scenario, stream, labels=_LABELS) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert _decoded_labels(sent) == _LABELS, sent.headers + + +def test_guard_vertex_claude_without_labels_sends_no_labels_header(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat(gateway, wire, scenario, False) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert _LABELS_HEADER not in sent.headers, sent.headers + + +def test_regression_vertex_claude_caller_labels_header_wins_over_labels_param(gateway: Gateway) -> None: + own_header: Final = base64.b64encode(json.dumps({"team": "caller"}).encode()).decode() + with gateway.scenario() as scenario, wire_server(_peer) as wire: + sent: Final = _chat( + gateway, wire, scenario, False, labels=_LABELS, extra_headers={"X-Vertex-AI-Labels": own_header} + ) + + assert "labels" not in _JSON_OBJECT.validate_json(sent.body), sent.body + assert sent.headers[_LABELS_HEADER] == own_header, sent.headers diff --git a/tests/integration/routing/test_batch_mode_deployments_wire.py b/tests/integration/routing/test_batch_mode_deployments_wire.py new file mode 100644 index 00000000000..7d702f499f5 --- /dev/null +++ b/tests/integration/routing/test_batch_mode_deployments_wire.py @@ -0,0 +1,426 @@ +import json +import uuid +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Final, Literal +from urllib.parse import urlsplit + +import pytest +from integration._support.client import ( + DEPLOYMENT_IDS, + Gateway, + Scenario, + eventually, + gateway_from_environment, + list_value, + object_value, + string_value, +) +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import owned_config +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.shares_worker_state + +_AZURE_API_VERSION: Final = "2024-10-21" +_BACKEND: Final = "gpt-5.4-mini" +_BATCH_ENDPOINT: Final = "/v1/batch" +_CHAT_ENDPOINT: Final = "/v1/chat/completions" +_RESPONSES_ENDPOINT: Final = "/v1/responses" +_TEAM_GROUP: Final = "batch-team-mix" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_ENDPOINTS: Final = TypeAdapter(list[str] | None) +_OWNED_PROXY_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +Provider = Literal["azure", "openai"] +Mode = Literal["chat", "batch"] + + +@dataclass(frozen=True, slots=True) +class _Sku: + provider: Provider + mode: Mode + key: str + weight: int | None + endpoints: JsonValue + team_id: str | None + + +@dataclass(frozen=True, slots=True) +class _Create: + status: int + body: str + + +@dataclass(frozen=True, slots=True) +class _TeamRig: + owned: OwnedProxy + wire: Wire + team_key: str + other_team_key: str + no_team_key: str + chat: _Sku + batch: _Sku + + +def _sku( + provider: Provider, + mode: Mode, + *, + weight: int | None = None, + endpoints: JsonValue = None, + team_id: str | None = None, +) -> _Sku: + return _Sku(provider, mode, f"{mode}-sku-{uuid.uuid4().hex}", weight, endpoints, team_id) + + +def _serves_batches(key: str) -> bool: + return key.startswith("batch-sku-") + + +def _deployment(name: str, wire: Wire, sku: _Sku) -> dict[str, JsonValue]: + endpoint: Final[dict[str, JsonValue]] = ( + {"model": f"azure/{_BACKEND}", "api_base": wire.url, "api_version": _AZURE_API_VERSION} + if sku.provider == "azure" + else {"model": f"openai/{_BACKEND}", "api_base": wire.url + "/v1"} + ) + weight: Final[dict[str, JsonValue]] = {} if sku.weight is None else {"weight": sku.weight} + endpoints: Final[dict[str, JsonValue]] = {} if sku.endpoints is None else {"supported_endpoints": sku.endpoints} + team: Final[dict[str, JsonValue]] = {} if sku.team_id is None else {"team_id": sku.team_id} + return { + "model_name": name, + "litellm_params": {**endpoint, "api_key": sku.key, **weight}, + "model_info": {"mode": sku.mode, **endpoints, **team}, + } + + +def _json_reply(body: Mapping[str, JsonValue], status: int = 200) -> Reply: + return Reply(status=status, body=json.dumps(body).encode()) + + +def _provider_key(request: Request) -> str: + return request.headers.get("api-key") or request.headers.get("authorization", "").removeprefix("Bearer ") + + +def _is_azure(request: Request) -> bool: + return urlsplit(request.target).path.startswith("/openai/") + + +def _posted_to(request: Request, suffix: str) -> bool: + return request.method == "POST" and urlsplit(request.target).path.endswith(suffix) + + +def _vendor(request: Request) -> Reply: + key: Final = _provider_key(request) + if _posted_to(request, "/files"): + return _json_reply( + { + "id": f"file-{key}", + "object": "file", + "bytes": len(request.body), + "created_at": 1700000000, + "filename": "in.jsonl", + "purpose": "batch", + "status": "processed", + } + ) + if _posted_to(request, "/batches") and _is_azure(request) and not _serves_batches(key): + return _json_reply( + {"error": {"code": "invalid_deployment_type", "message": "The deployment does not support batch jobs."}}, + 400, + ) + if _posted_to(request, "/batches"): + return _json_reply( + { + "id": f"batch-{uuid.uuid4().hex}", + "object": "batch", + "endpoint": _CHAT_ENDPOINT, + "input_file_id": f"file-{key}", + "status": "validating", + "created_at": 1700000000, + "completion_window": "24h", + } + ) + if _posted_to(request, "/chat/completions") and _is_azure(request) and _serves_batches(key): + return _json_reply( + {"error": {"code": "OperationNotSupported", "message": "Batch deployments do not serve chat."}}, 400 + ) + if _posted_to(request, "/chat/completions"): + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _BACKEND, "served", stream=False) + return _json_reply({"data": [], "object": "list"}) + + +def _deployment_ids_on_a_worker(gateway: Gateway, name: str) -> tuple[str, frozenset[str]]: + response: Final = gateway.request( + "GET", + "/health/integration-deployment-ids", + params={"model_name": name}, + headers={"connection": "close"}, + ) + assert response.status_code == 200, response.text + return response.headers["x-integration-worker-pid"], frozenset( + DEPLOYMENT_IDS.validate_json(response.content)["model_ids"] + ) + + +def _wait_until_every_worker_serves(gateway: Gateway, name: str, identities: frozenset[str]) -> None: + from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + + def answers() -> tuple[tuple[str, frozenset[str]], ...]: + with ThreadPoolExecutor(max_workers=8) as pool: + return tuple(pool.map(lambda _: _deployment_ids_on_a_worker(gateway, name), range(16))) + + eventually( + answers, + lambda seen: ( + len({pid for pid, _ in seen}) >= gateway.workers and all(identities <= served for _, served in seen) + ), + seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ) + + +def _add_deployment(scenario: Scenario, deployment: Mapping[str, JsonValue]) -> str: + created: Final = scenario.gateway.post("/model/new", deployment) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return identity + + +def _group(scenario: Scenario, wire: Wire, skus: Sequence[_Sku]) -> str: + name: Final = f"integration-batch-{uuid.uuid4().hex}" + identities: Final = frozenset(_add_deployment(scenario, _deployment(name, wire, sku)) for sku in skus) + _wait_until_every_worker_serves(scenario.gateway, name, identities) + return name + + +def _managed_upload(gateway: Gateway, group: str, key: str) -> str: + line: Final = {"custom_id": "r1", "method": "POST", "url": _CHAT_ENDPOINT, "body": {"model": group}} + uploaded: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": group}, + {"file": ("in.jsonl", json.dumps(line).encode() + b"\n", "application/jsonl")}, + key=key, + ) + assert uploaded.status_code == 200, uploaded.text + return string_value(_JSON_OBJECT.validate_json(uploaded.content)["id"]) + + +def _create(gateway: Gateway, managed_file: str, key: str) -> _Create: + response: Final = gateway.request( + "POST", + "/v1/batches", + {"input_file_id": managed_file, "endpoint": _CHAT_ENDPOINT, "completion_window": "24h"}, + key=key, + ) + return _Create(response.status_code, response.text) + + +def _chat(gateway: Gateway, group: str) -> int: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": group, "messages": [{"role": "user", "content": "hello"}], "cache": {"no-cache": True}}, + ).status_code + + +def _keys_posted_to(requests: Sequence[Request], suffix: str) -> tuple[str, ...]: + return tuple(_provider_key(request) for request in requests if _posted_to(request, suffix)) + + +def _assert_creates_reach_only(wire: Wire, creates: Sequence[_Create], sku: _Sku) -> None: + assert [create.status for create in creates] == [200] * len(creates), creates + assert _keys_posted_to(wire.drain(), "/batches") == (sku.key,) * len(creates) + + +@pytest.mark.parametrize( + ("provider", "creates_reach"), + [("azure", "batch"), ("openai", "chat")], + ids=["azure-skips-chat", "openai-keeps-chat"], +) +def test_a_batch_create_on_a_mixed_group_reaches_the_deployment_its_provider_allows( + gateway: Gateway, provider: Provider, creates_reach: Mode +) -> None: + chat: Final = _sku(provider, "chat", weight=1) + batch: Final = _sku(provider, "batch", weight=0) + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + group: Final = _group(scenario, wire, (chat, batch)) + managed_file: Final = _managed_upload(gateway, group, gateway.key) + assert set(_keys_posted_to(wire.drain(), "/files")) == {chat.key, batch.key} + + creates: Final = tuple(_create(gateway, managed_file, gateway.key) for _ in range(3)) + + _assert_creates_reach_only(wire, creates, batch if creates_reach == "batch" else chat) + + +def test_concurrent_batch_creates_and_chats_each_reach_their_own_deployment(gateway: Gateway) -> None: + chat: Final = _sku("azure", "chat", weight=1) + batch: Final = _sku("azure", "batch", weight=0) + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + group: Final = _group(scenario, wire, (chat, batch)) + managed_file: Final = _managed_upload(gateway, group, gateway.key) + wire.drain() + + with ThreadPoolExecutor(max_workers=16) as pool: + burst: Final = tuple( + pool.map( + lambda turn: _create(gateway, managed_file, gateway.key) if turn % 2 else _chat(gateway, group), + range(24), + ) + ) + after: Final = tuple(_chat(gateway, group) for _ in range(6)) + received: Final = wire.drain() + + chat_statuses: Final = tuple(result for result in burst if isinstance(result, int)) + creates: Final = tuple(result for result in burst if isinstance(result, _Create)) + assert chat_statuses + after == (200,) * 18, (burst, after) + assert [create.status for create in creates] == [200] * 12, creates + assert _keys_posted_to(received, "/chat/completions") == (chat.key,) * 18 + assert _keys_posted_to(received, "/batches") == (batch.key,) * 12 + + +def _group_endpoints(gateway: Gateway, group: str, key: str | None = None) -> list[str] | None: + response: Final = gateway.request("GET", "/model_group/info", params={"model_group": group}, key=key) + assert response.status_code == 200, response.text + (entry,) = list_value(_JSON_OBJECT.validate_json(response.content)["data"]) + return _ENDPOINTS.validate_python(object_value(entry).get("supported_endpoints")) + + +@pytest.mark.parametrize( + ("shape", "expected"), + [ + pytest.param( + ( + ("azure", "chat", [_CHAT_ENDPOINT, _RESPONSES_ENDPOINT, _BATCH_ENDPOINT]), + ("azure", "batch", [_CHAT_ENDPOINT]), + ), + [_BATCH_ENDPOINT, _CHAT_ENDPOINT, _RESPONSES_ENDPOINT], + id="azure-chat-and-batch", + ), + pytest.param( + (("azure", "chat", [_CHAT_ENDPOINT, _BATCH_ENDPOINT]),), + [_CHAT_ENDPOINT], + id="azure-chat-only", + ), + pytest.param( + (("azure", "batch", [_CHAT_ENDPOINT, _RESPONSES_ENDPOINT]),), + [_BATCH_ENDPOINT], + id="azure-batch-only", + ), + pytest.param( + ( + ("openai", "chat", [_CHAT_ENDPOINT, _BATCH_ENDPOINT]), + ("openai", "batch", [_RESPONSES_ENDPOINT]), + ), + [_BATCH_ENDPOINT, _CHAT_ENDPOINT, _RESPONSES_ENDPOINT], + id="openai-chat-and-batch", + ), + ], +) +def test_model_group_info_lists_the_endpoints_its_deployments_serve( + gateway: Gateway, shape: tuple[tuple[Provider, Mode, list[str]], ...], expected: list[str] +) -> None: + skus: Final = tuple(_sku(provider, mode, endpoints=endpoints) for provider, mode, endpoints in shape) + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + group: Final = _group(scenario, wire, skus) + + listed: Final = _group_endpoints(gateway, group) + + assert listed is not None and sorted(listed) == expected, listed + + +@pytest.mark.parametrize("malformed", ["/v1/batch", [7], {"batch": True}], ids=["string", "number-list", "object"]) +def test_a_malformed_endpoint_list_keeps_group_info_and_chat_answering(gateway: Gateway, malformed: JsonValue) -> None: + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + group: Final = _group(scenario, wire, (_sku("openai", "chat", endpoints=malformed),)) + + _group_endpoints(gateway, group) + + assert _chat(gateway, group) == 200 + + +@pytest.fixture(scope="module") +def team_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_TeamRig]: + directory: Final = tmp_path_factory.mktemp("batch-team") + team_id: Final = f"integration-batch-team-{uuid.uuid4().hex}" + chat: Final = _sku("azure", "chat", endpoints=[_CHAT_ENDPOINT]) + batch: Final = _sku("azure", "batch", team_id=team_id) + with gateway_from_environment() as gateway, wire_server(_vendor) as wire: + config: Final = owned_config( + directory, + [_deployment(_TEAM_GROUP, wire, chat), _deployment(_TEAM_GROUP, wire, batch)], + litellm_settings={"public_model_groups": [_TEAM_GROUP]}, + router_settings={"routing_strategy": "usage-based-routing"}, + ) + with owned_proxy_process(gateway, directory, {}, config=config) as owned, owned.gateway.scenario() as scenario: + scenario.team(team_id=team_id) + yield _TeamRig( + owned=owned, + wire=wire, + team_key=scenario.key(team_id=team_id), + other_team_key=scenario.key(team_id=scenario.team()), + no_team_key=scenario.key(), + chat=chat, + batch=batch, + ) + + +def _caller_key(rig: _TeamRig, caller: str) -> str: + return {"team": rig.team_key, "other-team": rig.other_team_key, "no-team": rig.no_team_key}[caller] + + +@_OWNED_PROXY_TIMEOUT +@pytest.mark.parametrize( + ("caller", "lists_batch"), + [("team", True), ("other-team", False), ("no-team", False)], + ids=["team", "other-team", "no-team"], +) +def test_model_group_info_lists_a_team_owned_batch_endpoint_only_to_that_team( + team_rig: _TeamRig, caller: str, lists_batch: bool +) -> None: + listed: Final = _group_endpoints(team_rig.owned.gateway, _TEAM_GROUP, _caller_key(team_rig, caller)) + + assert listed is not None and _CHAT_ENDPOINT in listed, listed + assert (_BATCH_ENDPOINT in listed) is lists_batch, listed + + +@_OWNED_PROXY_TIMEOUT +def test_public_model_hub_never_lists_a_team_owned_batch_endpoint(team_rig: _TeamRig) -> None: + response: Final = team_rig.owned.gateway.client.get("/public/model_hub") + assert response.status_code == 200, response.text + (entry,) = [ + object_value(row) + for row in _JSON_LIST.validate_json(response.content) + if object_value(row)["model_group"] == _TEAM_GROUP + ] + + assert _BATCH_ENDPOINT not in (_ENDPOINTS.validate_python(entry.get("supported_endpoints")) or []), entry + + +@_OWNED_PROXY_TIMEOUT +def test_a_legacy_strategy_sends_a_team_batch_create_to_its_batch_deployment(team_rig: _TeamRig) -> None: + gateway: Final = team_rig.owned.gateway + team_rig.wire.drain() + managed_file: Final = _managed_upload(gateway, _TEAM_GROUP, team_rig.team_key) + assert team_rig.batch.key in _keys_posted_to(team_rig.wire.drain(), "/files") + + creates: Final = tuple(_create(gateway, managed_file, team_rig.team_key) for _ in range(3)) + + _assert_creates_reach_only(team_rig.wire, creates, team_rig.batch) + + +@_OWNED_PROXY_TIMEOUT +@pytest.mark.parametrize("caller", ["other-team", "no-team"]) +def test_a_legacy_strategy_keeps_other_callers_off_a_team_batch_deployment(team_rig: _TeamRig, caller: str) -> None: + gateway: Final = team_rig.owned.gateway + key: Final = _caller_key(team_rig, caller) + team_rig.wire.drain() + managed_file: Final = _managed_upload(gateway, _TEAM_GROUP, key) + team_rig.wire.drain() + + tuple(_create(gateway, managed_file, key) for _ in range(3)) + + assert _keys_posted_to(team_rig.wire.drain(), "/batches") == (team_rig.chat.key,) * 3 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/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py index 9fc21843b04..11798542088 100644 --- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py +++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py @@ -512,12 +512,55 @@ def test_other_prioritized_endpoints_are_served_once_and_billed( marker: Final = new_marker() answer: Final = call(gateway, endpoint, model, marker) identity: Final = assert_served(answer, marker) - if endpoint == "queue": + if endpoint in ("queue", "queue-stream"): assert answer.headers.get("x-litellm-priority") == "1", answer.headers assert_once(received_markers(wire), (marker,)) spend_row_lands(identity) +@pytest.mark.parametrize("field", ({}, {"priority": "high", "drop_params": True}), ids=("missing", "dropped")) +def test_queue_route_refuses_a_request_it_cannot_queue_before_the_upstream( + gateway: Gateway, rig_model: tuple[str, Wire], field: Mapping[str, JsonValue] +) -> None: + model, wire = rig_model + marker: Final = new_marker() + answer: Final = post(gateway, "/queue/chat/completions", chat_body(model, marker, **field)) + assert answer.status == 400, (answer.status, answer.text) + error: Final = JSON_OBJECT.validate_json(answer.text)["error"] + assert isinstance(error, dict), answer.text + assert (error["type"], error["param"]) == ("invalid_request_error", "priority"), answer.text + assert_never(received_markers(wire), (marker,)) + + +def anthropic_error(status: int) -> Reply: + body: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": f"scripted {status}"}} + return Reply(status=status, body=json.dumps(body).encode()) + + +@pytest.mark.parametrize( + ("status", "provider_model", "base_path", "respond"), + ( + (400, "openai/gpt-4o-mini", "/v1", answering_model_discovery(lambda _: openai_error(400))), + (402, "anthropic/claude-sonnet-5-5", "", lambda _: anthropic_error(402)), + ), + ids=("openai-400", "anthropic-402"), +) +def test_queue_route_answers_a_provider_bad_request_with_its_own_status( + gateway: Gateway, status: int, provider_model: str, base_path: str, respond: Callable[[Request], Reply] +) -> None: + marker: Final = new_marker() + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=provider_model, api_base=wire.url + base_path) + answer: Final = post(gateway, "/queue/chat/completions", chat_body(model, marker, priority=1)) + assert marker in received_markers(wire) + assert answer.status == status, (answer.status, answer.text) + error: Final = JSON_OBJECT.validate_json(answer.text)["error"] + assert isinstance(error, dict), answer.text + assert error["type"] == "invalid_request_error", answer.text + assert f"scripted {status}" in str(error["message"]), answer.text + assert "Authentication Error" not in str(error["message"]), answer.text + + def upstream_body_for(wire: Wire, marker: str) -> dict[str, JsonValue]: (upstream_request,) = tuple(request for request in wire.drain() if marker.encode() in request.body) return JSON_OBJECT.validate_json(upstream_request.body) @@ -751,6 +794,19 @@ def test_waiter_proceeds_once_the_entry_ahead_is_cleared(pair: Pair) -> None: assert_drained(pair.cache, group) +@pytest.mark.timeout(PAIR_CELL_TIMEOUT) +def test_request_pinned_to_a_deployment_id_is_served_ahead_of_an_earlier_queue_entry(pair: Pair) -> None: + deployment: Final = f"{PAIR_GROUPS[0]}-deployment" + write_queue(pair.cache, deployment, [GHOST_ENTRY]) + marker: Final = new_marker() + answer: Final = post(pair.second, "/v1/chat/completions", chat_body(deployment, marker, priority=2)) + assert_served(answer, marker) + assert_prioritized(answer) + assert answer.seconds < PROMPT_SECONDS, answer.seconds + assert queue_entries(pair.cache, deployment) == [GHOST_ENTRY] + assert_once(received_markers(pair.wire), (marker,)) + + def pinned_connection(gateway: Gateway) -> http.client.HTTPConnection: address: Final = urlsplit(str(gateway.client.base_url)) assert address.hostname is not None and address.port is not None, address @@ -1313,8 +1369,7 @@ def test_queue_route_without_priority_takes_the_default(defaulted: Defaulted, en marker: Final = new_marker() answer: Final = call(defaulted.pair.first, endpoint, group, marker, priority=None) assert_served(answer, marker) - if endpoint == "queue": - assert answer.headers.get("x-litellm-priority") == str(DEFAULT_PRIORITY), answer.headers + assert answer.headers.get("x-litellm-priority") == str(DEFAULT_PRIORITY), answer.headers assert "priority" not in upstream_body_for(defaulted.pair.wire, marker) diff --git a/tests/integration/routing/test_silent_model_shadow_wire.py b/tests/integration/routing/test_silent_model_shadow_wire.py new file mode 100644 index 00000000000..b00fbca645e --- /dev/null +++ b/tests/integration/routing/test_silent_model_shadow_wire.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import json +import re +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import unquote + +import httpx +import pytest +from integration._support.anthropic_sse import parse_sse +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.openai_wire import chat_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +Shape = Literal["single", "pair"] + +_MODEL: Final = "gpt-5.6" +_API_KEY: Final = "synthetic-shadow-key" +_ROUTER_ONLY_KEYS: Final = ("silent_model", "include_fallback_errors") +_ANSWER: Final = "answered by" +_NONCE: Final = re.compile(r"nonce=([0-9a-f]{32})") +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_CHOICES: Final = TypeAdapter(list[dict[str, JsonValue]]) +_FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"}) +_SPEND_ROW_SECONDS: Final = 70 +_TARGETS: Final[Mapping[Shape, tuple[str, ...]]] = MappingProxyType( + {"single": ("shadow-a",), "pair": ("shadow-a", "shadow-b")} +) +_CASES: Final = ( + pytest.param("single", False, id="non-stream-single"), + pytest.param("single", True, id="stream-single"), + pytest.param("pair", True, id="stream-pair"), +) + + +def _prompt(nonce: str) -> str: + return f"which deployment answers this? nonce={nonce}" + + +def _nonce_of(raw: str) -> str: + found: Final = _NONCE.search(raw) + assert found is not None, raw + return found.group(1) + + +def _identity(segment: str, nonce: str) -> str: + return f"chatcmpl-{segment}-{nonce}" + + +def _request(model: str, nonce: str, *, stream: bool) -> dict[str, JsonValue]: + body: Final[dict[str, JsonValue]] = { + "model": model, + "stream": stream, + "messages": [{"role": "user", "content": _prompt(nonce)}], + } + return {**body, "stream_options": {"include_usage": True}} if stream else body + + +def _segment_of(request: Request) -> str: + return unquote(request.target).lstrip("/").partition("/")[0] + + +def _peer(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + segment: Final = _segment_of(request) + body: Final = _JSON.validate_json(request.body) + return chat_reply( + _identity(segment, _nonce_of(request.body.decode())), + _MODEL, + f"{_ANSWER} {segment}", + stream=body.get("stream") is True, + ) + + +@dataclass(frozen=True, slots=True) +class _Posted: + segment: str + streamed: bool + leaked: tuple[str, ...] + + +def _posted(wire: Wire, nonce: str) -> tuple[_Posted, ...]: + return tuple( + _Posted( + _segment_of(request), + _JSON.validate_json(request.body).get("stream") is True, + tuple(key for key in _ROUTER_ONLY_KEYS if key in request.body.decode()), + ) + for request in wire.drain() + if request.method == "POST" and nonce in request.body.decode() + ) + + +def _delta_text(data: Mapping[str, JsonValue]) -> str: + deltas: Final = ( + _JSON.validate_python(choice.get("delta") or {}) for choice in _CHOICES.validate_python(data["choices"]) + ) + return "".join(str(delta.get("content") or "") for delta in deltas) + + +def _served(response: httpx.Response, *, stream: bool) -> tuple[frozenset[str], str]: + if not stream: + body: Final = _JSON.validate_json(response.content) + message: Final = _JSON.validate_python(_CHOICES.validate_python(body["choices"])[0]["message"]) + return frozenset({str(body["id"])}), str(message["content"]) + events: Final = tuple(event.data for event in parse_sse(response.text)) + return frozenset(str(data["id"]) for data in events), "".join(_delta_text(data) for data in events) + + +@dataclass(frozen=True, slots=True) +class _SpendRow: + request_id: str + model_group: str + spend: float + + +def _spend_row(gateway: Gateway, request_id: str) -> tuple[_SpendRow, ...]: + rows: Final = gateway.rows( + 'SELECT model_group, spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,) + ) + return tuple(_SpendRow(request_id, str(row["model_group"]), float(str(row["spend"]))) for row in rows) + + +def _spend_rows(gateway: Gateway, request_ids: tuple[str, ...]) -> tuple[_SpendRow, ...]: + return tuple(chain.from_iterable(_spend_row(gateway, request_id) for request_id in request_ids)) + + +@dataclass(frozen=True, slots=True) +class _Mirrors: + gateway: Gateway + wire: Wire + groups: Mapping[str, str] + + +@pytest.fixture(scope="module") +def mirrors() -> Iterator[_Mirrors]: + with gateway_from_environment() as gateway, wire_server(_peer) as wire, gateway.scenario() as scenario: + + def register(segment: str, silent_model: JsonValue) -> str: + return scenario.model( + model=f"openai/{_MODEL}", api_key=_API_KEY, api_base=f"{wire.url}/{segment}", silent_model=silent_model + ) + + shadow_a: Final = register("shadow-a", None) + shadow_b: Final = register("shadow-b", None) + groups: Final = MappingProxyType( + { + "shadow-a": shadow_a, + "shadow-b": shadow_b, + "primary-single": register("primary-single", shadow_a), + "primary-pair": register("primary-pair", [shadow_a, shadow_b]), + } + ) + wire.drain() + yield _Mirrors(gateway, wire, groups) + + +@pytest.mark.parametrize(("shape", "stream"), _CASES) +def test_each_shadow_target_gets_one_request_in_the_callers_stream_mode( + mirrors: _Mirrors, shape: Shape, stream: bool +) -> None: + nonce: Final = uuid.uuid4().hex + primary: Final = f"primary-{shape}" + targets: Final = _TARGETS[shape] + response: Final = mirrors.gateway.request( + "POST", + "/v1/chat/completions", + _request(mirrors.groups[primary], nonce, stream=stream), + headers=_FRESH_CONNECTION, + ) + assert response.status_code == 200, response.text + assert _served(response, stream=stream) == (frozenset({_identity(primary, nonce)}), f"{_ANSWER} {primary}") + + expected_groups: Final = MappingProxyType( + {_identity(segment, nonce): mirrors.groups[segment] for segment in (primary, *targets)} + ) + rows: Final = eventually( + lambda: _spend_rows(mirrors.gateway, tuple(expected_groups)), + lambda found: len(found) == len(expected_groups), + seconds=_SPEND_ROW_SECONDS, + ) + assert {row.request_id: row.model_group for row in rows} == expected_groups, rows + assert all(row.spend > 0 for row in rows), rows + + posted: Final = _posted(mirrors.wire, nonce) + assert sorted(item.segment for item in posted) == sorted((primary, *targets)), posted + assert all(item.streamed is stream for item in posted), posted + assert all(item.leaked == () for item in posted), posted diff --git a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py new file mode 100644 index 00000000000..da091b1c12e --- /dev/null +++ b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py @@ -0,0 +1,146 @@ +import base64 +import json +import uuid +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + +_JSON_TOOL: Final = "json_tool_call" +_FRAGMENT_SIZE: Final = 20 +_CUT_EVENTS: Final = 3 +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], + "additionalProperties": False, +} +_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "risk", "schema": _SCHEMA}, +} +_ANTHROPIC_MODEL: Final = "anthropic/claude-scripted-9" +_BEDROCK_INVOKE_MODEL: Final = "bedrock/invoke/anthropic.claude-sonnet-5" + + +def _payload(marker_hex: str) -> dict[str, JsonValue]: + return {"tier": 2, "summary": marker_hex} + + +def _tool_input(shape: str, marker_hex: str) -> dict[str, JsonValue]: + if shape == "wrapped": + return {"tool_name": _JSON_TOOL, "parameters": _payload(marker_hex)} + return _payload(marker_hex) + + +def _events(marker_hex: str, shape: str, cut: bool) -> tuple[tuple[str, dict[str, JsonValue]], ...]: + text: Final = json.dumps(_tool_input(shape, marker_hex)) + deltas: Final = tuple( + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": text[start : start + _FRAGMENT_SIZE]}, + }, + ) + for start in range(0, len(text), _FRAGMENT_SIZE) + ) + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{marker_hex}", + "type": "message", + "role": "assistant", + "model": "claude-scripted-9", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + } + complete: Final = ( + ("message_start", {"type": "message_start", "message": message}), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": f"toolu_{marker_hex}", "name": _JSON_TOOL, "input": {}}, + }, + ), + *deltas, + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + return complete[:-_CUT_EVENTS] if cut else complete + + +def _aws_chunk(data: Mapping[str, JsonValue]) -> bytes: + payload: Final = json.dumps({"bytes": base64.b64encode(json.dumps(data).encode()).decode()}).encode() + return aws_event_stream_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, payload + ) + + +def _peer(marker_hex: str, shape: str, cut: bool) -> Callable[[Request], Reply]: + events: Final = _events(marker_hex, shape, cut) + + def respond(request: Request) -> Reply: + assert marker_hex in request.body.decode(), request.body + if request.target.endswith("/invoke-with-response-stream"): + return Reply( + content_type="application/vnd.amazon.eventstream", + chunks=tuple(_aws_chunk(data) for _, data in events), + ) + assert request.target.endswith("/v1/messages"), request.target + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() for name, data in events), + ) + + return respond + + +def _credentials(model: str) -> dict[str, str]: + if model == _ANTHROPIC_MODEL: + return {"api_key": "synthetic-anthropic-key"} + return {"api_key": "synthetic-bedrock-bearer", "aws_region_name": "us-east-1"} + + +@pytest.mark.parametrize("model", [_ANTHROPIC_MODEL, _BEDROCK_INVOKE_MODEL], ids=["anthropic", "bedrock_invoke"]) +@pytest.mark.parametrize( + ("shape", "cut"), + [("plain", False), ("plain", True), ("wrapped", True)], + ids=["plain_complete", "plain_cut", "wrapped_cut"], +) +def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whole_answer( + model: str, shape: str, cut: bool +) -> None: + marker_hex: Final = uuid.uuid4().hex + with wire_server(_peer(marker_hex, shape, cut)) as wire: + response: Final = litellm.completion( + model=model, + messages=[{"role": "user", "content": f"Rate the risk {marker_hex}"}], + response_format=_RESPONSE_FORMAT, + max_tokens=64, + stream=True, + api_base=wire.url, + num_retries=0, + **_credentials(model), + ) + assert isinstance(response, CustomStreamWrapper), type(response) + text: Final = "".join(str(chunk.choices[0].delta.content or "") for chunk in tuple(response)) + assert len(wire.drain()) == 1 + assert json.loads(text) == _payload(marker_hex), text diff --git a/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py b/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py index 6b2f4cb8eea..2988f9b5f53 100644 --- a/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py +++ b/tests/integration/sdk/test_bedrock_converse_redacted_thinking_router_sdk.py @@ -88,3 +88,42 @@ def test_sync_router_reads_base_model_to_replay_redacted_thinking_to_an_applicat assistant: Final = next(message for message in received[-1].messages if message.role == "assistant") redacted: Final = ({"reasoningContent": {"redactedContent": REDACTED}},) if replays_redacted else () assert assistant.content == (*redacted, {"text": FIRST_ANSWER}) + + +def test_sync_router_replays_a_block_in_content_and_thinking_blocks_once() -> None: + redacted: Final[dict[str, JsonValue]] = {"type": "redacted_thinking", "data": REDACTED} + with wire_server(_respond) as wire: + router: Final = Router( + model_list=[ + { + "model_name": "claude-profile", + "litellm_params": { + "model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}", + "api_base": wire.url, + "api_key": "synthetic-bedrock-bearer", + "aws_region_name": "us-east-1", + }, + "model_info": {"base_model": "bedrock/us.anthropic.claude-sonnet-4-6"}, + } + ], + num_retries=0, + ) + + response: Final = router.completion( + model="claude-profile", + max_tokens=256, + messages=[ + {"role": "user", "content": "What's the weather in Paris?"}, + { + "role": "assistant", + "content": [redacted, {"type": "text", "text": FIRST_ANSWER}], + "thinking_blocks": [redacted], + }, + {"role": "user", "content": "Reply with the word ok."}, + ], + ) + + assert response.choices[0].message.content == FINAL_ANSWER + received: Final = tuple(_ConverseRequest.model_validate_json(request.body) for request in wire.drain()) + assistant: Final = next(message for message in received[-1].messages if message.role == "assistant") + assert assistant.content == ({"reasoningContent": {"redactedContent": REDACTED}}, {"text": FIRST_ANSWER}) diff --git a/tests/integration/sdk/test_ollama_embedding_wire.py b/tests/integration/sdk/test_ollama_embedding_wire.py new file mode 100644 index 00000000000..21b803b8728 --- /dev/null +++ b/tests/integration/sdk/test_ollama_embedding_wire.py @@ -0,0 +1,218 @@ +import json +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import JSON_OBJECT +from integration._support.timing import assert_cut_short +from integration._support.wire import Reply, Request, Wire, wire_server + +import litellm +from litellm.types.utils import EmbeddingResponse + +_MODEL: Final = "nomic-embed-text" +_EXPLICIT_KEY: Final = "synthetic-ollama-explicit-key" +_ENV_KEY: Final = "synthetic-ollama-env-key" +_GLOBAL_KEY: Final = "synthetic-global-key-never-sent" +_CALLER_TOKEN: Final = "Bearer synthetic-caller-token" +_DELAY_SECONDS: Final = 5.0 +_SHORT_DELAY_SECONDS: Final = 1.0 + +Mode = Literal["sync", "async"] +_MODES: Final[tuple[Mode, ...]] = ("sync", "async") + + +def _vector(index: int) -> list[float]: + return [0.25, float(index)] + + +def _embeddings(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + count: Final = len(body["input"]) if isinstance(body["input"], list) else 1 + return Reply( + body=json.dumps({"embeddings": [_vector(index) for index in range(count)], "prompt_eval_count": 3}).encode() + ) + + +def _held(release: threading.Event, seconds: float) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + release.wait(timeout=seconds) + return _embeddings(request) + + return respond + + +@pytest.fixture(autouse=True) +def _no_ambient_ollama_settings(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + + +@contextmanager +def _peer(respond: Callable[[Request], Reply] = _embeddings) -> Iterator[Wire]: + with wire_server(respond) as wire: + yield wire + + +async def _call(mode: Mode, **kwargs: object) -> EmbeddingResponse: + if mode == "sync": + return litellm.embedding(model=f"ollama/{_MODEL}", **kwargs) + return await litellm.aembedding(model=f"ollama/{_MODEL}", **kwargs) + + +async def _seconds_until_timeout(mode: Mode, **kwargs: object) -> float: + started: Final = time.monotonic() + with pytest.raises(litellm.Timeout): + await _call(mode, **kwargs) + return time.monotonic() - started + + +def _only(wire: Wire) -> Request: + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [("POST", "/api/embed")], [ + (request.method, request.target) for request in received + ] + return received[0] + + +def _values_of(request: Request, name: str) -> list[str]: + return [value for header, value in request.header_lines if header.lower() == name.lower()] + + +def _text() -> str: + return f"ollama embedding {uuid.uuid4().hex}" + + +@pytest.mark.parametrize("mode", _MODES) +@pytest.mark.parametrize( + ("explicit", "env", "expected"), + ( + (_EXPLICIT_KEY, None, _EXPLICIT_KEY), + (None, _ENV_KEY, _ENV_KEY), + (_EXPLICIT_KEY, _ENV_KEY, _EXPLICIT_KEY), + ), + ids=("explicit-key", "env-key", "explicit-beats-env"), +) +async def test_regression_embedding_sends_the_resolved_ollama_key_as_one_bearer_header( + monkeypatch: pytest.MonkeyPatch, mode: Mode, explicit: str | None, env: str | None, expected: str +) -> None: + if env is not None: + monkeypatch.setenv("OLLAMA_API_KEY", env) + texts: Final = [_text(), _text()] + with _peer() as wire: + response: Final = await _call(mode, input=texts, api_base=wire.url, api_key=explicit) + request: Final = _only(wire) + assert [row["embedding"] for row in response.data] == [_vector(0), _vector(1)], response + assert _values_of(request, "authorization") == [f"Bearer {expected}"], request.header_lines + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": texts}, request.body + + +@pytest.mark.parametrize("mode", _MODES) +@pytest.mark.parametrize("auth_name", ("Authorization", "authorization"), ids=("canonical", "lowercase")) +@pytest.mark.parametrize("header_kwarg", ("extra_headers", "headers")) +async def test_regression_caller_headers_reach_ollama_and_override_the_key_once( + mode: Mode, auth_name: str, header_kwarg: str +) -> None: + caller: Final[dict[str, str]] = {auth_name: _CALLER_TOKEN, "X-Embed-Tenant": "tenant-a"} + snapshot: Final = dict(caller) + text: Final = _text() + with _peer() as wire: + response: Final = await _call( + mode, input=text, api_base=wire.url, api_key=_EXPLICIT_KEY, **{header_kwarg: caller} + ) + request: Final = _only(wire) + assert [row["embedding"] for row in response.data] == [_vector(0)], response + assert _values_of(request, "authorization") == [_CALLER_TOKEN], request.header_lines + assert _values_of(request, "x-embed-tenant") == ["tenant-a"], request.header_lines + assert caller == snapshot, "the caller's header map was changed by the call" + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": [text]}, request.body + + +@pytest.mark.parametrize("mode", _MODES) +async def test_regression_extra_headers_win_over_headers_and_both_reach_ollama(mode: Mode) -> None: + extra: Final[dict[str, str]] = {"X-Embed-Tenant": "from-extra-headers"} + extra_snapshot: Final = dict(extra) + with _peer() as wire: + await _call( + mode, + input=_text(), + api_base=wire.url, + headers={"X-Embed-Tenant": "from-headers", "X-Embed-Trace": "trace-1"}, + extra_headers=extra, + ) + request: Final = _only(wire) + assert _values_of(request, "x-embed-tenant") == ["from-extra-headers"], request.header_lines + assert _values_of(request, "x-embed-trace") == ["trace-1"], request.header_lines + assert extra == extra_snapshot, "the caller's extra_headers map was changed by the call" + + +@pytest.mark.parametrize("mode", _MODES) +@pytest.mark.parametrize("timeout", (0.5, httpx.Timeout(0.5)), ids=("float", "httpx-timeout")) +async def test_regression_short_timeout_cuts_a_held_ollama_read(mode: Mode, timeout: float | httpx.Timeout) -> None: + release: Final = threading.Event() + with _peer(_held(release, _DELAY_SECONDS)) as wire: + try: + elapsed: Final = await _seconds_until_timeout(mode, input=_text(), api_base=wire.url, timeout=timeout) + finally: + release.set() + _only(wire) + assert_cut_short(elapsed, injected_delay=_DELAY_SECONDS) + + +@pytest.mark.parametrize("mode", _MODES) +@pytest.mark.parametrize("timeout_kwargs", ({}, {"timeout": None}), ids=("omitted", "none")) +async def test_guard_default_timeout_does_not_cut_a_slow_ollama_reply( + mode: Mode, timeout_kwargs: Mapping[str, None] +) -> None: + release: Final = threading.Event() + with _peer(_held(release, _SHORT_DELAY_SECONDS)) as wire: + response: Final = await _call(mode, input=_text(), api_base=wire.url, **timeout_kwargs) + _only(wire) + assert [row["embedding"] for row in response.data] == [_vector(0)], response + + +@pytest.mark.parametrize("mode", _MODES) +async def test_guard_no_ollama_key_sends_no_authorization_even_with_global_keys( + monkeypatch: pytest.MonkeyPatch, mode: Mode +) -> None: + monkeypatch.setattr(litellm, "api_key", _GLOBAL_KEY) + monkeypatch.setattr(litellm, "ollama_key", _GLOBAL_KEY, raising=False) + monkeypatch.setattr(litellm, "openai_key", _GLOBAL_KEY) + texts: Final = [_text(), _text()] + with _peer() as wire: + response: Final = await _call(mode, input=texts, api_base=f"{wire.url}/") + request: Final = _only(wire) + assert [row["embedding"] for row in response.data] == [_vector(0), _vector(1)], response + assert _values_of(request, "authorization") == [], request.header_lines + assert not any(_GLOBAL_KEY in value for _, value in request.header_lines), request.header_lines + assert JSON_OBJECT.validate_json(request.body) == {"model": _MODEL, "input": texts}, request.body + + +@pytest.mark.parametrize("mode", _MODES) +async def test_guard_non_text_input_is_rejected_before_ollama_is_called(mode: Mode) -> None: + valid: Final = _text() + with _peer() as wire: + with pytest.raises(litellm.BadRequestError): + await _call(mode, input=[[1, 2, 3]], api_base=wire.url, api_key=_EXPLICIT_KEY) + await _call(mode, input=valid, api_base=wire.url) + received: Final = wire.drain() + assert [JSON_OBJECT.validate_json(request.body)["input"] for request in received] == [[valid]] + + +@pytest.mark.parametrize("mode", _MODES) +async def test_guard_ollama_error_reaches_the_caller_as_its_status(mode: Mode) -> None: + marker: Final = f"scripted-unauthorized-{uuid.uuid4().hex}" + + def respond(_: Request) -> Reply: + return Reply(status=401, body=json.dumps({"error": marker}).encode()) + + with _peer(respond) as wire: + with pytest.raises(litellm.AuthenticationError) as caught: + await _call(mode, input=_text(), api_base=wire.url) + _only(wire) + assert marker in str(caught.value), caught.value diff --git a/tests/integration/sdk/test_router_silent_model_shadow_sdk.py b/tests/integration/sdk/test_router_silent_model_shadow_sdk.py new file mode 100644 index 00000000000..821e4ddd0b7 --- /dev/null +++ b/tests/integration/sdk/test_router_silent_model_shadow_sdk.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import re +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from queue import SimpleQueue +from typing import Final + +import pytest +from integration._support.client import eventually +from integration._support.openai_wire import chat_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm import CustomStreamWrapper, Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import ModelResponseStream + +_MODEL: Final = "gpt-5.6" +_API_KEY: Final = "synthetic-shadow-key" +_ANSWER: Final = "answered by" +_NONCE: Final = re.compile(r"nonce=([0-9a-f]{32})") +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_OBJECT: Final = TypeAdapter(dict[str, object]) +_SHADOW_SECONDS: Final = 20 +_TARGETS: Final = ( + pytest.param(("shadow-a",), id="single"), + pytest.param(("shadow-a", "shadow-b"), id="pair"), +) + + +def _prompt(nonce: str) -> str: + return f"which deployment answers this? nonce={nonce}" + + +def _peer(request: Request) -> Reply: + deployment: Final = request.target.lstrip("/").partition("/")[0] + raw: Final = request.body.decode() + found: Final = _NONCE.search(raw) + assert found is not None, raw + return chat_reply( + f"chatcmpl-{deployment}-{found.group(1)}", + _MODEL, + f"{_ANSWER} {deployment}", + stream=_JSON.validate_json(request.body).get("stream") is True, + ) + + +def _deployment(name: str, wire: Wire, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": f"openai/{_MODEL}", "api_base": f"{wire.url}/{name}", "api_key": _API_KEY, **extra}, + } + + +@dataclass(frozen=True, slots=True) +class _Posted: + deployment: str + streamed: bool + leaked: bool + + +def _posted(wire: Wire) -> tuple[_Posted, ...]: + return tuple( + _Posted( + request.target.lstrip("/").partition("/")[0], + _JSON.validate_json(request.body).get("stream") is True, + "silent_model" in request.body.decode(), + ) + for request in wire.drain() + if request.method == "POST" + ) + + +class _ShadowRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.shadow_groups: SimpleQueue[str] = SimpleQueue() + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + metadata: Final = _OBJECT.validate_python(_OBJECT.validate_python(kwargs["litellm_params"])["metadata"] or {}) + if metadata.get("is_silent_experiment") is True: + self.shadow_groups.put(str(metadata.get("model_group"))) + + def drained(self) -> tuple[str, ...]: + return tuple(self.shadow_groups.get_nowait() for _ in range(self.shadow_groups.qsize())) + + +def _delta(chunk: object) -> str: + assert isinstance(chunk, ModelResponseStream), chunk + return "".join(str(choice.delta.content or "") for choice in chunk.choices) + + +@pytest.mark.parametrize("targets", _TARGETS) +def test_sync_streamed_completion_streams_and_logs_each_shadow_target( + targets: tuple[str, ...], monkeypatch: pytest.MonkeyPatch +) -> None: + recorder: Final = _ShadowRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + with wire_server(_peer) as wire: + silent_model: Final[JsonValue] = targets[0] if len(targets) == 1 else [*targets] + router: Final = Router( + model_list=[ + _deployment("mirrored", wire, silent_model=silent_model), + *(_deployment(target, wire) for target in targets), + ], + num_retries=0, + disable_cooldowns=True, + ) + response: Final = router.completion( + model="mirrored", + messages=[{"role": "user", "content": _prompt(uuid.uuid4().hex)}], + stream=True, + ) + assert isinstance(response, CustomStreamWrapper), response + text: Final = "".join(_delta(chunk) for chunk in response) + eventually(wire.received.qsize, lambda count: count >= 1 + len(targets), seconds=_SHADOW_SECONDS) + posted: Final = _posted(wire) + eventually(recorder.shadow_groups.qsize, lambda count: count >= len(targets), seconds=_SHADOW_SECONDS) + logged: Final = recorder.drained() + assert text == f"{_ANSWER} mirrored" + assert sorted(item.deployment for item in posted) == sorted(("mirrored", *targets)), posted + assert all(item.streamed for item in posted), posted + assert not any(item.leaked for item in posted), posted + assert sorted(logged) == sorted(targets), logged diff --git a/tests/integration/security/test_files_batches_access.py b/tests/integration/security/test_files_batches_access.py index 7b6fd4652cf..119055c64c7 100644 --- a/tests/integration/security/test_files_batches_access.py +++ b/tests/integration/security/test_files_batches_access.py @@ -48,6 +48,7 @@ VERTEX_ACCESS_TOKEN: Final = "ya29.scripted-delete-token" AWS_ACCESS_KEY_ID: Final = "AKIAIOSFODNN7EXAMPLE" AWS_SECRET_ACCESS_KEY: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" AWS_REGION: Final = "us-east-1" +S3_OBJECT_BODY: Final = b'{"recordId":"r1","modelOutput":{"stop_reason":"end_turn"}}\n' _ForbiddenSurface = Literal[ "upload_sdk_extra_body", "upload_form", @@ -708,6 +709,86 @@ def test_non_admin_cannot_delete_a_raw_gcs_uri_before_vertex_auth_or_storage(gat assert token_requests == (), "non-admin raw GCS delete reached the Vertex token endpoint" +def test_non_admin_still_reads_an_ordinary_provider_file_id(gateway: Gateway) -> None: + with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="provider-key-" + uuid.uuid4().hex[:8] + ) + _wait_until_every_worker_serves(gateway, model) + member: Final = _member(scenario, model) + file_id: Final = "file-" + uuid.uuid4().hex + + read: Final = gateway.request("GET", f"/v1/files/{file_id}/content", key=member.key, params={"model": model}) + assert read.status_code == 200, read.text + assert read.content == b"file-content\n", read.content + assert [ + (request.method, urlsplit(request.target).path) for request in _drained_other_than_model_list_probes(wire) + ] == [("GET", f"/v1/files/{file_id}/content")] + + +@pytest.mark.parametrize( + ("method", "route_suffix", "verb"), + [("GET", "/content", "retrieved"), ("DELETE", "", "deleted")], + ids=["content", "delete"], +) +def test_non_admin_cannot_reach_vertex_with_a_twice_encoded_raw_gcs_uri( + gateway: Gateway, method: Literal["GET", "DELETE"], route_suffix: str, verb: str +) -> None: + with ( + wire_server(_vertex_token_backend) as token_wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model=VERTEX_MODEL, + api_key=None, + vertex_project=VERTEX_PROJECT, + vertex_location=VERTEX_LOCATION, + vertex_credentials=_vertex_credentials(token_wire.url), + gcs_bucket_name=BUCKET, + ) + _wait_until_every_worker_serves(gateway, model) + caller_user: Final = scenario.user(user_role="internal_user") + caller_key: Final = scenario.key(user_id=caller_user, models=[model]) + object_uri: Final = f"gs://{BUCKET}/litellm-vertex-files/litellm-vertex-files-obj.jsonl" + twice_encoded: Final = quote(quote(object_uri, safe=""), safe="") + token_wire.drain() + + denied: Final = gateway.request( + method, f"/v1/files/{twice_encoded}{route_suffix}", key=caller_key, params={"model": model} + ) + token_requests: Final = token_wire.drain() + assert denied.status_code == 403, f"{denied.text}; token requests={token_requests}" + assert _error_message(denied, 403) == ( + f"Raw cloud storage file ids can only be {verb} by a proxy admin key. " + "Use the LiteLLM managed file id returned when the file was created." + ), denied.text + assert token_requests == (), f"non-admin twice-encoded raw GCS {method} reached the Vertex token endpoint" + + +def _assert_signed_empty_body_s3_request(request: Request, method: Literal["GET", "DELETE"], wire_url: str) -> None: + assert request.method == method, request.method + request_date: Final = request.headers["x-amz-date"] + content_hash: Final = hashlib.sha256(b"").hexdigest() + assert request.headers["x-amz-content-sha256"] == content_hash, request.headers + aws_request: Final = AWSRequest( + method=method, + url=wire_url + request.target, + headers={ + "host": request.headers["host"], + "x-amz-date": request_date, + "x-amz-content-sha256": content_hash, + }, + ) + aws_request.context["timestamp"] = request_date + auth: Final = S3SigV4Auth(Credentials(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), "s3", AWS_REGION) + signature: Final = auth.signature(auth.string_to_sign(aws_request, auth.canonical_request(aws_request)), aws_request) + expected_authorization: Final = ( + f"AWS4-HMAC-SHA256 Credential={AWS_ACCESS_KEY_ID}/{request_date[:8]}/{AWS_REGION}/s3/aws4_request, " + f"SignedHeaders={auth.signed_headers(auth.headers_to_sign(aws_request))}, Signature={signature}" + ) + assert request.headers["authorization"] == expected_authorization, request.headers["authorization"] + + def test_proxy_admin_deletes_a_raw_s3_uri_with_one_signed_delete(gateway: Gateway) -> None: with wire_server(_s3_backend()) as wire, gateway.scenario() as scenario: model: Final = _bedrock_model(scenario, wire) @@ -721,29 +802,30 @@ def test_proxy_admin_deletes_a_raw_s3_uri_with_one_signed_delete(gateway: Gatewa assert parsed.deleted is True and parsed.id == object_uri, deleted.text (request,) = _drained_other_than_model_list_probes(wire) assert request.target == f"/{BUCKET}/litellm-bedrock-files-obj.jsonl", request.target - request_date: Final = request.headers["x-amz-date"] - content_hash: Final = hashlib.sha256(b"").hexdigest() - assert request.headers["x-amz-content-sha256"] == content_hash, request.headers - aws_request: Final = AWSRequest( - method="DELETE", - url=wire.url + request.target, - headers={ - "host": request.headers["host"], - "x-amz-date": request_date, - "x-amz-content-sha256": content_hash, - }, + _assert_signed_empty_body_s3_request(request, "DELETE", wire.url) + + +def _s3_object_backend(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=S3_OBJECT_BODY, content_type="application/octet-stream") + return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404) + + +def test_proxy_admin_reads_a_raw_s3_uri_with_one_signed_get(gateway: Gateway) -> None: + with wire_server(_s3_object_backend) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_model(scenario, wire) + _wait_until_every_worker_serves(gateway, model) + object_key: Final = f"litellm-batch-outputs/{uuid.uuid4().hex}.jsonl.out" + object_uri: Final = f"s3://{BUCKET}/{object_key}" + + read: Final = gateway.request( + "GET", "/v1/files/" + quote(object_uri, safe="") + "/content", params={"model": model} ) - aws_request.context["timestamp"] = request_date - auth: Final = S3SigV4Auth(Credentials(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), "s3", AWS_REGION) - signature: Final = auth.signature( - auth.string_to_sign(aws_request, auth.canonical_request(aws_request)), aws_request - ) - expected_authorization: Final = ( - f"AWS4-HMAC-SHA256 Credential={AWS_ACCESS_KEY_ID}/{request_date[:8]}/{AWS_REGION}/s3/aws4_request, " - f"SignedHeaders={auth.signed_headers(auth.headers_to_sign(aws_request))}, Signature={signature}" - ) - assert request.headers["authorization"] == expected_authorization, request.headers["authorization"] - assert request.method == "DELETE", request.method + storage_requests: Final = _drained_other_than_model_list_probes(wire) + assert read.status_code == 200, f"{read.text}; storage requests={storage_requests}" + assert read.content == S3_OBJECT_BODY, read.content + assert [request.target for request in storage_requests] == [f"/{BUCKET}/{object_key}"], storage_requests + _assert_signed_empty_body_s3_request(storage_requests[0], "GET", wire.url) def _granting_vertex_token_backend(request: Request) -> Reply: diff --git a/tests/integration/spend/test_baggage_tag_mappings.py b/tests/integration/spend/test_baggage_tag_mappings.py index 1bf90269484..e08348a074c 100644 --- a/tests/integration/spend/test_baggage_tag_mappings.py +++ b/tests/integration/spend/test_baggage_tag_mappings.py @@ -2,8 +2,9 @@ import itertools import json import uuid from collections.abc import Generator, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from hashlib import sha256 from pathlib import Path from types import MappingProxyType @@ -15,14 +16,15 @@ import yaml from integration._support.client import ( JSON_OBJECT, Gateway, + Scenario, eventually, gateway_from_environment, list_value, object_value, ) -from integration._support.database import scratch_database +from integration._support.database import read_rows, scratch_database, write_rows from integration._support.database_relay import dropped_connection_relay, plaintext_database_url -from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter from websockets.exceptions import InvalidStatus @@ -33,6 +35,8 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" PRICED: Final[dict[str, JsonValue]] = {"input_cost_per_token": 0.01, "output_cost_per_token": 0.01} WORKER_PROBES: Final = 6 +VLLM_UPSTREAM_MODEL: Final = "baggage-model" +PINNED_CONNECTION: Final = httpx.Limits(max_connections=1) JSON_LIST: Final = TypeAdapter(list[JsonValue]) GUARDRAIL_VERDICT: Final = b'{"action":"GUARDRAIL_INTERVENED","texts":["masked baggage prompt"]}' GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api" @@ -40,6 +44,7 @@ NO_FIELDS: Final[Mapping[str, JsonValue]] = MappingProxyType({}) Route = Literal["chat", "chat_litellm_metadata", "messages", "responses"] CallerTag = Literal["none", "header", "metadata_json_with_header", "litellm_metadata"] +PassThrough = Literal["custom", "anthropic", "vllm"] RejectedClientTag = Literal["none", "header", "body", "metadata_json", "litellm_metadata"] @@ -120,11 +125,31 @@ def _responses_reply() -> Reply: ) +def _anthropic_reply() -> Reply: + marker: Final = uuid.uuid4().hex + return Reply( + body=json.dumps( + { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": marker}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + } + ).encode() + ) + + def _provider(request: Request) -> Reply: if request.target == GUARDRAIL_PATH: return Reply(body=b'{"action":"NONE"}') if request.target == "/v1/responses": return _responses_reply() + if request.target == "/v1/messages": + return _anthropic_reply() return _chat_reply() @@ -183,7 +208,15 @@ def _mapped_models(marker: str, provider: Wire, upstream_url: str) -> list[JsonV } guarded: Final = _deployment(f"guarded-{marker}", provider, tags=[f"guarded-{marker}:on"]) pre_guarded: Final = _deployment(f"guarded-pre-{marker}", provider, tags=[f"guarded-{marker}:pre"]) - return [_deployment(f"chat-{marker}", provider), guarded, pre_guarded, *routed, realtime] + vllm: Final[dict[str, JsonValue]] = { + "model_name": f"vllm-{marker}", + "litellm_params": { + "model": f"hosted_vllm/{VLLM_UPSTREAM_MODEL}", + "api_base": provider.url, + "api_key": "synthetic-vllm", + }, + } + return [_deployment(f"chat-{marker}", provider), guarded, pre_guarded, *routed, realtime, vllm] def _mapped_rules(marker: str, *, tenant_required: bool) -> list[JsonValue]: @@ -239,9 +272,12 @@ def _rig( {"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}, [_deployment_tag_mode_guardrail(marker, provider)] if tag_mode_guardrail else (), ) - with owned_proxy( - base, directory, {"GENERIC_LOGGER_ENDPOINT": callbacks.url}, config=config, workers=workers - ) as proxy: + environment: Final = { + "GENERIC_LOGGER_ENDPOINT": callbacks.url, + "ANTHROPIC_API_BASE": provider.url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + } + with owned_proxy(base, directory, environment, config=config, workers=workers) as proxy: yield Rig(proxy, provider, callbacks, marker) @@ -676,6 +712,185 @@ def test_custom_passthrough_spend_row_carries_the_mapped_tags(mapped: Rig) -> No assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key) +def _passthrough_path(rig: Rig, scenario: Scenario, entrypoint: PassThrough) -> str: + match entrypoint: + case "custom": + path: Final = f"/baggage-passthrough-{uuid.uuid4().hex}" + created: Final = rig.proxy.post( + "/config/pass_through_endpoint", {"path": path, "target": f"{rig.provider.url}/echo", "auth": True} + ) + endpoint_id: Final = str(object_value(list_value(created["endpoints"])[0])["id"]) + scenario.cleanups.callback( + lambda: rig.proxy.request( + "DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id} + ) + ) + return path + case "anthropic": + return "/anthropic/v1/messages" + case "vllm": + return "/vllm/v1/chat/completions" + + +def _passthrough_body(rig: Rig, entrypoint: PassThrough, text: str) -> dict[str, JsonValue]: + model: Final = {"custom": rig.chat_model, "anthropic": "claude-sonnet-4-5", "vllm": f"vllm-{rig.marker}"} + body: Final[dict[str, JsonValue]] = { + "model": model[entrypoint], + "max_tokens": 16, + "messages": [{"role": "user", "content": text}], + } + match entrypoint: + case "vllm": + return {**body, "metadata": {"user_id": f"user-{text}"}} + case "custom" | "anthropic": + return body + + +def _upstream_body(entrypoint: PassThrough, body: Mapping[str, JsonValue]) -> JsonValue: + match entrypoint: + case "vllm": + return {**body, "model": VLLM_UPSTREAM_MODEL} + case "custom" | "anthropic": + return dict(body) + + +def _forwarded_bodies(rig: Rig) -> tuple[JsonValue, ...]: + return tuple( + JSON_OBJECT.validate_json(request.body) for request in rig.provider.drain() if request.target != GUARDRAIL_PATH + ) + + +@pytest.mark.parametrize("entrypoint", ["custom", "anthropic", "vllm"]) +def test_passthrough_forwards_the_client_body_unchanged_while_its_spend_row_carries_the_mapped_tags( + optional: Rig, entrypoint: PassThrough +) -> None: + tenant: Final = uuid.uuid4().hex + with optional.proxy.scenario() as scenario: + path: Final = _passthrough_path(optional, scenario, entrypoint) + key_fields: Final[dict[str, JsonValue]] = ( + {"allowed_passthrough_routes": [path]} if entrypoint == "custom" else {} + ) + plain_key: Final = scenario.key(**key_fields) + key: Final = scenario.key(**key_fields) + body: Final = _passthrough_body(optional, entrypoint, tenant) + optional.provider.drain() + plain: Final = optional.proxy.request("POST", path, body, key=plain_key) + plain_forwarded: Final = _forwarded_bodies(optional) + response: Final = optional.proxy.request( + "POST", path, body, key=key, headers={"baggage": f"entrypoint=a,tenant={tenant}"} + ) + assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text) + assert plain_forwarded == (_upstream_body(entrypoint, body),) + assert _forwarded_bodies(optional) == plain_forwarded + mapped_tags: Final = {optional.exact_tag, optional.tenant_tag(tenant)} + assert _spend_tags(optional.proxy, key) == _spend_tags(optional.proxy, plain_key) | mapped_tags + + +def _general_settings_mappings(database_url: str) -> JsonValue: + rows: Final = read_rows( + 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', + ("general_settings",), + database_url=database_url, + ) + return object_value(rows[0]["param_value"]).get("request_tag_mappings") if rows else None + + +def _serving_pid(proxy: Gateway) -> int: + response: Final = proxy.request("GET", "/health/integration-readiness") + return int(response.headers["x-integration-worker-pid"]) + + +def _status_on_one_worker(owned: OwnedProxy, rig: Rig, key: str) -> tuple[int, int] | None: + with httpx.Client( + base_url=str(owned.gateway.client.base_url), timeout=15, trust_env=False, limits=PINNED_CONNECTION + ) as client: + pinned: Final = replace(rig, proxy=replace(owned.gateway, client=client)) + before: Final = _serving_pid(pinned.proxy) + status: Final = _chat(pinned, key, None).status_code + return (before, status) if _serving_pid(pinned.proxy) == before else None + + +def _worker_statuses(owned: OwnedProxy, rig: Rig, key: str) -> frozenset[tuple[int, int]]: + def answers() -> frozenset[tuple[int, int]]: + with ThreadPoolExecutor(max_workers=8) as pool: + found: Final = tuple(pool.map(lambda _: _status_on_one_worker(owned, rig, key), range(8))) + return frozenset(answer for answer in found if answer is not None) + + return eventually(answers, lambda found: frozenset(pid for pid, _ in found) >= owned.worker_pids(), seconds=30) + + +def _mappings_on_a_worker(proxy: Gateway) -> tuple[int, JsonValue]: + response: Final = proxy.request( + "GET", + "/health/integration-general-setting", + params={"name": "request_tag_mappings"}, + headers={"connection": "close"}, + ) + assert response.status_code == 200, response.text + return int(response.headers["x-integration-worker-pid"]), JSON_OBJECT.validate_json(response.content)["value"] + + +def _mappings_settled_on_every_worker(owned: OwnedProxy, mappings: JsonValue) -> None: + def answers() -> tuple[tuple[int, JsonValue], ...]: + with ThreadPoolExecutor(max_workers=8) as pool: + return tuple(pool.map(lambda _: _mappings_on_a_worker(owned.gateway), range(16))) + + eventually( + answers, + lambda found: ( + frozenset(pid for pid, _ in found) >= owned.worker_pids() and all(value == mappings for _, value in found) + ), + seconds=70, + ) + + +def test_mappings_saved_through_the_settings_endpoint_apply_on_every_worker_and_an_invalid_stored_value_is_ignored( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex[:12] + rule: Final[dict[str, JsonValue]] = {"key": "tenant", "tag_prefix": f"tenant-{marker}:", "required": True} + with ExitStack() as stack: + provider: Final = stack.enter_context(wire_server(_provider)) + config: Final = _write_config(tmp_path, [_deployment(f"chat-{marker}", provider)], {}, {}) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2)) + proxy: Final = owned.gateway + database_url: Final = proxy.database_url + assert database_url is not None + scenario: Final = stack.enter_context(proxy.scenario()) + rig: Final = Rig(proxy, provider, provider, marker) + key: Final = scenario.key() + + def save(rules: list[JsonValue]) -> httpx.Response: + return proxy.request( + "POST", + "/config/field/update", + {"field_name": "request_tag_mappings", "field_value": rules, "config_type": "general_settings"}, + ) + + duplicate: Final = save([rule, {**rule, "required": False}]) + assert duplicate.status_code == 400, duplicate.text + assert _general_settings_mappings(database_url) is None + saved: Final = save([rule]) + assert saved.status_code == 200, saved.text + _mappings_settled_on_every_worker(owned, [rule]) + assert _worker_statuses(owned, rig, key) == frozenset((pid, 400) for pid in owned.worker_pids()) + tagged_key: Final = scenario.key() + provider.drain() + admitted: Final = _chat(rig, tagged_key, f"tenant={marker}") + assert admitted.status_code == 200, admitted.text + assert len(provider.drain()) == 1 + assert _spend_tags(proxy, tagged_key) == {f"tenant-{marker}:{marker}"} + + write_rows( + 'UPDATE "LiteLLM_Config" SET param_value = jsonb_set(param_value::jsonb, %s, %s::jsonb) ' + "WHERE param_name = %s", + ("{request_tag_mappings}", json.dumps([{"key": "tenant", "required": True}]), "general_settings"), + database_url=database_url, + ) + _mappings_settled_on_every_worker(owned, None) + assert _worker_statuses(owned, rig, key) == frozenset((pid, 200) for pid in owned.worker_pids()) + + def _rejected_client_tag(client_tag: RejectedClientTag, tag: str) -> tuple[Mapping[str, JsonValue], Mapping[str, str]]: match client_tag: case "none": 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/chat_completions/test_failures.py b/tests/test_litellm_rust/chat_completions/test_failures.py index 68a81096a7c..dda5a1bad90 100644 --- a/tests/test_litellm_rust/chat_completions/test_failures.py +++ b/tests/test_litellm_rust/chat_completions/test_failures.py @@ -67,3 +67,16 @@ async def test_empty_messages_fail_before_provider_io(server: RecordingServer, a await call(server, asynchronous=asynchronous, messages=[]) assert not server.requests + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_a_native_decline_reaches_the_dispatcher_before_provider_io( + server: RecordingServer, asynchronous: bool +) -> None: + server.expected_requests = 0 + + with pytest.raises(RuntimeError, match="declined the request: native Python chat_completions streaming"): + await call(server, asynchronous=asynchronous, stream=True) + + assert not server.requests 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/test_litellm_rust/responses/data.py b/tests/test_litellm_rust/responses/data.py index 0347fd7a57b..3379b0f80b1 100644 --- a/tests/test_litellm_rust/responses/data.py +++ b/tests/test_litellm_rust/responses/data.py @@ -70,6 +70,6 @@ async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: SYNC_DEPLOYMENT_HOOK: Final = gap("native sync calls skip the async pre-call deployment hook") STREAMING: Final = gap( - "native Responses streaming is unsupported, the catalog keeps it on Python", raises=litellm.APIConnectionError + "native Responses streaming is unsupported, the catalog keeps it on Python", raises=RuntimeError ) BODY_HOOKS: Final = gap("native Responses does not run the handler body hooks") diff --git a/tests/test_litellm_rust/responses/test_failures.py b/tests/test_litellm_rust/responses/test_failures.py index 4c4bb96c6ac..bfd6629fb33 100644 --- a/tests/test_litellm_rust/responses/test_failures.py +++ b/tests/test_litellm_rust/responses/test_failures.py @@ -56,3 +56,16 @@ async def test_provider_failure_is_terminal_and_shared_with_callbacks( assert len(failure) == 1 assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value assert not any("success" in name for name in recorder.names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_a_native_decline_reaches_the_dispatcher_before_provider_io( + server: RecordingServer, asynchronous: bool +) -> None: + server.expected_requests = 0 + + with pytest.raises(RuntimeError, match="declined the request: native Python responses streaming"): + await call(server, asynchronous=asynchronous, stream=True) + + assert not server.requests 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/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index 45697723af1..3f14098f6a9 100644 --- a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -1,7 +1,7 @@ import asyncio import os import traceback -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime, timezone from typing import Final @@ -30,6 +30,7 @@ from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( delete_project, project_info, ) +from litellm.proxy.management_helpers.model_quota_policy import ModelRoute, TeamDeployment, router_model_routes from litellm.proxy.proxy_server import ( LitellmUserRoles, ) @@ -1205,7 +1206,9 @@ def test_update_quota_existing_quotaless_model_rejected(): _raise_on_missing_project_model_quota_on_update(data, existing) -def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router: mock.MagicMock | None) -> None: +def _enforced_new_project_mocks( + monkeypatch: pytest.MonkeyPatch, team_models: list[str], llm_router: object, quota_setting: object = True +) -> None: from litellm.proxy._types import LiteLLM_TeamTable from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe @@ -1213,7 +1216,7 @@ def _enforced_new_project_mocks(monkeypatch, team_models: list[str], llm_router: monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock.MagicMock()) monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": True}) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enforce_project_model_quota": quota_setting}) monkeypatch.setattr(pe, "_validate_team_exists", mock.AsyncMock(return_value=team)) monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=True)) @@ -1604,6 +1607,559 @@ async def test_new_project_flag_on_access_group_model_returns_400(monkeypatch): assert "expand to multiple models at request time" in str(exc_info.value) +_CUSTOMER_QUOTA_POLICY: Final = { + "providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, + "models": {"veo-*": ["rpm"], "jev-latest": []}, +} + +_CUSTOMER_MODEL_PROVIDERS: Final = { + "gemini-3.8-flash": frozenset({"vertex_ai"}), + "veo-3.1-lite-generate-001": frozenset({"vertex_ai"}), + "runway-gen4-turbo": frozenset({"runwayml"}), +} + + +def _customer_route_of(model: str) -> ModelRoute: + return ModelRoute(providers=_CUSTOMER_MODEL_PROVIDERS.get(model, frozenset())) + + +def _quota_error( + models: list[str], + rpm: dict[str, int] | None = None, + tpm: dict[str, int] | None = None, + setting: object = _CUSTOMER_QUOTA_POLICY, + route_of: Callable[[str], ModelRoute] = _customer_route_of, +) -> str | None: + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _raise_on_missing_project_model_quota, + ) + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + policy: Final = parse_model_quota_setting(setting) + assert policy is not None + data: Final = NewProjectRequest(team_id="test-team", models=models, model_rpm_limit=rpm, model_tpm_limit=tpm) + try: + _raise_on_missing_project_model_quota(data, frozenset(), policy, route_of) + except HTTPException as e: + rejection: Final = e + else: + return None + assert rejection.status_code == 400 + return str(rejection.detail) + + +def test_quota_policy_exempts_provider_and_passthrough_name_but_keeps_vertex_strict(): + assert _quota_error(["runway-gen4-turbo", "jev-latest"]) is None + error: Final = _quota_error(["gemini-3.8-flash"], rpm={"gemini-3.8-flash": 10}) + assert error is not None + assert "gemini-3.8-flash (model_tpm_limit)" in error + + +def test_quota_policy_model_pattern_overrides_its_provider_per_limit(): + assert _quota_error(["veo-3.1-lite-generate-001"], rpm={"veo-3.1-lite-generate-001": 5}) is None + error: Final = _quota_error(["veo-3.1-lite-generate-001"], tpm={"veo-3.1-lite-generate-001": 5000}) + assert error is not None + assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error + assert "model_tpm_limit" not in error + + +@pytest.mark.parametrize( + ("model", "rpm", "tpm", "expected_missing"), + [ + ("runway-gen4-turbo", {"runway-gen4-turbo": 0}, None, "runway-gen4-turbo (model_rpm_limit)"), + ("jev-latest", None, {"jev-latest": -5}, "jev-latest (model_tpm_limit)"), + ( + "veo-3.1-lite-generate-001", + {"veo-3.1-lite-generate-001": 5}, + {"veo-3.1-lite-generate-001": 0}, + "veo-3.1-lite-generate-001 (model_tpm_limit)", + ), + ], +) +def test_quota_policy_rejects_a_non_positive_limit_set_on_an_exempt_model( + model: str, rpm: dict[str, int] | None, tpm: dict[str, int] | None, expected_missing: str +): + error: Final = _quota_error([model], rpm=rpm, tpm=tpm) + assert error is not None + assert expected_missing in error + + +def test_quota_policy_error_names_each_model_with_the_limits_it_misses(): + error: Final = _quota_error(["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo"]) + assert error is not None + assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in error + assert "veo-3.1-lite-generate-001 (model_rpm_limit)" in error + assert "runway-gen4-turbo" not in error + + +@pytest.mark.parametrize( + "setting, model, expected_missing", + [ + ({"default": ["tpm"]}, "unrouted-name", "unrouted-name (model_tpm_limit)"), + ({"default": [], "providers": {"runwayml": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"), + ( + {"providers": {"runwayml": []}, "models": {"runway-gen4-turbo": ["tpm"]}}, + "runway-gen4-turbo", + "(model_tpm_limit)", + ), + ({"models": {"runway-*": [], "runway-gen4-turbo": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"), + ({"models": {"veo-*": [], "veo-3.1-lite-*": ["rpm"]}}, "veo-3.1-lite-generate-001", "(model_rpm_limit)"), + ({"models": {"gpt-5-mini*": [], "gpt-*-*-*": ["rpm"]}}, "gpt-5-mini-fast", "(model_rpm_limit)"), + ], +) +def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default( + setting: object, model: str, expected_missing: str +): + error: Final = _quota_error([model], setting=setting) + assert error is not None + assert expected_missing in error + + +def test_quota_policy_model_entry_matches_names_literally_apart_from_star(): + setting: Final = {"models": {"claude-opus-5-5[1m]": [], "gpt.5": []}} + assert _quota_error(["claude-opus-5-5[1m]"], setting=setting) is None + assert _quota_error(["claude-opus-5-51"], setting=setting) is not None + assert _quota_error(["gptx5"], setting=setting) is not None + + +def test_quota_policy_model_on_several_providers_needs_every_required_limit(): + def mixed_providers(_model: str) -> ModelRoute: + return ModelRoute(providers=frozenset({"runwayml", "vertex_ai"})) + + error: Final = _quota_error(["video-pool"], rpm={"video-pool": 10}, route_of=mixed_providers) + assert error is not None + assert "video-pool (model_tpm_limit)" in error + + +def test_quota_policy_model_pattern_matches_from_the_start_like_the_router(): + setting: Final = {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"*-lite": []}} + assert _quota_error(["veo-3.1-lite-generate-001"], setting=setting) is None + assert _quota_error(["gemini-3.8-flash"], setting=setting) is not None + + +def _team_vertex_route(called_model: str) -> Callable[[str], ModelRoute]: + def route_of(_model: str) -> ModelRoute: + return ModelRoute(team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model=called_model)})) + + return route_of + + +def test_quota_policy_matches_a_team_model_by_the_model_it_calls(): + borrowed_name: Final = _quota_error(["jev-latest"], route_of=_team_vertex_route("gemini-3.8-flash")) + veo_call: Final = _quota_error( + ["video-team"], rpm={"video-team": 5}, route_of=_team_vertex_route("veo-3.1-generate-001") + ) + assert borrowed_name is not None and "jev-latest (model_rpm_limit and model_tpm_limit)" in borrowed_name + assert veo_call is None + + +def test_quota_policy_team_model_beside_a_shared_one_needs_both_routes_limits(): + def team_vertex_beside_shared_runway(_model: str) -> ModelRoute: + return ModelRoute( + providers=frozenset({"runwayml"}), + team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model="gemini-3.8-flash")}), + ) + + error: Final = _quota_error(["media-latest"], rpm={"media-latest": 5}, route_of=team_vertex_beside_shared_runway) + assert error is not None and "media-latest (model_tpm_limit)" in error + + +@pytest.mark.parametrize("setting", [True, "true", "false", 1, ["rpm"], {"enabled": True}, {"provider": {"runwayml": []}}]) +def test_quota_setting_truthy_non_policy_value_still_requires_both_limits(setting: object): + error: Final = _quota_error(["runway-gen4-turbo"], setting=setting) + assert error is not None + assert "runway-gen4-turbo (model_rpm_limit and model_tpm_limit)" in error + + +@pytest.mark.parametrize("setting", [False, None, {}, "", 0, []]) +def test_quota_setting_off_enforces_nothing(setting: object): + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + assert parse_model_quota_setting(setting) is None + + +def test_update_quota_policy_resolves_providers_for_the_resulting_models(): + import types + + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _raise_on_missing_project_model_quota_on_update, + ) + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + policy: Final = parse_model_quota_setting(_CUSTOMER_QUOTA_POLICY) + assert policy is not None + existing: Final = types.SimpleNamespace(models=["runway-gen4-turbo"], metadata={}) + assert ( + _raise_on_missing_project_model_quota_on_update( + UpdateProjectRequest(project_id="p", description="unrelated change"), + existing, + frozenset(), + policy, + _customer_route_of, + ) + is None + ) + with pytest.raises(HTTPException) as exc_info: + _raise_on_missing_project_model_quota_on_update( + UpdateProjectRequest(project_id="p", models=["runway-gen4-turbo", "gemini-3.8-flash"]), + existing, + frozenset(), + policy, + _customer_route_of, + ) + assert "gemini-3.8-flash (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.detail) + + +def _customer_router() -> "litellm.Router": + return litellm.Router( + model_list=[ + { + "model_name": "gemini-3.8-flash", + "litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}, + }, + { + "model_name": "veo-3.1-lite-generate-001", + "litellm_params": {"model": "vertex_ai/veo-3.1-lite-generate-001", "vertex_project": "qa-project"}, + }, + { + "model_name": "runway-gen4-turbo", + "litellm_params": {"model": "runwayml/gen4_turbo", "api_key": "qa-runway-key"}, + }, + ] + ) + + +def test_quota_providers_come_from_the_one_wildcard_route_requests_use(): + router: Final = litellm.Router( + model_list=[ + {"model_name": "gen4*", "litellm_params": {"model": "runwayml/gen4*", "api_key": "qa-runway-key"}}, + { + "model_name": "media-team-gen", + "litellm_params": {"model": "vertex_ai/*", "vertex_project": "qa-project"}, + "model_info": {"team_id": "media-team", "team_public_model_name": "gen*"}, + }, + ] + ) + + assert router_model_routes(router, "media-team")("gen4_turbo") == ModelRoute( + team_deployments=frozenset({TeamDeployment(provider="vertex_ai", called_model="4_turbo")}) + ) + assert router_model_routes(router, None)("gen4_turbo") == ModelRoute(providers=frozenset({"runwayml"})) + + +def test_quota_providers_read_the_provider_behind_a_known_api_base(): + groq_api_base: Final = "https://api.groq.com/openai/v1" + router: Final = litellm.Router( + model_list=[ + { + "model_name": "qa-llama", + "litellm_params": {"model": "qa-llama", "api_base": groq_api_base, "api_key": "qa-groq-key"}, + }, + { + "model_name": "qa-runway", + "litellm_params": {"model": "gen4_turbo", "custom_llm_provider": "runwayml", "api_base": groq_api_base}, + }, + ] + ) + route_of: Final = router_model_routes(router, None) + + assert route_of("qa-llama").providers == frozenset({"groq"}) + assert route_of("qa-runway").providers == frozenset({"runwayml"}) + + +def test_quota_policy_rejects_an_unknown_provider_id(): + from pydantic import ValidationError + + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + with pytest.raises(ValidationError, match=r"unknown provider ids \['runway'\]"): + parse_model_quota_setting({"providers": {"vertex_ai": ["rpm"], "runway": []}}) + + +def test_quota_policy_accepts_a_provider_registered_at_runtime(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "qa-custom-llm"]) + policy: Final = parse_model_quota_setting({"providers": {"qa-custom-llm": []}}) + + assert policy is not None + assert policy.providers["qa-custom-llm"] == frozenset() + + +_TEAM_WRITE_POLICY: Final = {"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, "models": {"jev-latest": []}} +_VERTEX_PARAMS: Final = {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"} +_RUNWAY_PARAMS: Final = {"model": "gen4_turbo", "custom_llm_provider": "runwayml"} + + +def _team_deployment(public_name: str, model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]: + return { + "model_name": f"model_name_media-team_{model_id}", + "litellm_params": dict(litellm_params), + "model_info": {"id": model_id, "team_id": "media-team", "team_public_model_name": public_name}, + } + + +@pytest.mark.parametrize( + ("setting", "public_name", "proxy_params", "team_routes", "replaced_id", "written_params", "added"), + [ + pytest.param( + _TEAM_WRITE_POLICY, + "jev-latest", + {"model": "openai/jev-latest"}, + (), + None, + _VERTEX_PARAMS, + {"rpm", "tpm"}, + id="team-add-on-vertex-behind-an-exempt-proxy-name", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "jev-latest", + {"model": "openai/jev-latest"}, + (), + None, + _RUNWAY_PARAMS, + set(), + id="team-add-on-an-exempt-provider", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "jev-latest", + {"model": "openai/jev-latest"}, + (), + None, + {"model": "aihubmix/jev-latest"}, + set(), + id="team-add-through-a-json-provider-matches-the-model-it-calls", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "media-latest", + _VERTEX_PARAMS, + (("t1", _RUNWAY_PARAMS),), + "t1", + None, + {"rpm", "tpm"}, + id="team-delete-falls-back-to-the-proxy-vertex-route", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "media-latest", + _VERTEX_PARAMS, + (("t1", _RUNWAY_PARAMS), ("t2", _RUNWAY_PARAMS)), + "t1", + None, + set(), + id="team-delete-keeps-a-sibling-team-route", + ), + pytest.param( + _TEAM_WRITE_POLICY, + "media-latest", + _VERTEX_PARAMS, + (("t1", _RUNWAY_PARAMS),), + "t1", + _VERTEX_PARAMS, + {"rpm", "tpm"}, + id="team-update-moves-the-route-to-vertex", + ), + pytest.param( + True, + "jev-latest", + {"model": "openai/jev-latest"}, + (), + None, + _VERTEX_PARAMS, + set(), + id="true-already-requires-both-limits", + ), + ], +) +def test_team_model_write_adds_only_the_limits_its_new_route_requires( + setting: object, + public_name: str, + proxy_params: Mapping[str, object], + team_routes: tuple[tuple[str, Mapping[str, object]], ...], + replaced_id: str | None, + written_params: Mapping[str, object] | None, + added: set[str], +): + from litellm.proxy.management_helpers.model_quota_policy import ( + limits_added_by_team_write, + parse_model_quota_setting, + ) + + policy: Final = parse_model_quota_setting(setting) + router: Final = litellm.Router( + model_list=[ + {"model_name": public_name, "litellm_params": dict(proxy_params)}, + *(_team_deployment(public_name, model_id, params) for model_id, params in team_routes), + ] + ) + + assert policy is not None + assert limits_added_by_team_write( + policy, router, "media-team", public_name, replaced_id, written_params, public_name + ) == frozenset(added) + + +def test_quota_route_of_an_alias_reads_the_team_deployments_its_target_routes_to(): + router: Final = litellm.Router( + model_list=[ + {"model_name": "media-latest", "litellm_params": dict(_RUNWAY_PARAMS)}, + _team_deployment("media-latest", "t1", _VERTEX_PARAMS), + ], + model_group_alias={"media": "media-latest"}, + ) + route_of: Final = router_model_routes(router, "media-team") + + assert route_of("media") == route_of("media-latest") + assert TeamDeployment(provider="vertex_ai", called_model="gemini-3.8-flash") in route_of("media").team_deployments + + +def _team_pattern_router() -> "litellm.Router": + return litellm.Router( + model_list=[ + {"model_name": "media-latest", "litellm_params": dict(_VERTEX_PARAMS)}, + {"model_name": "video-named", "litellm_params": dict(_VERTEX_PARAMS)}, + _team_deployment("video-*", "t1", {"model": "runwayml/*"}), + _team_deployment("video-veo-*", "t2", {"model": "vertex_ai/veo-*", "vertex_project": "qa-project"}), + _team_deployment("clip-*x", "t3", {"model": "runwayml/*"}), + _team_deployment("clip-x*", "t4", {"model": "vertex_ai/*", "vertex_project": "qa-project"}), + ], + model_group_alias={"media": "media-latest"}, + ) + + +@pytest.mark.parametrize( + ("public_name", "model", "served"), + [ + pytest.param("media-latest", "media-latest", True, id="the-name-itself"), + pytest.param("media-latest", "media", True, id="an-alias-of-the-name"), + pytest.param("media-latest", "video-named", False, id="another-name"), + pytest.param("video-*", "video-gen-1", True, id="a-name-the-pattern-matches"), + pytest.param("video-*", "video-veo-3", False, id="a-more-specific-team-pattern-takes-it"), + pytest.param("video-veo-*", "video-veo-3", True, id="the-most-specific-team-pattern"), + pytest.param("video-*", "video-named", False, id="a-named-deployment-shadows-the-pattern"), + pytest.param("video-*", "image-gen-1", False, id="a-name-the-pattern-misses"), + pytest.param("audio-*", "audio-gen-1", True, id="a-pattern-the-write-adds"), + pytest.param("clip-*x", "clip-xx", True, id="the-first-registered-of-two-tied-team-patterns"), + pytest.param("clip-x*", "clip-xx", False, id="a-tied-team-pattern-registered-first-takes-it"), + pytest.param("cli*-xx", "clip-xx", False, id="a-new-pattern-tied-with-a-registered-one"), + ], +) +def test_team_write_serves_the_project_names_a_team_request_routes_through_it( + public_name: str, model: str, served: bool +): + from litellm.proxy.management_helpers.model_quota_policy import team_write_serves + + assert team_write_serves(_team_pattern_router(), "media-team", public_name, model) is served + + +@pytest.mark.parametrize( + ("model", "added"), + [ + pytest.param("video-gen-1", {"rpm", "tpm"}, id="vertex-requires-both"), + pytest.param("video-veo-3", {"rpm"}, id="the-called-veo-model-entry-wins"), + ], +) +def test_team_pattern_write_reads_each_served_name_by_the_model_it_calls(model: str, added: set[str]): + from litellm.proxy.management_helpers.model_quota_policy import ( + limits_added_by_team_write, + parse_model_quota_setting, + ) + + policy: Final = parse_model_quota_setting( + {"default": [], "providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"veo-*": ["rpm"]}} + ) + + assert policy is not None + assert limits_added_by_team_write( + policy, + litellm.Router(model_list=[]), + "media-team", + "video-*", + None, + {"model": "vertex_ai/*", "vertex_project": "qa-project"}, + model, + ) == frozenset(added) + + +@pytest.mark.asyncio +async def test_new_project_quota_policy_reads_providers_from_router_deployments(monkeypatch: pytest.MonkeyPatch): + team_models: Final = ["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo", "jev-latest"] + _enforced_new_project_mocks( + monkeypatch, team_models=team_models, llm_router=_customer_router(), quota_setting=_CUSTOMER_QUOTA_POLICY + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=team_models)) + + assert exc_info.value.code == "400" + assert ( + "gemini-3.8-flash (model_rpm_limit and model_tpm_limit), veo-3.1-lite-generate-001 (model_rpm_limit)." + in str(exc_info.value.message) + ) + assert "runway-gen4-turbo" not in str(exc_info.value.message) + assert "jev-latest" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_new_project_rejects_malformed_quota_setting_instead_of_skipping_it(monkeypatch: pytest.MonkeyPatch): + _enforced_new_project_mocks( + monkeypatch, + team_models=["runway-gen4-turbo"], + llm_router=_customer_router(), + quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}}, + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=["runway-gen4-turbo"])) + + assert exc_info.value.code == "500" + assert "general_settings.enforce_project_model_quota is invalid" in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_new_project_checks_permission_before_reading_the_quota_setting(monkeypatch: pytest.MonkeyPatch): + from litellm_enterprise.proxy.management_endpoints import project_endpoints as pe + + _enforced_new_project_mocks( + monkeypatch, + team_models=["runway-gen4-turbo"], + llm_router=_customer_router(), + quota_setting={"providers": {"runwayml": []}, "model": {"jev-latest": []}}, + ) + monkeypatch.setattr(pe, "_check_user_permission_for_project", mock.AsyncMock(return_value=False)) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=["runway-gen4-turbo"])) + + assert exc_info.value.code == "403" + assert "enforce_project_model_quota" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_new_project_quota_policy_ignores_model_entries_for_a_team_model_name(monkeypatch: pytest.MonkeyPatch): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "model_name_test-team_qa", + "litellm_params": {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"}, + "model_info": {"team_id": "test-team", "team_public_model_name": "jev-latest"}, + }, + ] + ) + _enforced_new_project_mocks( + monkeypatch, team_models=["jev-latest"], llm_router=router, quota_setting=_CUSTOMER_QUOTA_POLICY + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_new_project(NewProjectRequest(team_id="test-team", models=["jev-latest"])) + + assert exc_info.value.code == "400" + assert "jev-latest (model_rpm_limit and model_tpm_limit)" in str(exc_info.value.message) + + _ACCESS_NOW: Final = datetime.now(timezone.utc) 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/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/images/test_main.py b/tests/unit/images/test_main.py index 38af99ff505..88d5ae1fc2e 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -3,14 +3,18 @@ import base64 import json from datetime import datetime from pathlib import Path -from typing import Final +from collections.abc import Generator +from queue import SimpleQueue +from typing import Final, Literal, TypedDict from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import AsyncOpenAI +from openai import AsyncOpenAI, OpenAI from openai.types.image import Image +from pydantic import JsonValue, TypeAdapter +from typing_extensions import ReadOnly import litellm from litellm.integrations.custom_logger import CustomLogger @@ -19,6 +23,212 @@ from litellm.types.utils import CallTypes, StandardLoggingPayload import os +class ImageRetryParameters(TypedDict, total=False): + max_retries: ReadOnly[int | None] + + +@pytest.fixture +def image_retry_delays() -> Generator[None, None, None]: + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + yield + + +def _image_retry_transport( + requests: SimpleQueue[httpx.Request], outcome: Literal["503", "timeout", "success"] +) -> httpx.MockTransport: + def respond(request: httpx.Request) -> httpx.Response: + requests.put(request) + assert request.method == "POST" + assert request.url.path == "/v1/images/generations" + body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content) + assert body == {"model": "gpt-image-1-mini", "prompt": "a red circle"} + if outcome == "timeout": + raise httpx.ReadTimeout("scripted image timeout", request=request) + if outcome == "503": + return httpx.Response(503, json={"error": {"message": "scripted unavailable", "type": "server_error"}}) + return httpx.Response(200, json={"created": 1700000000, "data": [{"b64_json": "aW1n"}]}) + + return httpx.MockTransport(respond) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("image_retry_delays") +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("outcome", ("503", "timeout")) +@pytest.mark.parametrize( + ("retries", "attempts"), + ((0, 1), (1, 2), (None, 3), ("unset", 3), ("global", 3)), + ids=("zero", "one", "none", "unset", "global"), +) +async def test_image_generation_honors_sdk_retry_budget( + monkeypatch: pytest.MonkeyPatch, + is_async: bool, + outcome: Literal["503", "timeout"], + retries: int | None | Literal["unset", "global"], + attempts: int, +) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + transport: Final = _image_retry_transport(requests, outcome) + if retries == "global": + monkeypatch.setattr(litellm, "max_retries", 0, raising=False) + retry_params: Final[ImageRetryParameters] = {} if isinstance(retries, str) else {"max_retries": retries} + error: Final = litellm.Timeout if outcome == "timeout" else litellm.ServiceUnavailableError + if is_async: + async with AsyncOpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + max_retries=5, + http_client=httpx.AsyncClient(transport=transport), + ) as client: + with pytest.raises(error): + await litellm.aimage_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + **retry_params, + ) + else: + with OpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + max_retries=5, + http_client=httpx.Client(transport=transport), + ) as client: + with pytest.raises(error): + litellm.image_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + **retry_params, + ) + assert requests.qsize() == attempts + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("image_retry_delays") +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("provider", ("litellm_proxy", "hosted_vllm")) +async def test_compatible_image_generation_honors_zero_retries(provider: str, is_async: bool) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + transport: Final = _image_retry_transport(requests, "503") + if is_async: + async with AsyncOpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.AsyncClient(transport=transport), + ) as client: + with pytest.raises(litellm.ServiceUnavailableError): + await litellm.aimage_generation( + model=f"{provider}/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries=0, + ) + else: + with OpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.Client(transport=transport), + ) as client: + with pytest.raises(litellm.ServiceUnavailableError): + litellm.image_generation( + model=f"{provider}/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries=0, + ) + assert requests.qsize() == 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("image_retry_delays") +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +async def test_image_generation_rejects_non_integer_retries_before_sending(is_async: bool) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + transport: Final = _image_retry_transport(requests, "success") + if is_async: + async with AsyncOpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.AsyncClient(transport=transport), + ) as client: + with pytest.raises(litellm.BadRequestError) as raised: + await litellm.aimage_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries="invalid", + ) + else: + with OpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.Client(transport=transport), + ) as client: + with pytest.raises(litellm.BadRequestError) as raised: + litellm.image_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries="invalid", + ) + assert raised.value.status_code == 422 + assert "max retries must be an int" in raised.value.message + assert requests.qsize() == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +async def test_image_generation_returns_success_without_forwarding_retry_control(is_async: bool) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + transport: Final = _image_retry_transport(requests, "success") + if is_async: + async with AsyncOpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.AsyncClient(transport=transport), + ) as client: + async_response: Final = await litellm.aimage_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries=0, + ) + async_data: Final = TypeAdapter(tuple[Image, ...]).validate_python(async_response.data) + assert async_data[0].b64_json == "aW1n" + else: + with OpenAI( + api_key="synthetic-image-key", + base_url="https://image.test/v1", + http_client=httpx.Client(transport=transport), + ) as client: + sync_response: Final = litellm.image_generation( + model="openai/gpt-image-1-mini", + prompt="a red circle", + api_key="synthetic-image-key", + api_base="https://image.test/v1", + client=client, + max_retries=0, + ) + sync_data: Final = TypeAdapter(tuple[Image, ...]).validate_python(sync_response.data) + assert sync_data[0].b64_json == "aW1n" + assert requests.qsize() == 1 + + def test_image_generation_keeps_an_internal_prefixed_kwarg_out_of_the_provider_request( respx_mock: respx.MockRouter, ) -> None: diff --git a/tests/unit/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py index 8af13fd6e7e..aad06d33edb 100644 --- a/tests/unit/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -1076,6 +1076,28 @@ def test_build_resource_includes_deployment_environment(): assert resource.attributes["deployment.environment"] == "prod" +@pytest.mark.parametrize("sampler", ["always_on", "always_off", "parentbased_always_on"]) +def test_request_exclusion_preserves_configured_sampling(monkeypatch: pytest.MonkeyPatch, sampler: str) -> None: + from opentelemetry.trace import NonRecordingSpan, SpanContext, set_span_in_context + + monkeypatch.setenv("OTEL_TRACES_SAMPLER", sampler) + provider: Final = providers.build_tracer_provider(OpenTelemetryV2Config(exporter="in_memory")) + tracer: Final = provider.get_tracer("request-exclusions") + sampled_parent: Final = set_span_in_context( + NonRecordingSpan(SpanContext(trace_id=1, span_id=2, is_remote=True, trace_flags=TraceFlags(1))) + ) + try: + with ctx_mod.request_trace_scope(True): + excluded: Final = tracer.start_span("excluded", context=sampled_parent) + assert not excluded.is_recording() + excluded.end() + included: Final = tracer.start_span("included", context=sampled_parent) + assert included.is_recording() == (sampler != "always_off") + included.end() + finally: + provider.shutdown() + + def test_build_tracer_provider_processor_selection(): cfg = OpenTelemetryV2Config(exporter="in_memory") simple = providers.build_tracer_provider(cfg, exporter=InMemorySpanExporter()) diff --git a/tests/unit/integrations/otel/test_otel_v2_metrics.py b/tests/unit/integrations/otel/test_otel_v2_metrics.py index 016dbcd824b..74921512825 100644 --- a/tests/unit/integrations/otel/test_otel_v2_metrics.py +++ b/tests/unit/integrations/otel/test_otel_v2_metrics.py @@ -24,6 +24,7 @@ dashboard already querying that histogram. import asyncio import json from datetime import datetime, timedelta +from typing import Final import pytest @@ -40,6 +41,10 @@ from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402 from litellm.integrations.otel.model.config import ( # noqa: E402 OpenTelemetryV2Config, ) +from litellm.integrations.otel.model.semconv import ( # noqa: E402 + PASSTHROUGH_OPERATION_KEY, + GenAIOperation, +) from litellm.integrations.otel.plumbing.metrics import ( # noqa: E402 GenAIMetricRecorder, create_genai_metrics, @@ -500,6 +505,17 @@ def test_agent_message_is_not_labelled_as_chat(call_type): assert dp.attributes[OPERATION_KEY] == "invoke_agent" +def test_passthrough_operation_labels_the_metrics() -> None: + reader: Final = InMemoryMetricReader() + kwargs, response_obj, start, end = _build_call(call_type="pass_through_endpoint") + labelled: Final = {**kwargs, PASSTHROUGH_OPERATION_KEY: GenAIOperation.RETRIEVAL} + + asyncio.run(_logger(reader, enable_metrics=True).async_log_success_event(labelled, response_obj, start, end)) + + for dp in _metrics_by_name(reader)[OPERATION_DURATION]: + assert dp.attributes[OPERATION_KEY] == "retrieval" + + def test_provider_name_is_filterable(): """gen_ai.provider.name is a member of the metric-attribute allowlist, so an operator can include or exclude it; an unlisted name raises instead.""" diff --git a/tests/unit/integrations/otel/test_otel_v2_mount.py b/tests/unit/integrations/otel/test_otel_v2_mount.py index 61f8b9ae0d9..32e846cc1b9 100644 --- a/tests/unit/integrations/otel/test_otel_v2_mount.py +++ b/tests/unit/integrations/otel/test_otel_v2_mount.py @@ -5,8 +5,11 @@ surface and the server-span + shared-provider behavior it produces. """ +import asyncio import contextvars +from contextlib import nullcontext from datetime import datetime, timezone +from typing import Final import pytest @@ -16,6 +19,7 @@ pytest.importorskip("opentelemetry.instrumentation.fastapi") fastapi = pytest.importorskip("fastapi") from fastapi.middleware.cors import CORSMiddleware # noqa: E402 +from fastapi import FastAPI, HTTPException, Request, Response # noqa: E402 # optional dependencies are skipped first from fastapi.responses import StreamingResponse # noqa: E402 from fastapi.testclient import TestClient # noqa: E402 from starlette.routing import Route # noqa: E402 @@ -36,6 +40,7 @@ from opentelemetry.sdk.trace.sampling import ALWAYS_OFF # noqa: E402 from opentelemetry.trace import SpanKind # noqa: E402 from litellm.constants import LITELLM_UI_ALLOW_HEADERS # noqa: E402 +from litellm.types.services import ServiceLoggerPayload, ServiceTypes # noqa: E402 # optional dependencies are skipped first from litellm.integrations.otel.model.config import ( # noqa: E402 OpenTelemetryV2Config, @@ -548,3 +553,139 @@ def test_the_unhandled_error_helper_names_the_span_it_is_given_in_any_context(mo headers = contextvars.Context().run(trace_response_headers, span) assert headers["traceresponse"] == f"00-{context.trace_id:032x}-{context.span_id:016x}-{context.trace_flags:02x}" + + +@pytest.mark.parametrize( + ("excluded_urls", "path", "root_path", "status", "excluded"), + [ + ("/health,/model/info", "/model/info", "", 200, True), + ("/model/info", "/model/info", "", 401, True), + ("/model/info", "/v1/models", "", 200, False), + (None, "/health/readiness", "", 200, True), + (None, "/model/info", "", 200, False), + ("", "/health/readiness", "", 200, False), + (r"^http://testserver/proxy/model/info$", "/model/info", "/proxy", 200, True), + (r"^http://other/model/info$", "/model/info", "", 200, False), + ], +) +@pytest.mark.parametrize("active_parent", [False, True]) +def test_url_exclusions_cover_auth_and_service_spans( + monkeypatch: pytest.MonkeyPatch, excluded_urls: str | None, path: str, + root_path: str, status: int, excluded: bool, active_parent: bool, +) -> None: + monkeypatch.setenv("LITELLM_OTEL_V2", "1") + if excluded_urls is None: + monkeypatch.delenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", raising=False) + else: + monkeypatch.setenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", excluded_urls) + app: Final = FastAPI(root_path=root_path) + logger: Final = OpenTelemetryV2(config=OpenTelemetryV2Config(exporter="in_memory")) + exporter: Final = InMemorySpanExporter() + logger.tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + + @app.get(path) + async def endpoint() -> Response: + with logger.start_phase_span(f"auth {path}"), logger.tracer.start_as_current_span("direct SDK span"): + await logger.async_service_success_hook( + payload=ServiceLoggerPayload( + is_error=False, service=ServiceTypes.DB, duration=1, + call_type="get_data", event_metadata=None, + ), + start_time=1, end_time=2, + ) + if status == 401: + raise HTTPException(status_code=status) + return Response(status_code=status) + + instrument_fastapi_app(app, tracer_provider=logger.tracer_provider) + with ( + logger.tracer.start_as_current_span("caller") if active_parent else nullcontext(), + TestClient(app) as client, + ): + response: Final = client.get(f"{root_path}{path}?poll=1") + + names: Final = tuple(span.name for span in exporter.get_finished_spans()) + expected_request: Final = () if excluded else ( + "postgres get_data", "direct SDK span", f"auth {path}", f"GET {path}", + ) + assert response.status_code == status + assert names == expected_request + (("caller",) if active_parent else ()) + + +async def test_exclusion_survives_deferred_work_without_suppressing_other_requests( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_OTEL_V2", "1") + monkeypatch.setenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", "/model/info") + app: Final = FastAPI() + logger: Final = OpenTelemetryV2(config=OpenTelemetryV2Config(exporter="in_memory")) + exporter: Final = InMemorySpanExporter() + logger.tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + finish: Final = asyncio.Event() + deferred: Final[asyncio.Queue[asyncio.Task[None]]] = asyncio.Queue() + + async def after_response(path: str) -> None: + await finish.wait() + with logger.start_phase_span(f"deferred {path}"): + pass + + @app.get("/{path:path}") + async def endpoint(path: str) -> Response: + with logger.start_phase_span(f"auth {path}"): + if path == "model/info": + entered.set() + await release.wait() + await deferred.put(asyncio.create_task(after_response(path))) + return Response(status_code=200) + + instrument_fastapi_app(app, tracer_provider=logger.tracer_provider) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + excluded: Final = asyncio.create_task(client.get("/model/info")) + await entered.wait() + included: Final = await client.get("/v1/models") + release.set() + assert (await excluded).status_code == included.status_code == 200 + finish.set() + await asyncio.gather(await deferred.get(), await deferred.get()) + with logger.start_phase_span("independent SDK work"): + pass + + assert sorted(span.name for span in exporter.get_finished_spans()) == [ + "GET /{path:path}", "auth v1/models", "deferred v1/models", "independent SDK work", + ] + + +@pytest.mark.parametrize("path", ["/excluded", "/included"]) +def test_url_exclusions_cover_unhandled_error_handlers_after_repeated_mount( + monkeypatch: pytest.MonkeyPatch, path: str, +) -> None: + monkeypatch.setenv("LITELLM_OTEL_V2", "1") + monkeypatch.setenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", "/excluded") + app: Final = FastAPI() + logger: Final = OpenTelemetryV2(config=OpenTelemetryV2Config(exporter="in_memory")) + exporter: Final = InMemorySpanExporter() + logger.tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + + async def error_handler(_request: Request, _exc: Exception) -> Response: + with logger.start_phase_span("error handler"), logger.tracer.start_as_current_span("error SDK span"): + return Response("handled", status_code=500) + + app.add_exception_handler(Exception, error_handler) + + @app.get(path) + async def endpoint() -> None: + with logger.start_phase_span("auth"): + raise ValueError("unhandled") + + instrument_fastapi_app(app, tracer_provider=logger.tracer_provider) + instrument_fastapi_app(app, tracer_provider=logger.tracer_provider) + with TestClient(app, raise_server_exceptions=False) as client: + response: Final = client.get(path) + + assert response.status_code == 500 + assert response.text == "handled" + assert tuple(span.name for span in exporter.get_finished_spans()) == ( + () if path == "/excluded" else ("auth", "error SDK span", "error handler", "GET /included") + ) diff --git a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py index 318859cb3ed..2d0b4701c9f 100644 --- a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py @@ -38,6 +38,7 @@ from litellm.integrations.otel.model.payloads import ( RequestIdentity, _upstream_address_port, ) +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY from litellm.integrations.otel.model.spans import ( SPAN_REGISTRY, LiteLLMSpanKind, @@ -223,6 +224,16 @@ def test_operation_resolution(): assert resolve_operation("call_mcp_tool") is GenAIOperation.EXECUTE_TOOL +def test_passthrough_operation_wins_over_the_call_type() -> None: + assert resolve_operation("pass_through_endpoint", GenAIOperation.RETRIEVAL) is GenAIOperation.RETRIEVAL + assert resolve_operation("pass_through_endpoint", "retrieval") is GenAIOperation.CHAT + event: Final = LLMCallEvent.from_dict( + {"call_type": "pass_through_endpoint", "model": "unknown", PASSTHROUGH_OPERATION_KEY: GenAIOperation.RETRIEVAL} + ) + assert event.operation is GenAIOperation.RETRIEVAL + assert event.provisional_span_name == "retrieval unknown" + + @pytest.mark.parametrize("call_type", ["vector_store_search", "avector_store_search"]) def test_vector_store_search_is_a_retrieval_operation(call_type): """A vector-store search is a retrieval, so its duration and cost must not diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index f908c685382..f0fca94c370 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -48,6 +48,7 @@ from litellm.integrations.opentelemetry import ( RAW_REQUEST_SPAN_NAME, _normalize_team_metadata_keys, ) +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy import proxy_server from litellm.proxy._types import ManagementEndpointLoggingPayload, ProxyException, SpanAttributes @@ -1048,6 +1049,15 @@ class TestOpenTelemetrySemconvStability(unittest.TestCase): kwargs = {"model": "claude-sonnet-4-5", "call_type": "unknown"} self.assertEqual(h._get_span_name(kwargs), "chat claude-sonnet-4-5") + def test_passthrough_operation_names_the_span(self) -> None: + h: Final = self._make(env="gen_ai_latest_experimental") + kwargs: Final = { + "model": "unknown", + "call_type": "pass_through_endpoint", + PASSTHROUGH_OPERATION_KEY: GenAIOperation.RETRIEVAL, + } + self.assertEqual(h._get_span_name(kwargs), "retrieval unknown") + def test_generation_name_metadata_overrides_semconv_name(self): h = self._make(env="gen_ai_latest_experimental") kwargs = { 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/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 7291c4858f6..e2452a928a4 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,4 +1,3 @@ - from typing import Final import pytest @@ -19,9 +18,7 @@ def test_web_search_cost_low(): web_search_options=web_search_options, model_info=model_info ) - assert ( - cost == model_info["search_context_cost_per_query"]["search_context_size_low"] - ) + assert cost == model_info["search_context_cost_per_query"]["search_context_size_low"] def test_web_search_cost_medium(): @@ -32,10 +29,7 @@ def test_web_search_cost_medium(): web_search_options=web_search_options, model_info=model_info ) - assert ( - cost - == model_info["search_context_cost_per_query"]["search_context_size_medium"] - ) + assert cost == model_info["search_context_cost_per_query"]["search_context_size_medium"] def test_web_search_cost_high(): @@ -46,9 +40,7 @@ def test_web_search_cost_high(): web_search_options=web_search_options, model_info=model_info ) - assert ( - cost == model_info["search_context_cost_per_query"]["search_context_size_high"] - ) + assert cost == model_info["search_context_cost_per_query"]["search_context_size_high"] def test_get_web_search_options_reads_tools_when_options_are_null() -> None: @@ -60,9 +52,7 @@ def test_get_web_search_options_reads_tools_when_options_are_null() -> None: def test_get_web_search_options_accepts_non_list_tool_iterables(): - options = StandardBuiltInToolCostTracking.get_web_search_options( - {"tools": ({"type": "web_search_preview"},)} - ) + options = StandardBuiltInToolCostTracking.get_web_search_options({"tools": ({"type": "web_search_preview"},)}) assert options == {"type": "web_search_preview"} @@ -76,25 +66,15 @@ def test_get_file_search_tool_call_does_not_validate_tool_payload(): # Test file search cost calculation def test_file_search_cost(): file_search = FileSearchTool(type="file_search") - cost = StandardBuiltInToolCostTracking.get_cost_for_file_search( - file_search=file_search - ) + cost = StandardBuiltInToolCostTracking.get_cost_for_file_search(file_search=file_search) assert cost == 0.0025 # $2.50/1000 calls = 0.0025 per call # Test edge cases def test_none_inputs(): # Test with None inputs - assert ( - StandardBuiltInToolCostTracking.get_cost_for_web_search( - web_search_options=None, model_info=None - ) - == 0.0 - ) - assert ( - StandardBuiltInToolCostTracking.get_cost_for_file_search(file_search=None) - == 0.0 - ) + assert StandardBuiltInToolCostTracking.get_cost_for_web_search(web_search_options=None, model_info=None) == 0.0 + assert StandardBuiltInToolCostTracking.get_cost_for_file_search(file_search=None) == 0.0 # Test the main get_cost_for_built_in_tools method @@ -119,9 +99,7 @@ def test_get_cost_for_built_in_tools_file_search(): Test that the cost for a file search is 0.00 when no response object is provided """ model = "gpt-4" - standard_built_in_tools_params = StandardBuiltInToolsParams( - file_search=FileSearchTool(type="file_search") - ) + standard_built_in_tools_params = StandardBuiltInToolsParams(file_search=FileSearchTool(type="file_search")) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( model=model, @@ -144,9 +122,7 @@ def test_get_cost_for_anthropic_web_search_with_server_tool_use_dict(): usage = Usage(server_tool_use={"web_search_requests": 1}) assert isinstance(usage.server_tool_use, ServerToolUse) - assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( - response_object=None, usage=usage - ) + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call(response_object=None, usage=usage) def test_anthropic_web_search_zero_requests_from_raw_response_charges_zero(): @@ -257,9 +233,7 @@ def test_completion_cost_includes_web_search_without_standard_built_in_tools_par ) assert web_search_cost > 0, "Web search cost should be non-zero" - assert ( - cost >= web_search_cost - ), f"completion_cost ({cost}) should include web search cost ({web_search_cost})" + assert cost >= web_search_cost, f"completion_cost ({cost}) should include web search cost ({web_search_cost})" @pytest.mark.parametrize( @@ -286,18 +260,14 @@ def test_gemini_3x_web_search_billed_per_query(model, local_model_cost_map): web_search_requests = 2 model_info = litellm.get_model_info(model) assert model_info["web_search_billing_unit"] == "per_query" - per_query_cost = model_info["search_context_cost_per_query"][ - "search_context_size_medium" - ] + per_query_cost = model_info["search_context_cost_per_query"]["search_context_size_medium"] expected_cost = per_query_cost * web_search_requests usage = Usage( prompt_tokens=11, completion_tokens=100, total_tokens=111, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=11, web_search_requests=web_search_requests - ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=11, web_search_requests=web_search_requests), ) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( @@ -309,8 +279,7 @@ def test_gemini_3x_web_search_billed_per_query(model, local_model_cost_map): ) assert cost == pytest.approx(expected_cost), ( - f"Expected {web_search_requests} x ${per_query_cost} = ${expected_cost} " - f"per_query search fee, got ${cost}" + f"Expected {web_search_requests} x ${per_query_cost} = ${expected_cost} per_query search fee, got ${cost}" ) @@ -353,17 +322,13 @@ def test_gemini_2x_web_search_still_billed_per_prompt(local_model_cost_map): model = "vertex_ai/gemini-2.5-flash" model_info = litellm.get_model_info(model) assert not model_info.get("web_search_billing_unit") - expected_cost = model_info["search_context_cost_per_query"][ - "search_context_size_medium" - ] + expected_cost = model_info["search_context_cost_per_query"]["search_context_size_medium"] usage = Usage( prompt_tokens=11, completion_tokens=100, total_tokens=111, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=11, web_search_requests=2 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=11, web_search_requests=2), ) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( @@ -375,8 +340,7 @@ def test_gemini_2x_web_search_still_billed_per_prompt(local_model_cost_map): ) assert cost == pytest.approx(expected_cost), ( - f"Expected flat ${expected_cost} per_prompt search fee (2 queries clamped to 1), " - f"got ${cost}" + f"Expected flat ${expected_cost} per_prompt search fee (2 queries clamped to 1), got ${cost}" ) @@ -402,9 +366,7 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model( prompt_tokens=11, completion_tokens=100, total_tokens=111, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=11, web_search_requests=2 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=11, web_search_requests=2), ) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( @@ -427,7 +389,6 @@ def _openai_responses_with_web_search_calls(model, num_calls): ResponseFunctionWebSearch, ) - output = [ ResponseFunctionWebSearch( id=f"ws_{i}", @@ -458,9 +419,7 @@ def test_openai_responses_web_search_multiplied_by_call_count(local_model_cost_m from litellm.types.utils import Usage model = "gpt-4o-search-preview" - per_call = litellm.get_model_info(model)["search_context_cost_per_query"][ - "search_context_size_medium" - ] + per_call = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"] usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) for num_calls in (1, 3): @@ -487,9 +446,7 @@ def test_web_search_call_count_reads_dict_output_items(local_model_cost_map): from litellm.types.utils import Usage model = "gpt-4o-search-preview" - per_call = litellm.get_model_info(model)["search_context_cost_per_query"][ - "search_context_size_medium" - ] + per_call = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"] response = ResponsesAPIResponse.model_validate( { @@ -498,10 +455,7 @@ def test_web_search_call_count_reads_dict_output_items(local_model_cost_map): "model": model, "object": "response", "status": "completed", - "output": [ - {"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} - for i in range(3) - ], + "output": [{"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} for i in range(3)], } ) assert all(isinstance(item, dict) for item in response.output) @@ -514,9 +468,7 @@ def test_web_search_call_count_reads_dict_output_items(local_model_cost_map): standard_built_in_tools_params=None, ) - assert cost == pytest.approx(3 * per_call), ( - f"3 dict-shaped web searches must bill 3 x ${per_call}, got ${cost}" - ) + assert cost == pytest.approx(3 * per_call), f"3 dict-shaped web searches must bill 3 x ${per_call}, got ${cost}" # Note: File search integration test removed due to complex annotation detection logic @@ -596,8 +548,6 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 - - def _openai_responses_response(model, output, usage=None, tools=None): return ResponsesAPIResponse.model_validate( { @@ -815,6 +765,31 @@ def test_responses_usage_tool_usage_web_search_billed_without_output_item(local_ assert cost == pytest.approx(num_requests * per_call) +@pytest.mark.parametrize("model", ["azure/gpt-5-mini", "azure/gpt-5.6-luna"]) +def test_azure_responses_web_search_billed_per_bing_request(local_model_cost_map, model): + """Azure OpenAI web_search is Grounding with Bing: $0.014 per reported Bing request, including rows + whose OpenAI twin prices search but the Azure row used to bill it at $0.""" + for num_requests in (1, 5): + response = _openai_responses_response( + model, + [dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": {"web_search": {"num_requests": num_requests}}, + }, + ) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=None, + custom_llm_provider="azure", + standard_built_in_tools_params=None, + ) + assert cost == pytest.approx(num_requests * 0.014) + + def _image_gen_token_usage(input_text, input_image, output_image, output_text): return { "input_tokens": 10, @@ -1012,6 +987,8 @@ def test_responses_image_gen_total_without_token_details_falls_back_to_per_image ) ) assert cost > 0 + + @pytest.mark.parametrize( ("model_key", "web_search_requests", "medium_cost_per_query", "expected_cost"), [ 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_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_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py index e59252e1571..99fa849879c 100644 --- a/tests/unit/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/unit/litellm_core_utils/test_fallback_generalizations.py @@ -7,7 +7,11 @@ its end-to-end wiring into provider routing (get_llm_provider) and model-info resolution (get_model_info) including the shipped rules in the bundled cost map. """ +import json import logging +from collections.abc import Callable +from pathlib import Path +from typing import Final import pytest @@ -21,6 +25,8 @@ from litellm.litellm_core_utils.fallback_generalizations import ( set_fallback_generalizations, ) +REPO_ROOT: Final = Path(__file__).parents[3] + @pytest.fixture def restore_generalizations(): @@ -148,7 +154,7 @@ def test_fill_missing_requires_per_rule_opt_in(restore_generalizations): } ] ) - assert match_fill_missing_generalizations("acme-1", "openai") == {"supports_vision": True} + assert match_fill_missing_generalizations("acme-1", "openai") is None restore_generalizations( [ @@ -203,89 +209,13 @@ def test_reinstalling_rules_replaces_compiled_rules(restore_generalizations): # --------------------------------------------------------------------------- # -# Engine: install-time validation and legacy-schema shim +# Engine: install-time validation # --------------------------------------------------------------------------- # -def test_legacy_mixed_rule_acts_as_both_kinds(restore_generalizations): - """A legacy rule mixing ``litellm_provider`` with capability keys routes AND - contributes its full model_info (provider included) to the capability union.""" - restore_generalizations( - [ - { - "name": "legacy-mixed", - "pattern": r"^acme-", - "model_info": {"litellm_provider": "anthropic", "supports_vision": True}, - }, - {"name": "new-caps", "pattern": r"^acme-pro-", "model_info": {"supports_reasoning": True}}, - ] - ) - assert match_routing_generalization("acme-pro-1") == "anthropic" - assert match_capability_generalizations("acme-pro-1") == { - "litellm_provider": "anthropic", - "supports_vision": True, - "supports_reasoning": True, - } - - -LEGACY_MAIN_RULES = [ - { - "name": "anthropic-claude-adaptive-thinking", - "pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))", - "description": "Claude opus/sonnet/haiku at version 4.6 or higher: 4.6 through 4.99, then any 5.x, 6.x or later major. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Turns on adaptive thinking for new families with no code change.", - "extends": "anthropic-claude", - "model_info": {"supports_adaptive_thinking": True}, - }, - { - "name": "anthropic-claude", - "pattern": "^claude-[a-z]+-\\d+[-.]\\d+(?:-\\d{8})?$", - "description": "Any Claude family-major-minor id, optionally with an 8-digit date suffix, anchored to the whole name. Version-neutral fallback that gives an unmapped Claude provider routing and baseline capabilities; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.", - "model_info": { - "litellm_provider": "anthropic", - "mode": "chat", - "max_input_tokens": 200000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "supports_function_calling": True, - "supports_parallel_function_calling": True, - "supports_vision": True, - "supports_tool_choice": True, - "supports_assistant_prefill": True, - "supports_prompt_caching": True, - "supports_response_schema": True, - "supports_reasoning": True, - "supports_pdf_input": True, - "supports_system_messages": True, - }, - }, -] - - -def test_legacy_main_schema_keeps_unmapped_claude_working(restore_generalizations): - """Pins the remote-map transition window: a released proxy running this engine - against main's old-schema block (mixed provider+capability rule plus ``extends``, - copied verbatim above) must keep unmapped-Claude inference and info resolution - working until the new-schema JSON reaches main.""" - restore_generalizations([dict(rule) for rule in LEGACY_MAIN_RULES]) - litellm.get_model_info.cache_clear() - - _, provider, _, _ = litellm.get_llm_provider(model="claude-opus-9-9") - assert provider == "anthropic" - - info = litellm.get_model_info("claude-opus-9-9") - assert info["litellm_provider"] == "anthropic" - assert info["supports_adaptive_thinking"] is True - assert info["supports_function_calling"] is True - assert info["max_input_tokens"] == 200000 - assert not info.get("input_cost_per_token") - - low = litellm.get_model_info("claude-opus-4-0") - assert low["litellm_provider"] == "anthropic" - assert low["supports_function_calling"] is True - assert low.get("supports_adaptive_thinking") is None - - -def test_non_string_provider_rule_warns_and_is_skipped(restore_generalizations, warning_messages): +def test_non_string_provider_rule_warns_and_is_skipped( + restore_generalizations: Callable[[list | None], None], warning_messages: list[str] +) -> None: restore_generalizations([{"name": "bad-provider", "pattern": r"^acme-", "model_info": {"litellm_provider": 42}}]) assert any("bad-provider" in message for message in warning_messages) assert match_routing_generalization("acme-1") is None @@ -306,6 +236,40 @@ def test_malformed_rules_are_skipped_not_fatal(restore_generalizations): assert match_capability_generalizations("claude-opus-9-9") == {"good": True} +@pytest.mark.parametrize( + "rules", + [ + [ + { + "name": "stale-rule", + "pattern": r"^acme-", + "model_info": {"litellm_provider": "anthropic", "supports_vision": True}, + } + ], + [ + {"name": "parent", "pattern": r"^zeta-only$", "model_info": {"supports_vision": True}}, + { + "name": "stale-rule", + "pattern": r"^acme-", + "extends": "parent", + "model_info": {"supports_reasoning": True}, + }, + ], + ], + ids=["provider-mixed-with-capabilities", "extends"], +) +def test_old_format_rule_is_skipped_with_a_warning_pointing_at_a_current_cost_map( + restore_generalizations: Callable[[list | None], None], warning_messages: list[str], rules: list[dict] +) -> None: + restore_generalizations(rules) + assert any( + "stale-rule" in m and "LITELLM_MODEL_COST_MAP_URL" in m and "LITELLM_LOCAL_MODEL_COST_MAP " in m + for m in warning_messages + ) + assert match_routing_generalization("acme-1") is None + assert match_capability_generalizations("acme-1") is None + + # --------------------------------------------------------------------------- # # End-to-end: provider routing + model-info resolution # --------------------------------------------------------------------------- # @@ -435,6 +399,18 @@ def test_exact_entries_fill_only_missing_fields(restore_generalizations, monkeyp # --------------------------------------------------------------------------- # +@pytest.mark.parametrize( + "cost_map_path", + ["model_prices_and_context_window.json", "litellm/model_prices_and_context_window_backup.json"], +) +def test_every_shipped_rule_installs_without_a_skip_warning( + restore_generalizations: Callable[[list | None], None], warning_messages: list[str], cost_map_path: str +) -> None: + rules: Final = json.loads((REPO_ROOT / cost_map_path).read_text())["fallback_generalizations"]["rules"] + restore_generalizations(rules) + assert warning_messages == [] + + @pytest.fixture def shipped_cost_map(monkeypatch): """Activate the bundled cost map so the shipped rules are installed.""" diff --git a/tests/unit/litellm_core_utils/test_get_litellm_params.py b/tests/unit/litellm_core_utils/test_get_litellm_params.py index 89173956dc6..738ab774b2f 100644 --- a/tests/unit/litellm_core_utils/test_get_litellm_params.py +++ b/tests/unit/litellm_core_utils/test_get_litellm_params.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.get_litellm_params import ( stored_control_options, ) from litellm.types.litellm_params import ControlOptions +from litellm.utils import get_non_default_completion_params def _funnel_kwargs_completion_forwards(monkeypatch, wif_kwargs: dict[str, object]) -> dict[str, object]: @@ -411,6 +412,14 @@ class TestMetadataFallsBackToLitellmMetadata: result["metadata"].pop("trace_id") assert litellm_metadata == {"trace_id": "trace-1"} + def test_session_id_derived_from_trace_id_adds_no_provider_params(self): + from_trace_id: Final = get_litellm_params(metadata={"trace_id": "trace-1"}) + from_session_id: Final = get_litellm_params(metadata={"session_id": "session-1"}) + added_by_trace_id: Final = from_trace_id.keys() - from_session_id.keys() + + assert added_by_trace_id + assert not added_by_trace_id & get_non_default_completion_params(from_trace_id).keys() + @pytest.mark.parametrize( "value, expected", 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_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/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index ef002a73e52..4d5d6cc17a9 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -2378,10 +2378,8 @@ def test_anthropic_model_supports_effort_param_rejects_non_supporting_models(mod "model", [ "claude-opus-4-6", - "claude-opus-4-7", "claude-opus-4-8", "claude-opus-4-6-20260205", - "claude-opus-4-7-20260416", ], ) def test_anthropic_model_supports_speed_param_recognizes_supporting_models(model): @@ -2397,6 +2395,10 @@ def test_anthropic_model_supports_speed_param_recognizes_supporting_models(model "vertex_ai/claude-opus-4-8", "azure_ai/claude-opus-4-8", "anthropic.claude-opus-4-8", + # Anthropic removed Opus 4.7 fast mode on 2026-07-24 and answers 400 for `speed` on it + # (code.claude.com/docs/en/fast-mode, platform.claude.com/docs/en/build-with-claude/fast-mode, read 2026-10-28) + "claude-opus-4-7", + "claude-opus-4-7-20260416", ], ) def test_anthropic_model_supports_speed_param_rejects_non_supporting_models(model): @@ -4677,8 +4679,9 @@ def test_fast_mode_parameter_mapping(): assert result["speed"] == "fast" -def test_anthropic_drop_params_strips_speed_for_unsupported_models(): - """``drop_params=True`` strips unsupported ``speed`` for non-Opus models.""" +@pytest.mark.parametrize("model", ["claude-sonnet-4-6", "claude-opus-4-7"]) +def test_anthropic_drop_params_strips_speed_for_unsupported_models(model: str): + """``drop_params=True`` strips ``speed`` for models without fast mode.""" config = AnthropicConfig() messages = [{"role": "user", "content": "Hello"}] @@ -4686,7 +4689,7 @@ def test_anthropic_drop_params_strips_speed_for_unsupported_models(): litellm.drop_params = True try: result = config.transform_request( - model="claude-sonnet-4-6", + model=model, messages=messages, optional_params={"speed": "fast", "max_tokens": 1024}, litellm_params={}, 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/azure/test_azure_exception_mapping.py b/tests/unit/llms/azure/test_azure_exception_mapping.py index 8670a9df2ea..cbbf70c691b 100644 --- a/tests/unit/llms/azure/test_azure_exception_mapping.py +++ b/tests/unit/llms/azure/test_azure_exception_mapping.py @@ -1,13 +1,16 @@ import json +from collections.abc import Awaitable, Callable, Mapping +from typing import Final from unittest.mock import MagicMock, patch import pytest import httpx +import openai import respx import litellm -from litellm.exceptions import ContentPolicyViolationError +from litellm.exceptions import BadRequestError, ContentPolicyViolationError, ContextWindowExceededError from litellm.litellm_core_utils.exception_mapping_utils import exception_type @@ -490,3 +493,142 @@ class TestAzureExceptionMapping: error = exc_info.value assert "encrypted_content_affinity" in error.message assert "enable_pre_call_checks" in error.message + + +_AZURE_API_BASE: Final = "https://azure-context.openai.azure.com" +_AZURE_CALL_KWARGS: Final = { + "api_base": _AZURE_API_BASE, + "api_version": "2024-10-21", + "api_key": "test-azure-key", + "max_retries": 0, +} + + +_CONTEXT_LENGTH_ERROR: Final = { + "message": "Your input exceeds the context window of this model. Please adjust your input and try again.", + "type": "invalid_request_error", + "param": "input", + "code": "context_length_exceeded", +} + + +def _azure_error_response(status_code: int, error: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + status_code, text=json.dumps({"error": error}, indent=2), headers={"content-type": "application/json"} + ) + + +async def _azure_chat(*, stream: bool) -> None: + messages: Final = [{"role": "user", "content": "hi"}] + if stream: + stream_response: Final = await litellm.acompletion( + model="azure/gpt-4o", messages=messages, stream=True, **_AZURE_CALL_KWARGS + ) + assert isinstance(stream_response, litellm.CustomStreamWrapper) + async for _chunk in stream_response: + pass + else: + litellm.completion(model="azure/gpt-4o", messages=messages, **_AZURE_CALL_KWARGS) + + +async def _azure_responses(*, is_async: bool) -> None: + if is_async: + await litellm.aresponses(model="azure/gpt-4o", input="hi", **_AZURE_CALL_KWARGS) + else: + litellm.responses(model="azure/gpt-4o", input="hi", **_AZURE_CALL_KWARGS) + + +async def _azure_embedding() -> None: + litellm.embedding(model="azure/text-embedding-3-small", input="hi", **_AZURE_CALL_KWARGS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call", + [ + pytest.param(lambda: _azure_chat(stream=False), id="chat"), + pytest.param(lambda: _azure_chat(stream=True), id="async-chat-stream"), + pytest.param(lambda: _azure_responses(is_async=False), id="responses"), + pytest.param(lambda: _azure_responses(is_async=True), id="async-responses"), + pytest.param(_azure_embedding, id="embedding"), + ], +) +@pytest.mark.parametrize( + ("code", "expected"), + [("context_length_exceeded", ContextWindowExceededError), ("invalid_value", BadRequestError)], +) +@respx.mock +async def test_azure_400_is_classified_by_error_code( + monkeypatch: pytest.MonkeyPatch, call: Callable[[], Awaitable[None]], code: str, expected: type[Exception] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx.post(url__startswith=_AZURE_API_BASE).mock( + return_value=_azure_error_response(400, {**_CONTEXT_LENGTH_ERROR, "code": code}) + ) + + with pytest.raises(BadRequestError) as exc_info: + await call() + + assert type(exc_info.value) is expected + + +@pytest.mark.asyncio +@respx.mock +async def test_azure_context_length_code_triggers_context_window_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx.post(url__startswith=_AZURE_API_BASE).mock( + return_value=_azure_error_response(400, _CONTEXT_LENGTH_ERROR) + ) + router: Final = litellm.Router( + model_list=[ + {"model_name": "small", "litellm_params": {"model": "azure/gpt-4o", **_AZURE_CALL_KWARGS}}, + {"model_name": "large", "litellm_params": {"model": "gpt-4.1", "mock_response": "served by large"}}, + ], + context_window_fallbacks=[{"small": ["large"]}], + num_retries=0, + ) + + response: Final = await router.acompletion( # pyright: ignore[reportUnknownMemberType] # untyped **kwargs + model="small", messages=[{"role": "user", "content": "hi"}] + ) + + assert isinstance(response, litellm.ModelResponse) + assert response.choices[0].message.content == "served by large" + assert len(route.calls) == 1 + + +def _azure_chat_error(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> openai.APIError: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx.post(url__startswith=_AZURE_API_BASE).mock(return_value=response) + with pytest.raises(openai.APIError) as exc_info: + litellm.completion(model="azure/gpt-4o", messages=[{"role": "user", "content": "hi"}], **_AZURE_CALL_KWARGS) + return exc_info.value + + +@respx.mock +def test_azure_context_window_error_keeps_upstream_code_and_param(monkeypatch: pytest.MonkeyPatch) -> None: + error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, _CONTEXT_LENGTH_ERROR)) + + assert isinstance(error, ContextWindowExceededError) + assert error.code == "context_length_exceeded" + assert error.param == "input" + + +@pytest.mark.parametrize("status_code", [429, 500]) +@respx.mock +def test_azure_non_400_with_context_length_code_keeps_its_status( + monkeypatch: pytest.MonkeyPatch, status_code: int +) -> None: + error: Final = _azure_chat_error(monkeypatch, _azure_error_response(status_code, _CONTEXT_LENGTH_ERROR)) + + assert not isinstance(error, ContextWindowExceededError) + assert getattr(error, "status_code", None) == status_code + + +@respx.mock +def test_azure_content_policy_wins_over_context_length_code(monkeypatch: pytest.MonkeyPatch) -> None: + policy_error: Final = {**_CONTEXT_LENGTH_ERROR, "inner_error": {"code": "ResponsibleAIPolicyViolation"}} + + error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, policy_error)) + + assert type(error) is ContentPolicyViolationError diff --git a/tests/unit/llms/base_llm/responses/test_transformation.py b/tests/unit/llms/base_llm/responses/test_transformation.py index 8d21af83668..50d918fa5be 100644 --- a/tests/unit/llms/base_llm/responses/test_transformation.py +++ b/tests/unit/llms/base_llm/responses/test_transformation.py @@ -1,14 +1,17 @@ """The shared Responses API config contract.""" import asyncio, importlib, json +from collections.abc import Mapping import httpx import pytest +import respx import litellm from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.router import GenericLiteLLMParams from litellm import Router +from pydantic import BaseModel, ConfigDict from litellm.constants import STREAM_SSE_DONE_STRING from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig @@ -879,38 +882,6 @@ def test_combine_responses_fallback_usage_passthrough_for_unknown_event(): partial = ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2) Router._combine_responses_fallback_usage(other, partial) -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_build_responses_continuation_input_from_string(): - out = Router._build_responses_continuation_input("Hello world", "partial assistant text") - assert len(out) == 3 - assert out[0]["role"] == "user" - assert out[0]["content"][0]["text"] == "Hello world" - assert out[1]["role"] == "developer" - assert out[2]["role"] == "assistant" - assert out[2]["content"][0]["text"] == "partial assistant text" - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_build_responses_continuation_input_from_list_preserves_items(): - existing: List[Any] = [ - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "msg1"}], - } - ] - out = Router._build_responses_continuation_input(existing, "partial") - assert len(out) == 3 - assert out[0]["content"][0]["text"] == "msg1" - assert out[1]["role"] == "developer" - assert out[2]["role"] == "assistant" - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_build_responses_continuation_input_from_none(): - out = Router._build_responses_continuation_input(None, "partial") - assert len(out) == 2 - assert out[0]["role"] == "developer" - assert out[1]["role"] == "assistant" - @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_aresponses_streaming_iterator_passthrough(): @@ -1355,9 +1326,573 @@ async def test_aresponses_fallback_uses_continuation_input_after_partial_content continuation = mock_fallback.await_args.kwargs["kwargs"]["input"] assert isinstance(continuation, list) assert continuation[0]["content"][0]["text"] == "original question" - assert continuation[-2]["role"] == "developer" - assert continuation[-1]["role"] == "assistant" - assert continuation[-1]["content"][0]["text"] == "partial answer" + assert continuation[-2]["role"] == "assistant" + assert continuation[-2]["content"][0]["text"] == "partial answer" + assert continuation[-1]["role"] == "user" + + +def _sse(*events: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events), + ) + + +def _openai_stream_cut_after(partial: str) -> httpx.Response: + return _sse( + *( + ({ + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": partial, + },) + if partial else () + ), + { + "type": "error", "sequence_number": 1, + "error": {"type": "server_error", "code": "internal_error", "message": "interrupted"}, + }, + ) + + +def _anthropic_stream(text: str) -> httpx.Response: + return _sse( + {"type": "message_start", "message": { + "id": "msg_fallback", "type": "message", "role": "assistant", + "model": "claude-fallback", "content": [], "usage": {"input_tokens": 10, "output_tokens": 0}, + }}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}, + {"type": "message_stop"}, + ) + + +def _primary_with_fallback(fallback_model: str) -> Router: + return Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "openai/gpt-6-astra", "api_key": "synthetic"}}, + {"model_name": "fallback", "litellm_params": {"model": fallback_model, "api_key": "synthetic"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + + +def _streamed_text(events: tuple[Any, ...]) -> str: + return "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) + + +@pytest.fixture +def claude_with_prefill() -> str: + model: Final = "anthropic/claude-prefill-capable-test-model" + litellm.register_model( + {model: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}}, + persist_across_reloads=False, + ) + litellm.get_model_info.cache_clear() + return model + + +@pytest.mark.asyncio +@pytest.mark.parametrize("partial", ["", "1, 2, "]) +async def test_responses_fallback_can_continue_on_a_provider_without_prefill( + partial: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + prompt: Final = "Print the integers 1 through 5, separated by commas." + primary: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_after(partial) + ) + + def no_prefill_provider(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + if body["messages"][-1]["role"] != "user": + return httpx.Response( + 400, + json={"type": "error", "error": { + "type": "invalid_request_error", "message": "The conversation must end with a user message", + }}, + ) + return _anthropic_stream("3, 4, 5" if partial else "1, 2, 3, 4, 5") + + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=no_prefill_provider) + router: Final = _primary_with_fallback("anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses(model="primary", input=prompt, stream=True, max_output_tokens=100) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events) == "1, 2, 3, 4, 5" + assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert primary.call_count == fallback.call_count == 1 + messages: Final = json.loads(fallback.calls.last.request.content)["messages"] + assert messages[0] == {"role": "user", "content": [{"type": "text", "text": prompt}]} + if partial: + assert [message["role"] for message in messages] == ["user", "assistant", "user"] + assert messages[-2] == {"role": "assistant", "content": [{"type": "text", "text": partial}]} + assert json.dumps(partial) in messages[-1]["content"][0]["text"] + else: + assert len(messages) == 1 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_continues_a_claude_that_takes_prefill_from_its_own_partial( + claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, ")) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream(" 3, 4, 5") + ) + router: Final = _primary_with_fallback(claude_with_prefill) + try: + stream: Final = await router.aresponses( + model="primary", input="Print the integers 1 through 5.", stream=True, max_output_tokens=100 + ) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events) == "1, 2, 3, 4, 5" + assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + messages: Final = json.loads(fallback.calls.last.request.content)["messages"] + assert [message["role"] for message in messages] == ["user", "assistant"] + assert messages[-1]["content"][0]["text"].startswith("1, 2,") + finally: + router.discard() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback_takes_prefill", [True, False], ids=["prefill", "request_to_continue"]) +async def test_responses_fallback_continuation_drops_reasoning_the_fallback_cannot_decrypt( + fallback_takes_prefill: bool, + claude_with_prefill: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, ")) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream(" 3, 4, 5") + ) + router: Final = _primary_with_fallback(claude_with_prefill if fallback_takes_prefill else "anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses( + model="primary", + input=[ + {"role": "user", "content": "Print the integers 1 through 5."}, + {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "primary-only-ciphertext"}, + ], + stream=True, + max_output_tokens=100, + ) + _ = [event async for event in stream] + assert fallback.call_count == 1 + assert "primary-only-ciphertext" not in fallback.calls.last.request.content.decode() + finally: + router.discard() + + +class _Numbers(BaseModel): + model_config = ConfigDict(frozen=True) + + numbers: tuple[int, ...] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "constraint", + [ + {"text": {"format": {"type": "json_schema", "name": "numbers", "strict": True, "schema": {"type": "object"}}}}, + {"text": {"format": {"type": "json_object"}}}, + {"tool_choice": "required", "tools": [{"type": "function", "name": "lookup", "parameters": {"type": "object"}}]}, + {"text_format": _Numbers}, + ], + ids=["json_schema", "json_object", "required_tool", "sdk_text_format"], +) +async def test_responses_fallback_surfaces_the_primary_error_once_constrained_output_was_streamed( + constraint: dict[str, object], + claude_with_prefill: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_after('{"numbers": [1, 2, ') + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("3]}") + ) + router: Final = _primary_with_fallback(claude_with_prefill) + try: + stream: Final = await router.aresponses( + model="primary", input="Numbers 1 to 3 as JSON.", stream=True, max_output_tokens=100, **constraint + ) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_continues_within_the_callers_output_ceiling( + claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, ")) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("3, 4, 5") + ) + router: Final = _primary_with_fallback(claude_with_prefill) + try: + stream: Final = await router.aresponses(model="primary", input="Count to 5.", stream=True, max_output_tokens=100) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events) == "1, 2, 3, 4, 5" + assert json.loads(fallback.calls.last.request.content)["max_tokens"] == 100 - litellm.token_counter( + text="1, 2, " + ) + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_used_the_output_ceiling( + claude_with_prefill: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(return_value=_openai_stream_cut_after("1, 2, ")) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("3, 4, 5") + ) + router: Final = _primary_with_fallback(claude_with_prefill) + try: + stream: Final = await router.aresponses( + model="primary", input="Count to 5.", stream=True, max_output_tokens=litellm.token_counter(text="1, 2, ") + ) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +def _openai_stream_cut_mid_tool_call(preamble: str, *, announce_item: bool = True) -> httpx.Response: + text: Final = { + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": preamble, + } + item_added: Final = { + "type": "response.output_item.added", "sequence_number": 1, "output_index": 1, + "item": { + "type": "function_call", "id": "fc_primary", "call_id": "call_primary", + "name": "lookup", "arguments": "", "status": "in_progress", + }, + } + arguments: Final = { + "type": "response.function_call_arguments.delta", "sequence_number": 2, + "item_id": "fc_primary", "output_index": 1, "delta": '{"city": "Par', + } + error: Final = { + "type": "error", "sequence_number": 3, + "error": {"type": "server_error", "code": "internal_error", "message": "interrupted"}, + } + return _sse(*((text, item_added, arguments, error) if announce_item else (text, arguments, error))) + + +def _anthropic_stream_cut_mid_tool_call(preamble: str, tool_block: Mapping[str, object] | None = None) -> httpx.Response: + return _sse( + {"type": "message_start", "message": { + "id": "msg_primary", "type": "message", "role": "assistant", + "model": "claude-primary", "content": [], "usage": {"input_tokens": 10, "output_tokens": 0}, + }}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": preamble}}, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": tool_block or { + "type": "tool_use", "id": "toolu_primary", "name": "lookup", "input": {}, + }}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": '{"city": "Par'}}, + {"type": "error", "error": {"type": "overloaded_error", "message": "interrupted"}}, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("announce_item", [True, False], ids=["announced_item", "arguments_only"]) +async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_tool_call( + announce_item: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_mid_tool_call("Let me look that up. ", announce_item=announce_item) + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("It is sunny in Paris.") + ) + router: Final = _primary_with_fallback("anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses( + model="primary", + input="What is the weather in Paris?", + stream=True, + tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}], + ) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +def _openai_stream_cut_after_event(preamble: str, event: Mapping[str, object]) -> httpx.Response: + return _sse( + { + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": preamble, + }, + {**event, "sequence_number": 1}, + { + "type": "error", "sequence_number": 2, + "error": {"type": "server_error", "code": "internal_error", "message": "interrupted"}, + }, + ) + + +def _openai_stream_cut_after_output_item(preamble: str, item: Mapping[str, object]) -> httpx.Response: + return _openai_stream_cut_after_event( + preamble, {"type": "response.output_item.added", "output_index": 1, "item": item} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "item", + [ + {"type": "shell_call", "id": "sh_1", "call_id": "call_1", "action": {"commands": ["ls"]}, "status": "in_progress"}, + { + "type": "apply_patch_call", "id": "apc_1", "call_id": "call_1", "status": "in_progress", + "operation": {"type": "update_file", "path": "a.py", "diff": "@@"}, + }, + { + "type": "local_shell_call", "id": "lsh_1", "call_id": "call_1", "status": "in_progress", + "action": {"type": "exec", "command": ["ls"], "env": {}}, + }, + { + "type": "computer_call", "id": "cu_1", "call_id": "call_1", "status": "in_progress", + "action": {"type": "screenshot"}, "pending_safety_checks": [], + }, + {"type": "mcp_approval_request", "id": "mcpr_1", "name": "lookup", "arguments": "{}", "server_label": "docs"}, + ], + ids=lambda item: item["type"], +) +async def test_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_client_run_tool_call( + item: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_after_output_item("Let me check the repo. ", item) + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("The repo has two files.") + ) + router: Final = _primary_with_fallback("anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses(model="primary", input="What is in the repo?", stream=True) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server_run_event", + [ + { + "type": "response.output_item.added", "output_index": 1, + "item": { + "type": "mcp_call", "id": "mcp_1", "name": "ask_question", "arguments": "{}", + "server_label": "docs", "status": "in_progress", + }, + }, + {"type": "response.mcp_call_arguments.delta", "output_index": 1, "item_id": "mcp_1", "delta": '{"question": "set'}, + ], + ids=["mcp_call_item", "mcp_call_arguments"], +) +async def test_responses_fallback_continues_after_a_server_run_mcp_call( + server_run_event: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_openai_stream_cut_after_event("Let me check the docs. ", server_run_event) + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream("The docs cover setup.") + ) + router: Final = _primary_with_fallback("anthropic/claude-opus-5-5") + try: + stream: Final = await router.aresponses(model="primary", input="What do the docs cover?", stream=True) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events).endswith("The docs cover setup.") + assert fallback.call_count == 1 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_continues_when_allowed_tools_in_auto_mode_leaves_the_reply_open( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + openai: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + _openai_stream_cut_after("1, 2, "), + _sse( + { + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_fallback", "output_index": 0, "content_index": 0, "delta": "3, 4, 5", + }, + { + "type": "response.completed", "sequence_number": 1, + "response": { + "id": "resp_fallback", "object": "response", "created_at": 0, "status": "completed", + "model": "gpt-6.1-terra", "output": [], + "usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10}, + }, + }, + ), + ] + ) + router: Final = _primary_with_fallback("openai/gpt-6.1-terra") + try: + stream: Final = await router.aresponses( + model="primary", + input="Print the integers 1 through 5.", + stream=True, + tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}], + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "lookup"}]}, + ) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events) == "1, 2, 3, 4, 5" + assert openai.call_count == 2 + continued_input: Final = json.loads(openai.calls.last.request.content)["input"] + assert [item["role"] for item in continued_input[-2:]] == ["assistant", "user"] + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_responses_fallback_surfaces_a_content_filter_that_cut_the_stream_mid_answer( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=_sse( + { + "type": "response.output_text.delta", "sequence_number": 0, + "item_id": "msg_primary", "output_index": 0, "content_index": 0, "delta": "Step one is to obtain the", + }, + { + "type": "error", "sequence_number": 1, + "error": { + "type": "invalid_request_error", "code": "content_policy_violation", + "message": "The response was halted mid-stream.", "param": None, + }, + }, + ) + ) + fallback: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_stream(" reagents.") + ) + router: Final = Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "openai/gpt-6-astra", "api_key": "synthetic"}}, + {"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + content_policy_fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + try: + stream: Final = await router.aresponses(model="primary", input="How do I make it?", stream=True) + with pytest.raises(litellm.ContentPolicyViolationError, match="halted mid-stream"): + [event async for event in stream] + assert fallback.call_count == 0 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_bridged_responses_fallback_surfaces_the_primary_error_once_the_cut_stream_began_a_tool_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + provider: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + side_effect=[ + _anthropic_stream_cut_mid_tool_call("Let me look that up. "), + _anthropic_stream("It is sunny in Paris."), + ] + ) + router: Final = Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "anthropic/claude-primary", "api_key": "synthetic"}}, + {"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + try: + stream: Final = await router.aresponses( + model="primary", + input="What is the weather in Paris?", + stream=True, + tools=[{"type": "function", "name": "lookup", "parameters": {"type": "object"}}], + ) + with pytest.raises(litellm.InternalServerError, match="interrupted"): + [event async for event in stream] + assert provider.call_count == 1 + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_bridged_responses_fallback_continues_after_a_server_run_tool_search( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + provider: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + side_effect=[ + _anthropic_stream_cut_mid_tool_call( + "Let me find the right tool. ", + {"type": "server_tool_use", "id": "srvtoolu_primary", "name": "tool_search_tool_regex", "input": {}}, + ), + _anthropic_stream("It is sunny in Paris."), + ] + ) + router: Final = Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "anthropic/claude-primary", "api_key": "synthetic"}}, + {"model_name": "fallback", "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_key": "synthetic"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + try: + stream: Final = await router.aresponses( + model="primary", + input="What is the weather in Paris?", + stream=True, + tools=[ + {"type": "tool_search", "execution": "server"}, + {"type": "function", "name": "lookup", "parameters": {"type": "object"}, "defer_loading": True}, + ], + ) + events: Final = tuple([event async for event in stream]) + assert _streamed_text(events).endswith("It is sunny in Paris.") + assert provider.call_count == 2 + finally: + router.discard() + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 204b458cadc..2971405885d 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -4,15 +4,19 @@ and non-string objects (e.g. Pydantic BaseModel events from the Responses API) p """ import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping +from itertools import chain from typing import Final import httpx import pytest from pydantic import BaseModel +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, convert_model_response_to_streaming from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler +from litellm.llms.openai.common_utils import OpenAIError from litellm.types.utils import Choices, GenericStreamingChunk, Message, ModelResponse @@ -336,3 +340,98 @@ def test_fake_stream_chunk_indexes_thinking_blocks_by_position_without_touching_ {"type": "thinking", "thinking": "Step 2.", "signature": "sig-b", "index": 2}, ] assert [block.get("index") for block in response.choices[0].message.thinking_blocks] == [None, None, None] + + +def _sse_chunk(delta: Mapping[str, object], finish_reason: str | None = None) -> str: + return "data: " + json.dumps( + { + "id": "chatcmpl-done", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-small-latest", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + ) + + +async def _aiter_lines(lines: tuple[str, ...]) -> AsyncIterator[str]: + for line in lines: + yield line + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_done_inside_streamed_json_is_content_not_the_end_of_stream(sync_mode: bool) -> None: + lines: Final = ( + _sse_chunk({"role": "assistant", "reasoning_content": "Mark it [DONE] once tests pass"}), + _sse_chunk({"content": "Checklist: "}), + _sse_chunk({"content": "- [DONE] write tests"}), + _sse_chunk({"content": " - [ ] ship it"}), + _sse_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "set_status", "arguments": ""}, + } + ] + } + ), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '{"status": '}}]}), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '"[DONE]"'}}]}), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": ', "n": 1}'}}]}), + _sse_chunk({}, "tool_calls"), + "data: [DONE]", + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=OpenAIChatCompletionStreamingHandler( + streaming_response=iter(lines) if sync_mode else _aiter_lines(lines), sync_stream=sync_mode + ), + model="mistral/mistral-small-latest", + custom_llm_provider="mistral", + logging_obj=Logging( + model="mistral/mistral-small-latest", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=1_790_000_000.0, + litellm_call_id="call-id", + function_id="function-id", + ), + ) + + chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + + deltas: Final = [chunk.choices[0].delta for chunk in chunks if chunk.choices] + tool_calls: Final = tuple(chain.from_iterable(delta.tool_calls or () for delta in deltas)) + assert "".join(getattr(delta, "reasoning_content", None) or "" for delta in deltas) == ( + "Mark it [DONE] once tests pass" + ) + assert "".join(delta.content or "" for delta in deltas) == "Checklist: - [DONE] write tests - [ ] ship it" + assert json.loads("".join(tool_call.function.arguments or "" for tool_call in tool_calls)) == { + "status": "[DONE]", + "n": 1, + } + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [ + "tool_calls" + ] + + +@pytest.mark.parametrize("line", ["data: [DONE]", "data:[DONE]", "[DONE]", "data: [DONE]", "data: [DONE] "]) +def test_done_marker_ends_the_stream(line: str) -> None: + chunks: Final = list(BaseModelResponseIterator(streaming_response=iter((line,)), sync_stream=True)) + + assert chunks == [ + GenericStreamingChunk(text="", is_finished=True, finish_reason="stop", usage=None, index=0, tool_use=None) + ] + + +def test_error_frame_mentioning_done_raises_instead_of_ending_the_stream() -> None: + line: Final = 'data: {"error": {"message": "prompt too long [DONE]", "code": 400}}' + + with pytest.raises(OpenAIError, match=r"prompt too long \[DONE\]") as error: + list(OpenAIChatCompletionStreamingHandler(streaming_response=iter((line,)), sync_stream=True)) + + assert error.value.status_code == 400 diff --git a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py index 301fdaeb245..9ee8c14f5d6 100644 --- a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py +++ b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py @@ -121,3 +121,105 @@ def test_a_rejection_logs_the_retry_and_advises_opting_out_only_when_native_outp assert _MODEL in caplog.text assert ("supports_bedrock_native_structured_output: false" in caplog.text) is advises_opting_out + + +_STRICT_FALLBACK: Final = NativeOutputFallback(optional_params={"tools": "without_strict"}, litellm_params={}) +# Bedrock's 400s for strict tools, seen live on us.anthropic.claude-sonnet-4-6 from 2026-10-20 to 2026-10-28 and quoted +# in issue #34388, each answered 200 once the request dropped strict +_STRICT_TOOLS_REJECTIONS: Final = tuple( + BedrockError(status_code=400, message=f'{{"message":"The model returned the following errors: {text}"}}') + for text in ( + "Too many strict tools (21). The maximum number of strict tools supported is 20.", + "Schemas contains too many optional parameters (27), which would make grammar compilation inefficient.", + "Schemas contains too many parameters with union types (18 parameters with type arrays or anyOf). " + "This causes exponential compilation cost.", + "The compiled grammar is too large, which would cause performance issues.", + "tools.0.custom: Invalid schema: Schema type is missing for schema: {'description': 'tag'}", + "tools.3.custom: For 'integer' type, property 'minimum' is not supported", + "Compiled grammar size (410.4MB) exceeds maximum allowed size (300MB)", + "tools.1.custom.strict: Extra inputs are not permitted", + "Grammar compilation timed out.", + ) +) +_TOO_MANY_STRICT_TOOLS: Final = _STRICT_TOOLS_REJECTIONS[0] +# Bedrock's 400s for requests with a strict tool that fail for other reasons, seen live on +# us.anthropic.claude-sonnet-4-6 on 2026-10-21, which echo tool names and tool call ids +_OTHER_400S_WITH_A_STRICT_TOOL: Final = tuple( + BedrockError(status_code=400, message=f'{{"message":"{text}"}}') + for text in ( + "The tool check_grammar is already defined at toolConfig.tools.1.", + "Expected toolResult blocks at messages.0.content for the following Ids: call_grammar_9", + "tool_use ids were found without tool_result blocks immediately after: call_grammar_1", + ) +) + + +def _call_with_strict_tools_fallback( + bedrock: ScriptedBedrock, fallback: NativeOutputFallback | None, strict_tools_fallback: NativeOutputFallback | None +) -> object: + return _resolve( + call_with_native_output_fallback( + bedrock, + None if fallback is None else lambda: fallback, + _MODEL, + strict_tools_fallback=None if strict_tools_fallback is None else lambda: strict_tools_fallback, + ) + ) + + +@_MODES +@pytest.mark.parametrize("rejection", _STRICT_TOOLS_REJECTIONS, ids=lambda error: error.message[60:100]) +def test_a_strict_tools_rejection_is_sent_again_once_with_the_strict_tools_fallback( + asynchronous: bool, rejection: BedrockError +) -> None: + bedrock: Final = ScriptedBedrock(errors=(rejection, None), asynchronous=asynchronous) + + response: Final = _call_with_strict_tools_fallback(bedrock, None, _STRICT_FALLBACK) + + assert isinstance(response, ModelResponse) and response.id == "answer-2" + assert bedrock.sent == (None, _STRICT_FALLBACK) + + +@_MODES +@pytest.mark.parametrize( + ("error", "strict_tools_fallback"), + [ + (_OTHER_400, _STRICT_FALLBACK), + *((error, _STRICT_FALLBACK) for error in _OTHER_400S_WITH_A_STRICT_TOOL), + (BedrockError(status_code=500, message=_TOO_MANY_STRICT_TOOLS.message), _STRICT_FALLBACK), + (_TOO_MANY_STRICT_TOOLS, None), + ], + ids=[ + "other_400", + "duplicate_tool_name", + "missing_tool_result", + "orphan_tool_use", + "server_error", + "no_strict_tools_fallback", + ], +) +def test_no_second_request_without_both_a_strict_tools_rejection_and_a_strict_tools_fallback( + asynchronous: bool, error: BedrockError, strict_tools_fallback: NativeOutputFallback | None +) -> None: + bedrock: Final = ScriptedBedrock(errors=(error,), asynchronous=asynchronous) + + with pytest.raises(BedrockError) as raised: + _call_with_strict_tools_fallback(bedrock, None, strict_tools_fallback) + + assert raised.value is error + assert bedrock.sent == (None,) + + +@pytest.mark.parametrize( + ("error", "retried_with"), + [(_REJECTION, _FALLBACK), (_TOO_MANY_STRICT_TOOLS, _STRICT_FALLBACK)], + ids=["native_output_rejection", "strict_tools_rejection"], +) +def test_each_rejection_is_sent_again_with_its_own_fallback_when_both_are_set( + error: BedrockError, retried_with: NativeOutputFallback +) -> None: + bedrock: Final = ScriptedBedrock(errors=(error, None), asynchronous=False) + + _call_with_strict_tools_fallback(bedrock, _FALLBACK, _STRICT_FALLBACK) + + assert bedrock.sent == (None, retried_with) 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/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 7bde1805e18..dd89d9c46b6 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -6,7 +6,9 @@ from typing import Final from litellm.llms.bedrock.common_utils import( AmazonBedrockGlobalConfig, BedrockModelInfo, + NativeOutputFallback, bedrock_native_output_is_opt_in, + bedrock_strict_tools_fallback, deployment_model_info, extract_model_name_from_bedrock_arn, get_bedrock_base_model, @@ -2150,3 +2152,80 @@ def test_bedrock_response_stream_shape_is_unavailable_when_botocore_describes_th assert shape is None (logged,) = warning.call_args_list assert "only-in-the-description" not in logged.args[0] % logged.args[1:] + + +def test_bedrock_strict_tools_fallback_drops_only_strict_and_keeps_everything_else(): + strict_tool: Final = { + "type": "function", + "cache_control": {"type": "ephemeral"}, + "function": {"name": "lookup", "description": "Look up", "strict": True, "parameters": {"type": "object"}}, + } + plain_tool: Final = {"type": "function", "function": {"name": "plain", "parameters": {"type": "object"}}} + anthropic_tool: Final = {"name": "native", "input_schema": {"type": "object"}} + bedrock_tool: Final = {"toolSpec": {"name": "spec", "strict": True, "inputSchema": {"json": {"type": "object"}}}} + litellm_params: Final = {"drop_params": True} + + fallback: Final = bedrock_strict_tools_fallback( + "us.anthropic.claude-sonnet-4-6", + {"tools": [strict_tool, plain_tool, anthropic_tool, bedrock_tool], "maxTokens": 16, "stream": True}, + litellm_params, + ) + + assert fallback == NativeOutputFallback( + optional_params={ + "tools": [ + {**strict_tool, "function": {"name": "lookup", "description": "Look up", "parameters": {"type": "object"}}}, + plain_tool, + anthropic_tool, + {"toolSpec": {"name": "spec", "inputSchema": {"json": {"type": "object"}}}}, + ], + "maxTokens": 16, + "stream": True, + }, + litellm_params=litellm_params, + ) + assert strict_tool["function"]["strict"] is True + assert bedrock_tool["toolSpec"]["strict"] is True + + +@pytest.mark.parametrize( + "optional_params", + [ + {"maxTokens": 16}, + {"tools": [{"type": "function", "function": {"name": "plain", "parameters": {}}}]}, + {"tools": [{"toolSpec": {"name": "spec", "inputSchema": {"json": {}}}}]}, + ], + ids=["no_tools", "no_strict_tool", "no_strict_tool_spec"], +) +def test_bedrock_strict_tools_fallback_is_none_without_a_strict_tool(optional_params: dict[str, object]) -> None: + assert bedrock_strict_tools_fallback("us.anthropic.claude-sonnet-4-6", optional_params, {}) is None + + +_CLAUDE_PROFILE_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude" + + +@pytest.mark.parametrize( + ("model", "litellm_params"), + [ + ("us.anthropic.claude-sonnet-4-6", {}), + ("us.amazon.nova-pro-v1:0", {}), + (_CLAUDE_PROFILE_ARN, {"base_model": "bedrock/us.anthropic.claude-sonnet-4-6"}), + (_CLAUDE_PROFILE_ARN, {}), + ], + ids=["claude", "nova", "claude_profile_with_base_model", "profile_without_base_model"], +) +@pytest.mark.parametrize("strict", [True, False]) +def test_bedrock_strict_tools_fallback_retries_only_when_the_request_sent_strict( + model: str, litellm_params: dict[str, object], strict: bool +) -> None: + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + tools: Final = [{"type": "function", "function": {"name": "lookup", "strict": strict, "parameters": {}}}] + request: Final = AmazonConverseConfig().transform_request( + model, [{"role": "user", "content": "hi"}], {"tools": tools}, {**litellm_params}, {} + ) + sent_strict: Final = any("strict" in tool["toolSpec"] for tool in request["toolConfig"]["tools"]) + + fallback: Final = bedrock_strict_tools_fallback(model, {"tools": tools}, litellm_params) + + assert (fallback is not None) == sent_strict diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 8985d0b7fd8..a6f7d613f28 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -16,6 +16,7 @@ import pytest import litellm from litellm.llms.chatgpt.authenticator import Authenticator +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig from litellm.llms.openai.common_utils import OpenAIError from litellm.main import responses_api_bridge_check @@ -584,3 +585,54 @@ class TestChatGPTResponsesMultiAccount: assert headers_b["ChatGPT-Account-Id"] == "acct-b" assert headers_default["Authorization"] == "Bearer token-default" assert headers_default["ChatGPT-Account-Id"] == "acct-default" + + +@pytest.mark.parametrize("call", ["responses", "completion"]) +def test_prompt_cache_key_keeps_one_session_header_across_requests( + call: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "token", "account_id": "acct", "expires_at": 4102444800}) + ) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + completed: Final = { + "type": "response.completed", + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.5", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + } + upstream: Final = MagicMock( + side_effect=lambda _request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode(), + ) + ) + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(upstream))) + for _ in range(2): + if call == "responses": + litellm.responses(model="chatgpt/gpt-5.5", input="hi", prompt_cache_key="conv-123", client=client) + else: + litellm.completion( + model="chatgpt/gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + prompt_cache_key="conv-123", + client=client, + ) + + assert [call.args[0].headers["session_id"] for call in upstream.call_args_list] == ["conv-123", "conv-123"] diff --git a/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py index 6b280d97089..375c140187e 100644 --- a/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py +++ b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py @@ -1,11 +1,17 @@ import base64 import json +from collections.abc import Mapping +from typing import Final +import h11 import pytest +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.llms.chatgpt.common_utils import ( extract_chatgpt_account_id, get_chatgpt_client_credential, + get_chatgpt_session_id, is_chatgpt_oauth_key, select_chatgpt_client_credential_headers, ) @@ -88,3 +94,72 @@ def test_bearer_scheme_is_case_insensitive(scheme): def test_other_authentication_schemes_are_not_chatgpt_logins(): assert not is_chatgpt_oauth_key(f"Basic {CHATGPT_JWT}") + + +class TestGetChatGPTSessionId: + @pytest.mark.parametrize( + "explicit", + [ + {"litellm_session_id": "chosen"}, + {"session_id": "chosen"}, + {"metadata": {"session_id": "chosen"}}, + get_litellm_params(litellm_session_id="chosen", metadata={"trace_id": "chosen"}), + ], + ) + def test_explicit_session_id_wins_over_prompt_cache_key(self, explicit: Mapping[str, object]) -> None: + assert get_chatgpt_session_id({**explicit, "prompt_cache_key": "conv-123"}) == "chosen" + + @pytest.mark.parametrize( + "request_scoped", + [ + {"litellm_trace_id": "trace-1"}, + {"litellm_call_id": "call-1"}, + get_litellm_params(metadata={"trace_id": "trace-1"}), + get_litellm_params(litellm_metadata={"trace_id": "trace-1"}), + {"litellm_session_id": "gen-1", "metadata": {"session_id": "gen-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + { + "litellm_session_id": "gen-1", + "litellm_metadata": {"session_id": "gen-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, + ], + ) + def test_prompt_cache_key_wins_over_request_scoped_ids(self, request_scoped: Mapping[str, object]) -> None: + assert get_chatgpt_session_id({**request_scoped, "prompt_cache_key": "conv-123"}) == "conv-123" + + @pytest.mark.parametrize( + "mapped_header", + [ + { + "litellm_session_id": "header-session", + "metadata": {"trace_id": "header-session", "session_id": "header-session"}, + }, + get_litellm_params(metadata={"trace_id": "header-session", "session_id": "header-session"}), + ], + ) + def test_proxy_mapped_session_header_wins_over_body_session_id(self, mapped_header: Mapping[str, object]) -> None: + assert ( + get_chatgpt_session_id({**mapped_header, "session_id": "body-session", "prompt_cache_key": "conv-123"}) + == "header-session" + ) + + @pytest.mark.parametrize("prompt_cache_key", [{"tenant": "a"}, 12345, True, ["x", "y"]]) + def test_non_string_prompt_cache_key_is_ignored(self, prompt_cache_key: object) -> None: + assert get_chatgpt_session_id({"prompt_cache_key": prompt_cache_key, "litellm_call_id": "call-1"}) == "call-1" + + @pytest.mark.parametrize( + "first, second", + [ + ("a\nb", "a\tb"), + (json.loads('"conv-\\ud800"'), json.loads('"conv-\\udc00"')), + ("conv-123 ", " conv-123"), + ], + ) + def test_header_unsafe_cache_keys_map_to_distinct_stable_header_values(self, first: str, second: str) -> None: + first_id: Final = get_chatgpt_session_id({"prompt_cache_key": first}) + second_id: Final = get_chatgpt_session_id({"prompt_cache_key": second}) + + assert first_id is not None and second_id is not None + assert first_id == get_chatgpt_session_id({"prompt_cache_key": first}) + assert first_id != second_id + for value in (first_id, second_id): + h11.Request(method="POST", target="/", headers=[("host", "chatgpt.com"), ("session_id", value)]) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 297e64629a8..f10f5a67b2f 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -66,7 +66,13 @@ from litellm.responses.streaming_iterator import ( ) from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ModelResponse, + ProviderSpecificHeader, + TranscriptionResponse, +) from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault @@ -1607,6 +1613,89 @@ async def test_async_anthropic_messages_handler_header_priority(): assert captured_headers["X-Provider-Only"] == "keep-this-too" +class _AnthropicMessagesUpstream: + def __init__(self) -> None: + self.anthropic_beta: tuple[str, ...] = () + + def respond(self, request: httpx.Request) -> httpx.Response: + self.anthropic_beta = tuple(request.headers.get_list("anthropic-beta")) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + +_BOTH_FOUNDRY_BETAS: Final = ("context-1m-2025-08-07,context-management-2025-06-27",) + + +@pytest.mark.parametrize( + ("forwarded_headers", "deployment_headers", "client_beta", "expected"), + [ + pytest.param( + None, + {"anthropic-beta": "context-1m-2025-08-07"}, + "context-management-2025-06-27", + _BOTH_FOUNDRY_BETAS, + id="deployment-and-client-betas-are-unioned", + ), + pytest.param( + None, + {"anthropic-beta": ["context-1m-2025-08-07"]}, + "context-management-2025-06-27", + _BOTH_FOUNDRY_BETAS, + id="deployment-list-and-client-string-are-unioned", + ), + pytest.param( + {"anthropic-beta": "interleaved-thinking-2025-05-14"}, + {"anthropic-beta": "context-1m-2025-08-07"}, + "context-management-2025-06-27", + ("context-1m-2025-08-07,context-management-2025-06-27,interleaved-thinking-2025-05-14",), + id="forwarded-deployment-and-client-betas-are-unioned", + ), + pytest.param( + None, + {"anthropic-beta": "context-1m-2025-08-07"}, + "claude-code-20250219", + ("context-1m-2025-08-07",), + id="unsupported-client-beta-is-dropped-and-deployment-beta-kept", + ), + ], +) +@pytest.mark.asyncio +async def test_messages_sends_deployment_and_client_anthropic_betas_upstream( + forwarded_headers: dict[str, str] | None, + deployment_headers: dict[str, str | list[str]], + client_beta: str, + expected: tuple[str, ...], +) -> None: + upstream: Final = _AnthropicMessagesUpstream() + + await litellm.anthropic.messages.acreate( + model="azure_ai/claude-sonnet-5-5", + messages=[{"role": "user", "content": "Hello"}], + max_tokens=16, + api_base="https://test.services.ai.azure.com/anthropic", + api_key="test-key", + headers=forwarded_headers, + extra_headers=deployment_headers, + provider_specific_header=ProviderSpecificHeader( + custom_llm_provider="azure_ai", extra_headers={"anthropic-beta": client_beta} + ), + client=AsyncHTTPHandler(transport=httpx.MockTransport(upstream.respond)), + ) + + assert upstream.anthropic_beta == expected + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_drops_top_level_and_nested_params(): """ diff --git a/tests/unit/llms/databricks/test_databricks_cost_calculator.py b/tests/unit/llms/databricks/test_databricks_cost_calculator.py index 7120a130462..34d3806d6eb 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: diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 20c36878fef..5f2af34e513 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -512,6 +512,40 @@ def test_map_openai_image_params_to_gemini_forwards_a_json_object_image_config() assert _map_image_config('{"aspectRatio": "1:1"}') == {"imageConfig": {"aspectRatio": "1:1"}} +@pytest.mark.parametrize("prefix", ("", "gemini/", "vertex_ai/")) +@pytest.mark.parametrize( + "model, size, expected", + ( + ("gemini-3-pro-image", "256x256", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3-pro-image", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3-pro-image-preview", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3-pro-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "2K"}), + ("gemini-3-pro-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "4K"}), + ("gemini-3.1-flash-image", "512x512", {"aspectRatio": "1:1", "imageSize": "512"}), + ("gemini-3.1-flash-image-preview", "512x512", {"aspectRatio": "1:1", "imageSize": "512"}), + ("gemini-3.1-flash-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "2K"}), + ("gemini-3.1-flash-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "4K"}), + ("gemini-3.1-flash-lite-image", "256x256", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3.1-flash-lite-image", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3.1-flash-lite-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-3.1-flash-lite-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "1K"}), + ("gemini-2.5-flash-image", "512x512", {"aspectRatio": "1:1"}), + ("gemini-2.5-flash-image", "1024x256", {"aspectRatio": "21:9"}), + ("gemini-2.5-flash-image", "256x1024", {"aspectRatio": "9:16"}), + ("gemini-2.5-flash-image", "8000x1000", {"aspectRatio": "21:9"}), + ("gemini-2.5-flash-image", "1000x8000", {"aspectRatio": "9:16"}), + ("gemini-3.1-flash-image", "1024x256", {"aspectRatio": "4:1", "imageSize": "512"}), + ), +) +def test_size_mapping_respects_model_image_limits( + prefix: str, model: str, size: str, expected: dict[str, str] +) -> None: + """Limits as of 2026-10-26, https://ai.google.dev/gemini-api/docs/image-generation#aspect_ratios_and_image_size.""" + assert map_openai_image_params_to_gemini({"size": size}, prefix + model, ["size"]) == { + "imageConfig": expected + } + + @pytest.mark.parametrize("raw_image_config", ["[1, 2]", '"16:9"', "null"]) def test_map_openai_image_params_to_gemini_drops_or_refuses_a_non_object_image_config(raw_image_config: str): with contextlib.suppress(litellm.UnsupportedParamsError): diff --git a/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py index 6fd2b006fef..00a3cbfdb7f 100644 --- a/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py +++ b/tests/unit/llms/huggingface/rerank/test_huggingface_rerank_transformation.py @@ -5,11 +5,14 @@ Based on the test patterns from other rerank providers and the current HuggingFa import asyncio import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler def assert_response_shape(response, custom_llm_provider): @@ -219,11 +222,11 @@ def test_huggingface_rerank_return_documents(mock_post): assert "text" in result["document"] -def test_huggingface_rerank_config(): +def test_huggingface_rerank_config() -> None: """Test HuggingFaceRerankConfig class functionality.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig - config = HuggingFaceRerankConfig() + config: Final = HuggingFaceRerankConfig() # Test complete URL generation assert config.get_complete_url(None, "test") == "https://api-inference.huggingface.co/rerank" @@ -233,14 +236,16 @@ def test_huggingface_rerank_config(): assert custom_url == "https://custom.huggingface.co/rerank" # Test supported parameters - supported_params = config.get_supported_cohere_rerank_params("test") + supported_params: Final = config.get_supported_cohere_rerank_params("test") assert "query" in supported_params assert "documents" in supported_params assert "top_n" in supported_params assert "return_documents" in supported_params + assert "truncate" in supported_params + assert "truncation_direction" in supported_params # Test parameter mapping - params = config.map_cohere_rerank_params( + params: Final = config.map_cohere_rerank_params( non_default_params={ "query": "hello", "documents": ["hello", "world"], @@ -259,24 +264,28 @@ def test_huggingface_rerank_config(): assert params["return_text"] is True -def test_request_transformation(): +def test_request_transformation() -> None: """Test request transformation logic.""" from litellm.llms.huggingface.rerank.transformation import HuggingFaceRerankConfig from litellm.types.rerank import OptionalRerankParams - config = HuggingFaceRerankConfig() + config: Final = HuggingFaceRerankConfig() - optional_params = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) + optional_params: Final = OptionalRerankParams(query="hello", texts=["hello", "world"], top_n=2, return_text=True) - request_body = config.transform_rerank_request(model="test", optional_rerank_params=optional_params, headers={}) + request_body: Final = config.transform_rerank_request( + model="test", optional_rerank_params=optional_params, headers={} + ) - assert request_body["query"] == "hello" - assert request_body["texts"] == ["hello", "world"] - assert request_body["top_n"] == 2 - assert request_body["return_text"] is True - assert request_body["raw_scores"] is False - assert request_body["truncate"] is False - assert request_body["truncation_direction"] == "Right" + expected_request: Final = { + "query": "hello", + "texts": ["hello", "world"], + "top_n": 2, + "return_text": True, + "raw_scores": False, + "truncation_direction": "Right", + } + assert request_body == expected_request def test_response_transformation(): @@ -345,20 +354,20 @@ def test_validate_environment(): @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_huggingface_rerank_request_payload(mock_post): +def test_huggingface_rerank_request_payload(mock_post: MagicMock) -> None: """Test that the request payload is correctly formatted for HuggingFace API.""" - mock_response_data = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}] + mock_response_data: Final = [{"index": 0, "score": 0.9}, {"index": 1, "score": 0.1}] def return_val(): return mock_response_data - mock_response = MagicMock() + mock_response: Final = MagicMock() mock_response.json = return_val mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} mock_post.return_value = mock_response - response = litellm.rerank( + response: Final = litellm.rerank( model="huggingface/BAAI/bge-reranker-base", query="hello", documents=["hello", "world"], @@ -370,28 +379,58 @@ def test_huggingface_rerank_request_payload(mock_post): mock_post.assert_called_once() # Verify URL - call_url = mock_post.call_args.kwargs["url"] + call_url: Final = mock_post.call_args.kwargs["url"] assert call_url == "https://api-inference.huggingface.co/rerank" # Verify headers - headers = mock_post.call_args.kwargs["headers"] + headers: Final = mock_post.call_args.kwargs["headers"] assert "Bearer test_api_key" in headers["Authorization"] assert headers["content-type"] == "application/json" # Verify request body - request_data = json.loads(mock_post.call_args.kwargs["data"]) - expected_request = { + request_data: Final = json.loads(mock_post.call_args.kwargs["data"]) + expected_request: Final = { "query": "hello", "texts": ["hello", "world"], "raw_scores": False, "return_text": True, - "truncate": False, "truncation_direction": "Right", "top_n": 2, } - for key, value in expected_request.items(): - assert request_data[key] == value + assert request_data == expected_request assert response.results is not None assert len(response.results) == 2 + + +def test_huggingface_rerank_forwards_truncate_options() -> None: + expected_request: Final = { + "query": "hello", + "texts": ["hello", "world"], + "raw_scores": False, + "return_text": False, + "truncation_direction": "Left", + "truncate": True, + } + + def respond(request: httpx.Request) -> httpx.Response: + assert request.url == "https://api-inference.huggingface.co/rerank" + assert request.headers["Authorization"] == "Bearer test_api_key" + assert json.loads(request.content) == expected_request + return httpx.Response(200, json=[{"index": 0, "score": 0.9}]) + + with httpx.Client(transport=httpx.MockTransport(respond)) as http_client: + response: Final = litellm.rerank( + model="huggingface/BAAI/bge-reranker-base", + query="hello", + documents=["hello", "world"], + api_key="test_api_key", + return_documents=False, + truncate=True, + truncation_direction="Left", + client=HTTPHandler(client=http_client), + ) + + assert response.results is not None + assert response.results[0]["relevance_score"] == 0.9 diff --git a/tests/unit/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py index 10f57be33fe..b76e1707805 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -1,4 +1,4 @@ -from typing import List, cast +from typing import Final, List, cast from unittest.mock import MagicMock, patch import pytest @@ -8,8 +8,10 @@ from litellm.llms.mistral.chat.transformation import ( MistralChatResponseIterator, MistralConfig, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.mistral.transformation import VertexAIMistralConfig from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse +from litellm.utils import supports_reasoning @pytest.mark.asyncio @@ -805,7 +807,7 @@ class TestMistralStripsOutputOnlyFields: Regression for https://github.com/BerriAI/litellm/issues/30835. """ - def test_assistant_reasoning_content_is_dropped(self): + def test_assistant_reasoning_content_is_sent_as_a_thinking_chunk(self): messages = cast( List[AllMessageValues], [ @@ -824,11 +826,10 @@ class TestMistralStripsOutputOnlyFields: MistralConfig().transform_messages(messages=messages, model="mistral-medium-3-5"), ) - assistant_message = result[-1] - assert "reasoning_content" not in assistant_message - assert "thinking_blocks" not in assistant_message - assert assistant_message["content"] == "Follow-up" - assert assistant_message["role"] == "assistant" + assert result[-1] == { + "role": "assistant", + "content": [_think_chunk("Some internal reasoning text."), {"type": "text", "text": "Follow-up"}], + } def test_tool_search_reply_drops_litellm_response_fields(self): tool_search_reply = { @@ -920,6 +921,508 @@ class TestMistralStripsOutputOnlyFields: assert "reasoning_content" not in result[-1] +def _think_chunk(text: str) -> dict[str, object]: + return {"type": "thinking", "thinking": [{"type": "text", "text": text}], "closed": True} + + +def _mistral_request_messages( + messages: list[dict[str, object]], config: MistralConfig | None = None, model: str = "mistral-medium-3-5" +) -> list[object]: + request: Final = (config or MistralConfig()).transform_request( + model=model, + messages=cast(List[AllMessageValues], messages), + optional_params={}, + litellm_params={}, + headers={}, + ) + return request["messages"] + + +@pytest.mark.parametrize("content", [None, ""], ids=["null", "empty"]) +def test_mistral_sends_reasoning_of_a_tool_call_turn_as_a_thinking_chunk(content: str | None): + tool_call: Final = {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": content, + "reasoning_content": "I should call get_weather.", + "tool_calls": [tool_call], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [_think_chunk("I should call get_weather.")], + "tool_calls": [tool_call], + } + assert result[2] == {"role": "tool", "tool_call_id": "call_1", "content": "sunny"} + + +def test_mistral_sends_reasoning_ahead_of_a_text_part_list(): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Question?"}, + { + "role": "assistant", + "content": [{"type": "text", "text": "First. "}, {"type": "text", "text": "Second."}], + "reasoning_content": "Two parts.", + }, + {"role": "user", "content": "And then?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [_think_chunk("Two parts."), {"type": "text", "text": "First. Second."}], + } + + +def test_mistral_drops_a_reasoning_only_assistant_turn_between_a_tool_call_and_its_reply(): + tool_call: Final = {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "read the log"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + {"role": "assistant", "content": None, "reasoning_content": "The log is empty."}, + {"role": "tool", "tool_call_id": "call_1", "content": ""}, + ] + ) + + assert result == [ + {"role": "user", "content": "read the log"}, + {"role": "assistant", "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": "call_1", "content": ""}, + ] + + +@pytest.mark.parametrize( + "assistant_content, reasoning_fields, expected_content", + [ + pytest.param( + [{"type": "thinking", "thinking": "R", "signature": "c2ln"}, {"type": "text", "text": "A"}], + {"reasoning_content": "R"}, + [_think_chunk("R"), {"type": "text", "text": "A"}], + id="anthropic_thinking_part_with_reasoning_content", + ), + pytest.param( + [{"type": "thinking", "thinking": "R", "signature": "c2ln"}, {"type": "text", "text": "A"}], + {}, + "A", + id="anthropic_thinking_part_alone", + ), + pytest.param( + {"type": "text", "text": "A"}, + {"reasoning_content": "R"}, + {"type": "text", "text": "A"}, + id="content_that_is_not_a_list", + ), + ], +) +def test_mistral_never_forwards_a_thinking_part_mistral_cannot_read( + assistant_content: object, reasoning_fields: dict[str, object], expected_content: object +): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": assistant_content, **reasoning_fields}, + {"role": "user", "content": "Bye"}, + ] + ) + + assert result[1] == {"role": "assistant", "content": expected_content} + + +@pytest.mark.parametrize( + "reasoning_fields, expected_reasoning", + [ + pytest.param( + { + "thinking_blocks": [ + {"type": "thinking", "thinking": "Let", "signature": "mistral"}, + {"type": "thinking", "thinking": " me", "signature": "mistral"}, + {"type": "thinking", "thinking": " think", "signature": "mistral"}, + ] + }, + "Let me think", + id="mistral_stream_fragments", + ), + pytest.param( + { + "thinking_blocks": [ + {"type": "thinking", "thinking": "Step one.", "signature": "c2lnLTE="}, + {"type": "redacted_thinking", "data": "ZW5jcnlwdGVk"}, + {"type": "thinking", "thinking": " Step two.", "signature": "c2lnLTI="}, + ] + }, + "Step one. Step two.", + id="anthropic_blocks_with_redacted", + ), + pytest.param( + {"reasoning_content": " ", "thinking_blocks": [{"type": "thinking", "thinking": "From blocks."}]}, + "From blocks.", + id="blank_reasoning_content_falls_back_to_blocks", + ), + ], +) +def test_mistral_sends_replayed_thinking_blocks_as_one_thinking_chunk( + reasoning_fields: dict[str, object], expected_reasoning: str +): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Question?"}, + {"role": "assistant", "content": "Answer.", **reasoning_fields}, + {"role": "user", "content": "And then?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [_think_chunk(expected_reasoning), {"type": "text", "text": "Answer."}], + } + + +@pytest.mark.parametrize( + "native_content, reasoning_fields", + [ + pytest.param( + [ + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}]}, + {"type": "text", "text": "I have picked it."}, + ], + {}, + id="native_chunk", + ), + pytest.param( + [ + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}]}, + {"type": "text", "text": "I have picked it."}, + ], + {"reasoning_content": "I pick 518362.", "thinking_blocks": [{"type": "thinking", "thinking": "I pick"}]}, + id="native_chunk_and_reasoning_content", + ), + pytest.param( + [ + {"type": "text", "text": "I have "}, + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}]}, + {"type": "text", "text": "picked it."}, + ], + {}, + id="native_chunk_between_text_parts", + ), + ], +) +def test_mistral_keeps_a_native_thinking_chunk_once( + native_content: list[dict[str, object]], reasoning_fields: dict[str, object] +): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Pick a number."}, + {"role": "assistant", "content": native_content, **reasoning_fields}, + {"role": "user", "content": "Which one?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}]}, + {"type": "text", "text": "I have picked it."}, + ], + } + + +def test_mistral_keeps_a_native_thinking_chunk_with_no_text_part(): + native_chunk: Final = {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}], "closed": True} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Pick a number."}, + {"role": "assistant", "content": [native_chunk]}, + {"role": "user", "content": "Which one?"}, + ] + ) + + assert result[1] == {"role": "assistant", "content": [native_chunk]} + + +@pytest.mark.parametrize( + ("inner_items", "expected_content"), + [ + pytest.param( + [{"type": "text", "text": "I pick 518362."}, {"type": "reference", "reference_ids": [1]}], + [ + {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 518362."}], "closed": True}, + {"type": "text", "text": "I have picked it."}, + ], + id="text_and_reference", + ), + pytest.param( + [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}], + "I have picked it.", + id="only_image_url", + ), + pytest.param([1, 2], "I have picked it.", id="not_items"), + ], +) +def test_mistral_sends_only_the_text_items_of_a_native_thinking_chunk( + inner_items: list[object], expected_content: object +): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Pick a number."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": inner_items, "closed": True}, + {"type": "text", "text": "I have picked it."}, + ], + }, + {"role": "user", "content": "Which one?"}, + ] + ) + + assert result[1] == {"role": "assistant", "content": expected_content} + + +def test_mistral_sends_only_the_text_items_of_a_native_thinking_chunk_beside_an_image(): + image_part: Final = {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": [{"type": "text", "text": "What is this?"}, image_part]}, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": [{"type": "text", "text": "It has whiskers."}, {"type": "reference", "reference_ids": [1]}], + }, + {"type": "text", "text": "A cat."}, + ], + }, + {"role": "user", "content": "What color?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": [{"type": "text", "text": "It has whiskers."}]}, + {"type": "text", "text": "A cat."}, + ], + } + + +_OPEN_THINK_CHUNK: Final = { + "type": "thinking", + "thinking": [{"type": "text", "text": "I pick 482913."}], + "closed": False, +} +_THINK_CHUNK_WITHOUT_CLOSED: Final = {"type": "thinking", "thinking": [{"type": "text", "text": "I pick 482913."}]} +_PICKED_TEXT: Final = {"type": "text", "text": "I have picked it."} + + +@pytest.mark.parametrize( + ("assistant_turn", "later_messages", "expected_turn"), + [ + pytest.param( + {"role": "assistant", "content": [_OPEN_THINK_CHUNK, _PICKED_TEXT]}, + [{"role": "user", "content": "Which one?"}], + {"role": "assistant", "content": [_THINK_CHUNK_WITHOUT_CLOSED, _PICKED_TEXT]}, + id="history_turn", + ), + pytest.param( + {"role": "assistant", "content": [_OPEN_THINK_CHUNK], "prefix": True}, + [], + {"role": "assistant", "content": [_OPEN_THINK_CHUNK], "prefix": True}, + id="last_chunk_of_final_prefix_turn", + ), + pytest.param( + {"role": "assistant", "content": [_OPEN_THINK_CHUNK, _PICKED_TEXT], "prefix": True}, + [], + {"role": "assistant", "content": [_THINK_CHUNK_WITHOUT_CLOSED, _PICKED_TEXT], "prefix": True}, + id="final_prefix_turn_with_text_after_the_chunk", + ), + pytest.param( + {"role": "assistant", "content": [_PICKED_TEXT, _OPEN_THINK_CHUNK], "prefix": True}, + [], + {"role": "assistant", "content": [_THINK_CHUNK_WITHOUT_CLOSED, _PICKED_TEXT], "prefix": True}, + id="final_prefix_turn_with_text_before_the_chunk", + ), + pytest.param( + {"role": "assistant", "content": [_OPEN_THINK_CHUNK]}, + [], + {"role": "assistant", "content": [_THINK_CHUNK_WITHOUT_CLOSED]}, + id="final_turn_without_prefix", + ), + ], +) +def test_mistral_keeps_an_open_thinking_chunk_only_as_the_last_chunk_of_a_final_prefix_turn( + assistant_turn: dict[str, object], later_messages: list[dict[str, object]], expected_turn: dict[str, object] +): + result: Final = _mistral_request_messages( + [{"role": "user", "content": "Pick a number."}, assistant_turn, *later_messages] + ) + + assert result[1] == expected_turn + + +@pytest.mark.parametrize( + "reasoning_fields", + [ + pytest.param({}, id="none"), + pytest.param({"reasoning_content": None, "thinking_blocks": None}, id="null"), + pytest.param({"reasoning_content": ""}, id="empty_reasoning_content"), + pytest.param({"reasoning_content": " \n "}, id="blank_reasoning_content"), + pytest.param({"thinking_blocks": []}, id="empty_thinking_blocks"), + pytest.param({"thinking_blocks": [{"type": "redacted_thinking", "data": "ZW5j"}]}, id="only_redacted"), + pytest.param({"thinking_blocks": [{"type": "thinking", "thinking": " "}]}, id="blank_thinking_block"), + pytest.param({"thinking_blocks": "not a list"}, id="thinking_blocks_not_a_list"), + ], +) +def test_mistral_sends_an_assistant_turn_without_usable_reasoning_as_plain_text(reasoning_fields: dict[str, object]): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "Question?"}, + {"role": "assistant", "content": "Answer.", **reasoning_fields}, + {"role": "user", "content": "And then?"}, + ] + ) + + assert result[1] == {"role": "assistant", "content": "Answer."} + + +def test_mistral_still_flattens_text_part_lists_to_a_string(): + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": [{"type": "text", "text": "Hello "}, {"type": "text", "text": "there"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Hi "}, {"type": "text", "text": "back"}]}, + {"role": "user", "content": "Bye"}, + ] + ) + + assert result == [ + {"role": "user", "content": "Hello there"}, + {"role": "assistant", "content": "Hi back"}, + {"role": "user", "content": "Bye"}, + ] + + +def test_mistral_sends_reasoning_as_a_thinking_chunk_when_the_request_carries_an_image(): + data_uri: Final = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg==" + + result: Final = _mistral_request_messages( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": data_uri}}, + ], + }, + {"role": "assistant", "content": "A cat.", "reasoning_content": "It has whiskers."}, + {"role": "user", "content": "What color?"}, + ] + ) + + assert result[1] == { + "role": "assistant", + "content": [_think_chunk("It has whiskers."), {"type": "text", "text": "A cat."}], + } + + +@pytest.mark.parametrize( + "config", + [pytest.param(MistralConfig(), id="mistral"), pytest.param(VertexAIMistralConfig(), id="vertex_ai")], +) +@pytest.mark.parametrize( + "assistant_message", + [ + pytest.param( + {"role": "assistant", "content": "I have picked it.", "reasoning_content": "I pick 518362."}, + id="reasoning_content", + ), + pytest.param( + { + "role": "assistant", + "content": "I have picked it.", + "thinking_blocks": [{"type": "thinking", "thinking": "I pick 518362.", "signature": "mistral"}], + }, + id="thinking_blocks", + ), + pytest.param( + { + "role": "assistant", + "content": [_think_chunk("I pick 518362."), {"type": "text", "text": "I have picked it."}], + }, + id="native_chunk", + ), + ], +) +def test_mistral_drops_replayed_reasoning_for_a_model_without_reasoning( + config: MistralConfig, assistant_message: dict[str, object] +): + model: Final = "mistral-chat-model-without-reasoning" + assert not supports_reasoning(model=model, custom_llm_provider=config.custom_llm_provider) + + result: Final = _mistral_request_messages( + [{"role": "user", "content": "Pick a number."}, assistant_message, {"role": "user", "content": "Which one?"}], + config=config, + model=model, + ) + + assert result[1] == {"role": "assistant", "content": "I have picked it."} + + +@pytest.mark.parametrize( + "config", + [pytest.param(MistralConfig(), id="mistral"), pytest.param(VertexAIMistralConfig(), id="vertex_ai")], +) +def test_mistral_drops_a_native_chunk_without_text_for_a_model_without_reasoning(config: MistralConfig): + tool_call: Final = {"id": "call_1", "type": "function", "function": {"name": "get_time", "arguments": "{}"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": "What time is it?"}, + {"role": "assistant", "content": [_think_chunk("I should call get_time.")], "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": "call_1", "content": "12:00"}, + ], + config=config, + model="mistral-chat-model-without-reasoning", + ) + + assert result[1] == {"role": "assistant", "tool_calls": [tool_call]} + + +@pytest.mark.parametrize( + "config", + [pytest.param(MistralConfig(), id="mistral"), pytest.param(VertexAIMistralConfig(), id="vertex_ai")], +) +def test_mistral_drops_a_native_chunk_from_an_image_history_for_a_model_without_reasoning(config: MistralConfig): + image_part: Final = {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}} + + result: Final = _mistral_request_messages( + [ + {"role": "user", "content": [{"type": "text", "text": "What is this?"}, image_part]}, + { + "role": "assistant", + "content": [_think_chunk("It has whiskers."), {"type": "text", "text": "A cat."}], + }, + {"role": "user", "content": "What color?"}, + ], + config=config, + model="mistral-chat-model-without-reasoning", + ) + + assert result[1] == {"role": "assistant", "content": [{"type": "text", "text": "A cat."}]} + + def test_mistral_transform_request_keeps_an_empty_tool_message_after_its_tool_call(): messages: List[AllMessageValues] = cast( List[AllMessageValues], diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 1a0bbe171a9..df9f5fd9d14 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1064,6 +1064,9 @@ def test_ollama_embeddings(mock_embeddings): logging_obj=mock.ANY, model_response=mock.ANY, encoding=mock.ANY, + api_key=None, + headers={}, + timeout=600, ) except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -1088,6 +1091,9 @@ def test_ollama_aembeddings(mock_aembeddings): logging_obj=mock.ANY, model_response=mock.ANY, encoding=mock.ANY, + api_key=None, + headers={}, + timeout=600, ) except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/unit/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py index 6a1dfebc236..97f6f8468f4 100644 --- a/tests/unit/llms/ollama/test_ollama_embedding.py +++ b/tests/unit/llms/ollama/test_ollama_embedding.py @@ -1,3 +1,6 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -127,12 +130,183 @@ def test_prompt_eval_fallback_when_missing(mock_embedding_response, mock_encodin ], ) @respx.mock -def test_ollama_embeddings_ignore_a_trailing_slash_on_api_base(api_base: str) -> None: +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_ollama_embeddings_ignore_a_trailing_slash_on_api_base( + api_base: str, asynchronous: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") route: Final = respx.post("http://gateway.example.com/ollama/api/embed").mock( return_value=httpx.Response(200, json={"embeddings": [[0.1, 0.2], [0.3, 0.4]], "prompt_eval_count": 2}) ) - response: Final = litellm.embedding(model="ollama/all-minilm", api_base=api_base, input=["hello", "world"]) + response: Final = ( + await litellm.aembedding(model="ollama/all-minilm", api_base=api_base, input=["hello", "world"]) + if asynchronous + else litellm.embedding(model="ollama/all-minilm", api_base=api_base, input=["hello", "world"]) + ) - assert route.called + assert route.call_count == 1 assert [item["embedding"] for item in response.data] == [[0.1, 0.2], [0.3, 0.4]] + + +@pytest.fixture +def embedding_route(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.Route: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + return respx_mock.post("http://ollama.test/api/embed").respond( + 200, json={"embeddings": [[0.1, 0.2]], "prompt_eval_count": 2} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + ("api_key", "env_key", "authorization"), + [ + ("explicit-key", None, "Bearer explicit-key"), + (None, "env-key", "Bearer env-key"), + ("explicit-key", "env-key", "Bearer explicit-key"), + (None, None, None), + ], +) +async def test_sdk_embedding_forwards_resolved_auth_and_merged_headers( + embedding_route: respx.Route, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + api_key: str | None, + env_key: str | None, + authorization: str | None, +) -> None: + if env_key is None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + else: + monkeypatch.setenv("OLLAMA_API_KEY", env_key) + response: Final = ( + await litellm.aembedding( + model="ollama/nomic-embed-text", + input="hello", + api_base="http://ollama.test", + api_key=api_key, + headers={"X-Caller": "caller", "X-Shared": "original"}, + extra_headers={"X-Extra": "extra", "X-Shared": "override"}, + ) + if asynchronous + else litellm.embedding( + model="ollama/nomic-embed-text", + input="hello", + api_base="http://ollama.test", + api_key=api_key, + headers={"X-Caller": "caller", "X-Shared": "original"}, + extra_headers={"X-Extra": "extra", "X-Shared": "override"}, + ) + ) + request: Final = embedding_route.calls.last.request + + assert embedding_route.call_count == 1 + assert request.headers.get("Authorization") == authorization + assert (request.headers["X-Caller"], request.headers["X-Extra"], request.headers["X-Shared"]) == ( + "caller", + "extra", + "override", + ) + assert request.extensions["timeout"] == httpx.Timeout(600).as_dict() + assert json.loads(request.content) == {"model": "nomic-embed-text", "input": ["hello"]} + assert isinstance(response, EmbeddingResponse) + assert response.data[0]["embedding"] == [0.1, 0.2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("timeout", [0, 2.5, httpx.Timeout(connect=1, read=2, write=3, pool=4)]) +async def test_sdk_embedding_applies_timeout_to_http_request( + embedding_route: respx.Route, asynchronous: bool, timeout: float | httpx.Timeout +) -> None: + response: Final = ( + await litellm.aembedding( + model="ollama/nomic-embed-text", input=["hello"], api_base="http://ollama.test", timeout=timeout + ) + if asynchronous + else litellm.embedding( + model="ollama/nomic-embed-text", input=["hello"], api_base="http://ollama.test", timeout=timeout + ) + ) + + assert embedding_route.call_count == 1 + assert embedding_route.calls.last.request.extensions["timeout"] == httpx.Timeout(timeout).as_dict() + assert isinstance(response, EmbeddingResponse) + assert response.data[0]["embedding"] == [0.1, 0.2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + ("api_key", "headers", "authorization"), + [ + ("explicit-key", {"X-Custom": "custom"}, "Bearer explicit-key"), + ("explicit-key", {"Authorization": "Basic caller"}, "Basic caller"), + ("explicit-key", {"authorization": "Bearer caller"}, "Bearer caller"), + (None, {}, None), + ("", {}, None), + ], +) +@pytest.mark.parametrize("timeout_params", [{}, {"timeout": None}], ids=["omitted", "none"]) +async def test_handler_embedding_preserves_caller_headers_and_client_timeout( + embedding_route: respx.Route, + asynchronous: bool, + api_key: str | None, + headers: Mapping[str, str], + authorization: str | None, + timeout_params: Mapping[str, None], +) -> None: + caller_headers: Final = MappingProxyType(dict(headers)) + expected_timeout: Final = ( + litellm.module_level_aclient.client.timeout if asynchronous else litellm.module_level_client.client.timeout + ) + response: Final = ( + await ollama_aembeddings( + "http://ollama.test", + "nomic-embed-text", + ["hello"], + EmbeddingResponse(), + {}, + None, + None, + api_key=api_key, + headers=caller_headers, + **timeout_params, + ) + if asynchronous + else ollama_embeddings( + "http://ollama.test", + "nomic-embed-text", + ["hello"], + {}, + EmbeddingResponse(), + None, + api_key=api_key, + headers=caller_headers, + **timeout_params, + ) + ) + request: Final = embedding_route.calls.last.request + + assert embedding_route.call_count == 1 + assert request.headers.get_list("Authorization") == ([] if authorization is None else [authorization]) + assert dict(caller_headers) == headers + assert all(request.headers[name] == value for name, value in caller_headers.items()) + assert request.extensions["timeout"] == expected_timeout.as_dict() + assert json.loads(request.content) == {"model": "nomic-embed-text", "input": ["hello"]} + assert response.data[0]["embedding"] == [0.1, 0.2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_sdk_embedding_rejects_non_string_batch(embedding_route: respx.Route, asynchronous: bool) -> None: + if asynchronous: + with pytest.raises(litellm.BadRequestError, match="Invalid input for ollama embeddings"): + await litellm.aembedding(model="ollama/nomic-embed-text", input=[123], api_base="http://ollama.test") + else: + with pytest.raises(litellm.BadRequestError, match="Invalid input for ollama embeddings"): + litellm.embedding(model="ollama/nomic-embed-text", input=[123], api_base="http://ollama.test") + + assert not embedding_route.called diff --git a/tests/unit/llms/openai_like/test_zerogpu_provider.py b/tests/unit/llms/openai_like/test_zerogpu_provider.py new file mode 100644 index 00000000000..1021237c335 --- /dev/null +++ b/tests/unit/llms/openai_like/test_zerogpu_provider.py @@ -0,0 +1,239 @@ +import json +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +import respx +from openai import OpenAI + +import litellm +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.proxy.auth.model_checks import get_provider_models + +ZEROGPU_API_BASE: Final = "https://api.zerogpu.ai/v1" + +ZEROGPU_MODELS: Final = ( + "deepseek-v4.1-flash", + "glm-5.3-flash", + "gpt-4.1-mini", + "gpt-5.4-nano", + "gpt-5.6-luna", + "gpt-6-luna", + "gpt-oss-120b", + "LFM2.5-1.2B-Instruct", + "LFM2.5-1.2B-Thinking", + "llama-3.1-8b-instruct-fast", + "llama-3.2-3b-instruct-fast", + "qwen3-30b-a3b-fp8", + "qwen3.6-35b-a3b", +) + + +WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, +} + + +class _RecordingZeroGPU: + def __init__(self, usage: Mapping[str, object]) -> None: + self.usage: Final = usage + self.requests: tuple[httpx.Request, ...] = () + + def respond(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-zerogpu", + "object": "chat.completion", + "created": 1234567890, + "model": json.loads(request.content)["model"], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": dict(self.usage), + }, + ) + + def client(self) -> OpenAI: + return OpenAI( + api_key="zerogpu-test-key", + base_url=ZEROGPU_API_BASE, + http_client=httpx.Client(transport=httpx.MockTransport(self.respond)), + ) + + +def test_prefixed_model_reaches_zerogpu_chat_completions_without_the_prefix(): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + litellm.completion( + model="zerogpu/gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + assert len(zerogpu.requests) == 1 + assert str(zerogpu.requests[0].url) == f"{ZEROGPU_API_BASE}/chat/completions" + assert json.loads(zerogpu.requests[0].content) == { + "model": "gpt-oss-120b", + "messages": [{"role": "user", "content": "hello"}], + } + + +def test_responses_request_reaches_zerogpu_responses_natively(): + with respx.mock() as upstream: + route: Final = upstream.post(f"{ZEROGPU_API_BASE}/responses").respond( + 200, + json={ + "id": "resp_zerogpu", + "object": "response", + "created_at": 1_792_000_000, + "model": "gpt-oss-120b", + "status": "completed", + "output": [ + { + "id": "msg_zerogpu", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + }, + ) + response: Final = litellm.responses(model="zerogpu/gpt-oss-120b", input="hello", api_key="zerogpu-test-key") + + body: Final = json.loads(route.calls.last.request.content) + assert body["model"] == "gpt-oss-120b" + assert body["input"] == "hello" + assert response.output[0].content[0].text == "hi" + + +def test_anthropic_messages_request_is_bridged_to_zerogpu_chat_completions(): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + response: Final = litellm.anthropic.messages.create( + model="zerogpu/gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + client=client, + ) + + assert [str(request.url) for request in zerogpu.requests] == [f"{ZEROGPU_API_BASE}/chat/completions"] + assert json.loads(zerogpu.requests[0].content) == { + "model": "gpt-oss-120b", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + } + assert response["content"] == [{"type": "text", "text": "hi"}] + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_zerogpu_completion_is_charged_at_the_cost_map_rates(model: str): + usage: Final = { + "prompt_tokens": 1000, + "completion_tokens": 200, + "total_tokens": 1200, + "prompt_tokens_details": {"cached_tokens": 400}, + } + zerogpu: Final = _RecordingZeroGPU(usage) + + with zerogpu.client() as client: + response: Final = litellm.completion( + model=f"zerogpu/{model}", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + rates: Final = litellm.model_cost[f"zerogpu/{model}"] + cached_rate: Final = rates.get("cache_read_input_token_cost", rates["input_cost_per_token"]) + expected_cost: Final = ( + 600 * rates["input_cost_per_token"] + 400 * cached_rate + 200 * rates["output_cost_per_token"] + ) + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + assert expected_cost > 0 + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_trim_messages_fits_zerogpu_history_into_the_model_context_window(model: str): + budget: Final = 20 + trim_ratio: Final = budget / litellm.model_cost[f"zerogpu/{model}"]["max_input_tokens"] + messages: Final = [{"role": "user", "content": "lorem ipsum " * 300}] + + trimmed: Final = litellm.utils.trim_messages(messages, model=f"zerogpu/{model}", trim_ratio=trim_ratio) + + trimmed_tokens: Final = litellm.token_counter(model=f"zerogpu/{model}", messages=trimmed) + assert trimmed_tokens <= budget, trimmed + + +@pytest.mark.parametrize( + ("model", "forwards_tools", "forwards_reasoning_effort"), + [ + ("gpt-5.6-luna", True, True), + ("gpt-oss-120b", True, True), + ("qwen3.6-35b-a3b", True, True), + ("LFM2.5-1.2B-Instruct", False, False), + ], +) +def test_tools_and_reasoning_effort_reach_zerogpu_only_for_models_that_support_them( + model: str, forwards_tools: bool, forwards_reasoning_effort: bool +): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + litellm.completion( + model=f"zerogpu/{model}", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[WEATHER_TOOL], + tool_choice="auto", + reasoning_effort="high", + drop_params=True, + client=client, + ) + + tools: Final = {"tools": [WEATHER_TOOL], "tool_choice": "auto"} if forwards_tools else {} + reasoning: Final = {"reasoning_effort": "high"} if forwards_reasoning_effort else {} + assert json.loads(zerogpu.requests[0].content) == { + "model": model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + **tools, + **reasoning, + } + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_model_info_reports_the_reasoning_support_litellm_routes_with(model: str): + model_info: Final = litellm.get_model_info(f"zerogpu/{model}") + + assert bool(model_info.get("supports_reasoning")) == litellm.utils.supports_reasoning(f"zerogpu/{model}") + + +def test_zerogpu_key_and_base_come_from_zerogpu_env_vars(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ZEROGPU_API_KEY", "zerogpu-env-key") + monkeypatch.delenv("ZEROGPU_API_BASE", raising=False) + + assert get_llm_provider(model="zerogpu/gpt-oss-120b") == ( + "gpt-oss-120b", + "zerogpu", + "zerogpu-env-key", + ZEROGPU_API_BASE, + ) + + monkeypatch.setenv("ZEROGPU_API_BASE", "https://zerogpu.internal.example/v1") + + assert get_llm_provider(model="zerogpu/gpt-oss-120b")[3] == "https://zerogpu.internal.example/v1" + + +def test_zerogpu_wildcard_lists_every_zerogpu_catalog_model(): + catalog_models: Final = sorted( + key for key, row in litellm.model_cost.items() if row.get("litellm_provider") == "zerogpu" + ) + + assert catalog_models + assert sorted(get_provider_models("zerogpu") or []) == catalog_models diff --git a/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py b/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py index 534176e381a..b6d26aa1de6 100644 --- a/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py +++ b/tests/unit/llms/perplexity/responses/test_perplexity_responses_transformation.py @@ -8,10 +8,13 @@ Source: litellm/llms/perplexity/responses/transformation.py """ import json +from typing import Final import httpx import pytest +import respx +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -21,6 +24,17 @@ from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +_SUCCESS_RESPONSE: Final[dict[str, object]] = { + "id": "resp_moderation_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "openai/gpt-5.2", + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} + + class TestPerplexityResponsesTransformation: """Test Perplexity Responses API configuration and transformations""" @@ -60,9 +74,7 @@ class TestPerplexityResponsesTransformation: assert len(result["tools"]) == 1 assert result["tools"][0]["type"] == "function" assert result["tools"][0]["function"]["name"] == "get_weather" - assert ( - result["tools"][0]["function"]["description"] == "Get the current weather" - ) + assert result["tools"][0]["function"]["description"] == "Get the current weather" assert "parameters" in result["tools"][0]["function"] def test_web_search_tool_passthrough(self): @@ -130,9 +142,7 @@ class TestPerplexityResponsesTransformation: """tool_choice passes through""" config = PerplexityResponsesConfig() - params = ResponsesAPIOptionalRequestParams( - tool_choice="required", temperature=0.7 - ) + params = ResponsesAPIOptionalRequestParams(tool_choice="required", temperature=0.7) result = config.map_openai_params( response_api_optional_params=params, @@ -146,9 +156,7 @@ class TestPerplexityResponsesTransformation: """parallel_tool_calls passes through""" config = PerplexityResponsesConfig() - params = ResponsesAPIOptionalRequestParams( - parallel_tool_calls=True, temperature=0.7 - ) + params = ResponsesAPIOptionalRequestParams(parallel_tool_calls=True, temperature=0.7) result = config.map_openai_params( response_api_optional_params=params, @@ -492,9 +500,7 @@ class TestPerplexityResponsesTransformation: "id": "msg_123", "role": "assistant", "status": "completed", - "content": [ - {"type": "output_text", "text": "Hello!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], } ], "usage": { @@ -549,9 +555,7 @@ class TestPerplexityResponsesTransformation: "id": "msg_123", "role": "assistant", "status": "completed", - "content": [ - {"type": "output_text", "text": "Hello!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], } ], "usage": { @@ -587,3 +591,26 @@ class TestPerplexityResponsesTransformation: assert result.type == "response.completed" assert result.response.usage.cost == 0.0003 assert isinstance(result.response.usage.cost, float) + + +@pytest.mark.asyncio +async def test_aresponses_rejects_moderation_for_perplexity_without_drop_params( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "drop_params", False) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.perplexity.ai/v1/responses").mock( + return_value=httpx.Response(200, json=_SUCCESS_RESPONSE) + ) + + with pytest.raises(litellm.UnsupportedParamsError): + await litellm.aresponses( + model="perplexity/openai/gpt-5.2", + api_base="https://api.perplexity.ai", + api_key="fake-api-key", + input="hi", + moderation={"model": "omni-moderation-latest"}, + ) + + assert upstream.call_count == 0 diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 0f5c21d08c7..065d894e658 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -8289,11 +8289,12 @@ def test_gemini_image_generation_ignores_non_dict_image_config(): def test_gemini_image_generation_openai_size_maps_to_google_table( size: str, expected_aspect_ratio: str, expected_image_size: str ): + """Flash table as of 2026-10-26, https://ai.google.dev/gemini-api/docs/image-generation#aspect_ratios_and_image_size.""" from litellm.llms.gemini.common_utils import ( map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == { + assert map_openai_size_to_gemini_image_config(size, "gemini-3.1-flash-image") == { "aspectRatio": expected_aspect_ratio, "imageSize": expected_image_size, } @@ -8317,7 +8318,7 @@ def test_gemini_image_generation_openai_size_snaps_to_nearest_option( map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == { + assert map_openai_size_to_gemini_image_config(size, "gemini-3.1-flash-image") == { "aspectRatio": expected_aspect_ratio, "imageSize": expected_image_size, } diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index 756e6441f1f..6ae0bbd0d72 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -3,22 +3,118 @@ import json import os from datetime import datetime from io import BytesIO -from typing import Dict +from typing import Dict, Final from unittest.mock import MagicMock, patch import httpx import pytest -from pydantic import ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.gemini.common_utils import map_openai_size_to_gemini_image_config from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import ( VertexAIGeminiImageEditConfig, ) from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import ( VertexAIImagenImageEditConfig, ) +from litellm.types.router import GenericLiteLLMParams + + +@pytest.mark.parametrize("size", ("1536x1024", "1024x1536", "3072x2048", "256x256", "512x512", "4096x4096")) +@pytest.mark.parametrize("model", ("gemini-3.1-flash-image", "gemini-3-pro-image", "gemini-3.1-flash-lite-image")) +def test_vertex_gemini_edit_size_reaches_request(size: str, model: str) -> None: + config: Final = VertexAIGeminiImageEditConfig() + payload, _ = config.transform_image_edit_request( + model, + "mountains", + BytesIO(b"image"), + config.map_openai_params({"size": size}, model, False), + GenericLiteLLMParams(), + {}, + ) + request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload)) + expected: Final = map_openai_size_to_gemini_image_config(size, model) or {} + assert request["generationConfig"] == { + "response_modalities": ["IMAGE"], + **( + {"image_config": {"aspect_ratio": expected["aspectRatio"], "image_size": expected["imageSize"]}} + if expected + else {} + ), + } + + +@pytest.mark.parametrize( + "size, expected_ratio", + (("1536x1024", "4:3"), ("1024x1536", "3:4"), ("100000x1", "16:9")), +) +def test_vertex_imagen_edit_size_reaches_request(size: str, expected_ratio: str | None) -> None: + model: Final = "imagen-3.0-capability-001" + config: Final = VertexAIImagenImageEditConfig() + payload, _ = config.transform_image_edit_request( + model, + "mountains", + BytesIO(b"image"), + config.map_openai_params({"size": size, "n": 2}, model, False), + GenericLiteLLMParams(), + {}, + ) + request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload)) + assert request["parameters"] == { + "sampleCount": 2, + "editMode": "EDIT_MODE_INPAINT_INSERTION", + "editConfig": {"baseSteps": 50}, + "guidanceScale": 7.5, + "seed": None, + **({"aspectRatio": expected_ratio} if expected_ratio is not None else {}), + } + + +@pytest.mark.parametrize("size", ("auto", "unknown", "0x1024", "-1x1024", "1024x", "3:2")) +@pytest.mark.parametrize( + "config_type, model, config_key, expected", + ( + ( + VertexAIGeminiImageEditConfig, + "gemini-3.1-flash-image", + "generationConfig", + {"response_modalities": ["IMAGE"], "image_config": {"aspect_ratio": "1:1"}}, + ), + ( + VertexAIImagenImageEditConfig, + "imagen-3.0-capability-001", + "parameters", + { + "sampleCount": 1, + "editMode": "EDIT_MODE_INPAINT_INSERTION", + "editConfig": {"baseSteps": 50}, + "guidanceScale": 7.5, + "seed": None, + }, + ), + ), +) +def test_vertex_edit_preserves_legacy_size_fallback( + size: str, + config_type: type[VertexAIGeminiImageEditConfig] | type[VertexAIImagenImageEditConfig], + model: str, + config_key: str, + expected: dict[str, object], +) -> None: + config: Final = config_type() + payload, _ = config.transform_image_edit_request( + model, + "mountains", + BytesIO(b"image"), + config.map_openai_params({"size": size}, model, False), + GenericLiteLLMParams(), + {}, + ) + request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload)) + assert request[config_key] == expected class TestVertexAIGeminiImageEditTransformation: diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 5ac16283094..10ce24e969c 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -11,6 +11,7 @@ from pydantic import ValidationError import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.gemini.image_generation.transformation import GoogleImageGenConfig from litellm.llms.vertex_ai.image_generation import ( get_vertex_ai_image_generation_config, ) @@ -23,6 +24,98 @@ from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import from litellm.types.utils import ImageResponse +@pytest.mark.parametrize("size", ("1536x1024", "1024x1536", "3072x2048", "256x256", "512x512", "4096x4096")) +@pytest.mark.parametrize("model", ("gemini-3.1-flash-image", "gemini-3-pro-image", "gemini-3.1-flash-lite-image")) +def test_vertex_gemini_size_matches_direct_gemini_request(size: str, model: str) -> None: + vertex: Final = VertexAIGeminiImageGenerationConfig() + direct: Final = GoogleImageGenConfig() + vertex_request: Final = vertex.transform_image_generation_request( + model, "mountains", vertex.map_openai_params({"size": size}, {}, model, False), {}, {} + ) + direct_request: Final = direct.transform_image_generation_request( + model, "mountains", direct.map_openai_params({"size": size}, {}, model, False), {}, {} + ) + assert vertex_request["generationConfig"].get("imageConfig") == direct_request["generationConfig"].get( + "imageConfig" + ) + + +@pytest.mark.parametrize( + "size, expected_ratio", + (("1536x1024", "4:3"), ("1024x1536", "3:4"), ("100000x1", "16:9")), +) +def test_vertex_imagen_generation_size_reaches_request(size: str, expected_ratio: str | None) -> None: + model: Final = "imagen-4.0-generate-001" + config: Final = VertexAIImagenImageGenerationConfig() + request: Final = config.transform_image_generation_request( + model, "mountains", config.map_openai_params({"size": size, "n": 2}, {}, model, False), {}, {} + ) + assert request["parameters"] == { + "sampleCount": 2, + **({"aspectRatio": expected_ratio} if expected_ratio is not None else {}), + } + + +@pytest.mark.parametrize("size", ("auto", "unknown", "0x1024", "-1x1024", "1024x", "3:2")) +@pytest.mark.parametrize( + "config_type, model, config_key, expected", + ( + ( + VertexAIGeminiImageGenerationConfig, + "gemini-3.1-flash-image", + "generationConfig", + {"responseModalities": ["IMAGE"], "imageConfig": {"aspectRatio": "1:1"}}, + ), + ( + VertexAIImagenImageGenerationConfig, + "imagen-4.0-generate-001", + "parameters", + {"sampleCount": 1, "aspectRatio": "1:1"}, + ), + ), +) +def test_vertex_generation_preserves_legacy_size_fallback( + size: str, + config_type: type[VertexAIGeminiImageGenerationConfig] | type[VertexAIImagenImageGenerationConfig], + model: str, + config_key: str, + expected: dict[str, object], +) -> None: + config: Final = config_type() + request: Final = config.transform_image_generation_request( + model, "mountains", config.map_openai_params({"size": size}, {}, model, False), {}, {} + ) + assert request[config_key] == expected + + +@pytest.mark.parametrize("config_type", (VertexAIGeminiImageGenerationConfig, VertexAIImagenImageGenerationConfig)) +def test_vertex_generation_warns_when_size_ratio_is_approximated( + config_type: type[VertexAIGeminiImageGenerationConfig] | type[VertexAIImagenImageGenerationConfig], + caplog: pytest.LogCaptureFixture, +) -> None: + mapped: Final = config_type().map_openai_params({"size": "100000x1"}, {}, "gemini-3.1-flash-image", False) + assert "100000x1" in caplog.text + assert str(mapped["aspectRatio"]) in caplog.text + + +@pytest.mark.parametrize("size_first", (True, False)) +@pytest.mark.parametrize( + "overrides", + ( + {"aspect_ratio": "9:16", "image_size": "2K"}, + {"imageConfig": {"aspectRatio": "9:16", "imageSize": "2K"}}, + ), +) +def test_vertex_gemini_explicit_image_settings_override_size(size_first: bool, overrides: dict[str, object]) -> None: + model: Final = "gemini-3.1-flash-image" + config: Final = VertexAIGeminiImageGenerationConfig() + params: Final = {"size": "1536x1024", **overrides} if size_first else {**overrides, "size": "1536x1024"} + request: Final = config.transform_image_generation_request( + model, "mountains", config.map_openai_params(params, {}, model, False), {}, {} + ) + assert request["generationConfig"]["imageConfig"] == {"aspectRatio": "9:16", "imageSize": "2K"} + + class TestVertexAIGeminiImageGenerationConfig: def setup_method(self): """Set up test fixtures""" @@ -55,14 +148,13 @@ class TestVertexAIGeminiImageGenerationConfig: result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) assert result.get("aspectRatio") == "16:9" - def test_map_size_to_aspect_ratio(self): - """Test size to aspect ratio mapping""" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" - assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" - assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3" - assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4" - assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + @pytest.mark.parametrize( + "size, expected_ratio", + (("1024x1024", "1:1"), ("1792x1024", "16:9"), ("1024x1792", "9:16"), ("1280x896", "4:3"), ("896x1280", "3:4")), + ) + def test_map_openai_params_legacy_sizes(self, size: str, expected_ratio: str) -> None: + mapped: Final = self.config.map_openai_params({"size": size}, {}, "gemini-3.1-flash-image", False) + assert mapped["aspectRatio"] == expected_ratio def test_get_supported_openai_params_includes_native_gemini_params(self): """Test that native Gemini imageConfig params are supported""" @@ -474,11 +566,10 @@ class TestVertexAIImagenImageGenerationConfig: result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) assert result.get("aspectRatio") == "1:1" - def test_map_size_to_aspect_ratio(self): - """Test size to aspect ratio mapping""" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" - assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + @pytest.mark.parametrize("size, expected_ratio", (("1024x1024", "1:1"), ("1792x1024", "16:9"))) + def test_map_openai_params_legacy_sizes(self, size: str, expected_ratio: str) -> None: + mapped: Final = self.config.map_openai_params({"size": size}, {}, "imagen-4.0-generate-001", False) + assert mapped["aspectRatio"] == expected_ratio def test_transform_image_generation_request_basic(self): """Test basic request transformation""" diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 21bda6f0345..1c432ec79da 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -1,4 +1,5 @@ +import base64 import copy import json from collections.abc import Mapping @@ -40,6 +41,33 @@ def test_vertex_ai_anthropic_thinking_param(model, expected_thinking): assert "thinking" not in supported_openai_params +@pytest.mark.parametrize( + ("labels", "headers", "expected_header"), + [ + ({"team": "triage", "env": "prod"}, {}, {"team": "triage", "env": "prod"}), + ({"team": "triage", "cost": 3}, {}, {"team": "triage"}), + ({"cost": 3}, {}, None), + ("team=triage", {}, None), + ({"team": "triage"}, {"x-vertex-ai-labels": "e30="}, {}), + ], + ids=["string_labels", "non_string_values_dropped", "no_string_values", "not_a_mapping", "caller_header_wins"], +) +def test_vertex_anthropic_labels_move_from_body_to_header( + labels: object, headers: dict[str, str], expected_header: dict[str, str] | None +) -> None: + body: Final = VertexAIAnthropicConfig().transform_request( + model="claude-sonnet-5-5", + messages=[{"role": "user", "content": "Hi"}], + optional_params={"max_tokens": 16, "anthropic_version": "vertex-2023-10-16", "labels": labels}, + litellm_params={}, + headers=headers, + ) + sent_labels: Final = next((value for name, value in headers.items() if name.lower() == "x-vertex-ai-labels"), None) + + assert "labels" not in body + assert (None if sent_labels is None else json.loads(base64.b64decode(sent_labels))) == expected_header + + @pytest.mark.parametrize( ("stream", "operation"), [(False, "rawPredict"), (True, "streamRawPredict")], diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index c7bb05e5d30..037b0a6c29a 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -5,9 +5,11 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec import json import logging import unittest +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi.testclient import TestClient @@ -615,6 +617,184 @@ async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_e assert counter.await_args.kwargs["request"].tools is None +def _count_on_claude( + monkeypatch: pytest.MonkeyPatch, peer: Callable[[httpx.Request], httpx.Response], headers: Mapping[str, str] +) -> httpx.Response: + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + claude: Final = {"model": "anthropic/claude-zeta-9", "api_key": "test-key", "api_base": "https://count-tokens.test"} + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-claude", + "litellm_params": {**claude, "order": 2}, + "model_info": {"id": "count-log-fallback-deployment"}, + }, + { + "model_name": "count-log-claude", + "litellm_params": {**claude, "order": 1}, + "model_info": {"id": "count-log-routed-deployment"}, + }, + ] + ), + ) + monkeypatch.setattr( + anthropic_count_tokens_handler, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + return TestClient(app).post( + "/v1/messages/count_tokens", + json={"model": "count-log-claude", "messages": [{"role": "user", "content": "Hello"}]}, + headers=dict(headers), + ) + + +def test_count_tokens_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _count_on_claude( + monkeypatch, + lambda request: httpx.Response(200, json={"input_tokens": 77}, request=request), + {"x-litellm-call-id": "count-log-msg-claude"}, + ) + + assert response.json() == {"input_tokens": 77, "litellm_count_source": "provider"}, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-msg-claude", + "request_model": "count-log-claude", + "model_id": "count-log-routed-deployment", + "model_used": "claude-zeta-9", + "input_tokens": 77, + "count_source": "provider", + "status_code": 200, + }, + ) + + +def test_count_tokens_provider_rejection_logs_its_status_and_answers_the_same_error( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + rejection: Final = { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages.0.content.0: invalid base64 data"}, + } + + response: Final = _count_on_claude( + monkeypatch, + lambda request: httpx.Response(400, json=rejection, request=request), + {"x-litellm-call-id": "count-log-msg-error"}, + ) + + assert response.status_code == 400, response.text + assert response.json()["error"] == rejection["error"] + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-msg-error", + "request_model": "count-log-claude", + "model_id": "count-log-routed-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def _count_on_local_model(monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object]) -> httpx.Response: + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-local", + "litellm_params": {"model": "hosted_vllm/count-log-local", "api_base": "https://local.test/v1"}, + "model_info": {"id": "count-log-local-deployment"}, + } + ] + ), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + return TestClient(app).post( + "/v1/messages/count_tokens", json=dict(body), headers={"x-litellm-call-id": "count-log-local"} + ) + + +def test_count_tokens_unsupported_content_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _count_on_local_model( + monkeypatch, + {"model": "count-log-local", "messages": [{"role": "user", "content": [{"type": "future_block"}]}]}, + ) + + assert (response.status_code, response.json()["error"]["type"]) == (400, "invalid_request_error"), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-local", + "request_model": "count-log-local", + "model_id": "count-log-local-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +@pytest.mark.parametrize( + ("body", "request_model"), + [ + ({"messages": [{"role": "user", "content": "Hello"}]}, None), + ({"model": "count-log-local", "messages": []}, "count-log-local"), + ], +) +def test_count_tokens_rejected_before_counting_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], + body: Mapping[str, object], + request_model: str | None, +) -> None: + response: Final = _count_on_local_model(monkeypatch, body) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-local", + "request_model": request_model, + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 2d9c781b0d1..ae7168ed1e1 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -1537,6 +1537,97 @@ async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): setattr(_proxy_server_mod, attr, val) +@pytest.mark.parametrize( + ("default_budget_id", "expected_limits"), + [ + pytest.param("proxy-default-end-user-budget", (0, 7), id="default-configured"), + pytest.param(None, (None, None), id="no-default"), + ], +) +@pytest.mark.asyncio +async def test_cached_virtual_key_applies_proxy_default_budget_to_new_end_user( + monkeypatch: pytest.MonkeyPatch, + default_budget_id: str | None, + expected_limits: tuple[int | None, int | None], +) -> None: + """Cached virtual keys apply configured limits to new end users.""" + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.proxy_server import hash_token + + end_user_id: Final = "cached-virtual-key-end-user" + monkeypatch.setattr(litellm, "max_end_user_budget_id", default_budget_id) + + api_key: Final = "sk-cached-default-end-user-budget-test" + hashed_key: Final = hash_token(api_key) + cached_token: Final = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + budget: Final = LiteLLM_BudgetTable( + budget_id="proxy-default-end-user-budget", + rpm_limit=0, + tpm_limit=7, + ) + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget) + + user_api_key_cache: Final = UserApiKeyCache() + + mock_proxy_logging_obj: Final = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + monkeypatch.setattr(_proxy_server_mod, "prisma_client", mock_prisma_client) + monkeypatch.setattr(_proxy_server_mod, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(_proxy_server_mod, "proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr(_proxy_server_mod, "master_key", "sk-master-key") + monkeypatch.setattr(_proxy_server_mod, "general_settings", {}) + monkeypatch.setattr(_proxy_server_mod, "llm_model_list", []) + monkeypatch.setattr(_proxy_server_mod, "llm_router", None) + monkeypatch.setattr(_proxy_server_mod, "open_telemetry_logger", None) + monkeypatch.setattr(_proxy_server_mod, "model_max_budget_limiter", MagicMock()) + monkeypatch.setattr(_proxy_server_mod, "user_custom_auth", None) + monkeypatch.setattr(_proxy_server_mod, "jwt_handler", None) + monkeypatch.setattr(_proxy_server_mod, "litellm_proxy_admin_name", "admin") + + request: Final = Request( + scope={ + "type": "http", + "headers": [(b"x-litellm-end-user-id", end_user_id.encode())], + } + ) + request._url = URL(url="/chat/completions") + + with patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=cached_token, + ): + result: Final = await user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert (result.end_user_rpm_limit, result.end_user_tpm_limit) == expected_limits + if default_budget_id is not None: + assert result.end_user_id == end_user_id + + @pytest.mark.asyncio async def test_master_key_auth_sets_via_virtual_key_marker(): """Master-key requests must also be stamped by overwrite_user_with_key_hash; diff --git a/tests/unit/proxy/client/cli/test_agents.py b/tests/unit/proxy/client/cli/test_agents.py index 3481346076f..3c79b916202 100644 --- a/tests/unit/proxy/client/cli/test_agents.py +++ b/tests/unit/proxy/client/cli/test_agents.py @@ -1,10 +1,13 @@ +import functools import inspect import json import os import stat import subprocess import sys +from collections.abc import Mapping, Sequence from pathlib import Path +from typing import Final from unittest.mock import patch import click @@ -31,8 +34,14 @@ from litellm.proxy.client.cli.commands.agents import ( verify_subscription_key, ) from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError +from litellm.proxy.client.cli.commands.config import save_config from litellm.proxy.client.cli.main import cli as lite +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" @@ -189,6 +198,30 @@ def _overlay(args): return json.loads(Path(args[2]).read_text()) +def _codex_config(args): + assert set(args[0::2]) == {"-c"} + return tomllib.loads("\n".join(args[1::2])) + + +def _set_codex_key(config: Mapping[str, object], keys: Sequence[str], value: object) -> Mapping[str, object]: + head, *rest = keys + if not rest: + return {**config, head: value} + nested: Final = config.get(head, {}) + assert isinstance(nested, Mapping) + return {**config, head: _set_codex_key(nested, rest, value)} + + +def _with_codex_override(config: Mapping[str, object], override: str) -> Mapping[str, object]: + path, value = override.split("=", 1) + return _set_codex_key(config, path.split("."), tomllib.loads(f"value = {value}")["value"]) + + +def _codex_config_over(user_config_toml: str, args: Sequence[str]) -> Mapping[str, object]: + assert set(args[0::2]) == {"-c"} + return functools.reduce(_with_codex_override, args[1::2], tomllib.loads(user_config_toml)) + + class TestAgentProfile: def test_claude_is_anthropic(self): name, profiles = agent_profile("claude") @@ -254,6 +287,7 @@ class TestBuildAgentEnv: env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"})) assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert env["OPENAI_API_KEY"] == "sk-key" + assert "LITELLM_PROXY_API_KEY" not in env assert "ANTHROPIC_BASE_URL" not in env assert "ENABLE_TOOL_SEARCH" not in env assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in env @@ -317,11 +351,21 @@ class TestBuildAgentEnv: ) assert env["ANTHROPIC_CUSTOM_HEADERS"].splitlines() == ["x-team: blue", "x-litellm-api-key: Bearer sk-new"] - def test_subscription_leaves_openai_agents_on_the_bearer_key(self): - env = build_agent_env({}, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True) - assert env["OPENAI_API_KEY"] == "sk-key" + def test_subscription_routes_openai_agents_without_offering_the_key_as_an_openai_credential(self): + env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"}), subscription=True) + assert "OPENAI_BASE_URL" not in env + assert env["LITELLM_PROXY_API_KEY"] == "sk-key" + assert "OPENAI_API_KEY" not in env assert "ANTHROPIC_CUSTOM_HEADERS" not in env + def test_subscription_leaves_the_users_own_openai_key_and_base_url_as_they_were(self): + base = {"OPENAI_API_KEY": "sk-users-openai", "OPENAI_BASE_URL": "https://api.openai.com/v1"} + env = build_agent_env(base, "http://localhost:4000", "sk-key", frozenset({"openai"}), subscription=True) + assert env["OPENAI_API_KEY"] == "sk-users-openai" + assert env["OPENAI_BASE_URL"] == "https://api.openai.com/v1" + assert env["LITELLM_PROXY_API_KEY"] == "sk-key" + assert base == {"OPENAI_API_KEY": "sk-users-openai", "OPENAI_BASE_URL": "https://api.openai.com/v1"} + def test_default_mode_does_not_touch_custom_headers(self): env = build_agent_env( {"ANTHROPIC_CUSTOM_HEADERS": "x-team: blue"}, "http://localhost:4000", "sk-key", frozenset({"anthropic"}) @@ -354,6 +398,67 @@ class TestAgentLaunchArgs: agent_launch_args("codex", "http://localhost:4000") ) + def test_codex_default_provider_reads_the_key_from_openai_api_key(self): + assert _codex_config(agent_launch_args("codex", "http://localhost:4000/")) == { + "model_provider": "litellm", + "model_providers": { + "litellm": { + "name": "LiteLLM proxy", + "base_url": "http://localhost:4000/v1", + "wire_api": "responses", + "supports_websockets": False, + "requires_openai_auth": False, + "env_key": "OPENAI_API_KEY", + "http_headers": {}, + } + }, + } + + def test_codex_subscription_provider_keeps_the_chatgpt_login_and_sends_the_key_in_the_litellm_header(self): + assert _codex_config( + agent_launch_args("/usr/local/bin/codex", "http://localhost:4000/", subscription=True) + ) == { + "model_provider": "lite_subscription", + "model_providers": { + "lite_subscription": { + "name": "LiteLLM proxy", + "base_url": "http://localhost:4000/v1", + "wire_api": "responses", + "supports_websockets": False, + "requires_openai_auth": True, + "env_http_headers": {"x-litellm-api-key": "LITELLM_PROXY_API_KEY"}, + "http_headers": {}, + } + }, + } + + def test_codex_subscription_leaves_the_cookbook_provider_table_out_of_the_launch(self): + user_config = """ +model_provider = "litellm_subscription" + +[model_providers.litellm_subscription] +name = "LiteLLM subscription" +base_url = "https://litellm.example.com/v1" +wire_api = "responses" +requires_openai_auth = true +env_key = "OPENAI_API_KEY" +http_headers = { "x-team" = "qa" } +""" + config = _codex_config_over( + user_config, agent_launch_args("codex", "http://localhost:4000/", subscription=True) + ) + launched = config["model_providers"][config["model_provider"]] + assert "env_key" not in launched + assert launched["requires_openai_auth"] is True + assert launched["base_url"] == "http://localhost:4000/v1" + assert config["model_providers"]["litellm_subscription"] == tomllib.loads(user_config)["model_providers"][ + "litellm_subscription" + ] + + @pytest.mark.parametrize("agent", ["claude", "opencode", "pi", "mytool"]) + def test_subscription_adds_no_args_for_agents_without_a_proxy_provider(self, agent): + assert agent_launch_args(agent, "http://localhost:4000", subscription=True) == [] + def test_pi_gets_no_static_args(self): assert agent_launch_args("pi", "http://localhost:4000") == [] @@ -1286,6 +1391,25 @@ class TestRunAgent: with pytest.raises(AgentRunError): run_agent("http://localhost:4000", "sk-key", []) + @pytest.mark.parametrize("subscription", [True, False]) + def test_codex_stop_hook_is_added_only_when_codex_sends_the_key_as_its_openai_key(self, subscription): + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex"], + subscription=subscription, + base_env={"OPENAI_API_KEY": "sk-users-own-openai-key"}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: None, + verify_subscription=lambda *a: None, + sync_models=lambda *a: ModelSyncSkipped("not needed"), + warn=lambda message: None, + launcher=lambda p, a, e: launched.update(args=tuple(a)), + preparers={"codex": lambda *a: ("-c", "hooks.Stop=lite")}, + ) + assert ("hooks.Stop=lite" in launched["args"]) is not subscription + def test_subscription_launch_probes_with_the_header_and_exports_it(self): probed = [] launched = {} @@ -1574,6 +1698,56 @@ class TestRunAgent: read_settings=lambda path: {}, ) + def test_codex_subscription_launch_keeps_the_chatgpt_login_and_needs_no_settings_overlay(self): + probed = [] + overlays = [] + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex", "exec", "--settings", "kept"], + subscription=True, + base_env={"PATH": "/usr/bin"}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: probed.append(("bearer", a)), + verify_subscription=lambda *a: probed.append(("header", a)), + sync_models=lambda *a: {}, + preparers={}, + write_overlay=lambda text, environ: overlays.append(text), + launcher=lambda p, a, e: launched.update(path=p, args=tuple(a), env=dict(e)), + ) + assert probed == [("header", ("http://localhost:4000", "sk-key"))] + assert overlays == [] + assert launched["path"] == "/usr/local/bin/codex" + assert launched["args"][0] == "codex" + assert launched["args"][-3:] == ("exec", "--settings", "kept") + assert _codex_config(launched["args"][1:-3])["model_provider"] == "lite_subscription" + assert "OPENAI_BASE_URL" not in launched["env"] + assert launched["env"]["LITELLM_PROXY_API_KEY"] == "sk-key" + assert "OPENAI_API_KEY" not in launched["env"] + + @pytest.mark.parametrize( + "sync_result", + [ModelSyncArgs(("-c", 'model_catalog_json="/tmp/c.json"')), ModelSyncSkipped("lists no chat models")], + ) + def test_codex_subscription_launch_leaves_the_model_list_to_codex(self, sync_result): + warnings = [] + launched = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex"], + subscription=True, + base_env={}, + which=lambda name: "/usr/local/bin/codex", + verify_subscription=lambda *a: None, + sync_models=lambda *a: sync_result, + warn=warnings.append, + launcher=lambda p, a, e: launched.update(args=tuple(a)), + ) + assert warnings == [] + assert "model_catalog_json" not in json.dumps(launched["args"]) + def test_default_launch_still_probes_with_the_bearer(self): probed = [] run_agent( @@ -1964,7 +2138,11 @@ class TestAgentCommands: assert captured["command"] == ["opencode"] assert "routing OpenCode through proxy at http://localhost:4000" in result.output - def test_claude_subscription_flag_reaches_run_agent_and_is_not_forwarded(self): + @pytest.mark.parametrize( + ("agent", "user_args", "login"), + [("claude", ["--resume"], "claude.ai"), ("codex", ["exec", "hi"], "ChatGPT")], + ) + def test_subscription_flag_reaches_run_agent_and_is_not_forwarded(self, agent, user_args, login): captured = {} def fake_run_agent(base_url, api_key, command, **kwargs): @@ -1973,61 +2151,85 @@ class TestAgentCommands: with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent): result = self.runner.invoke( - _agent_command("claude"), - ["--subscription", "--resume"], + _agent_command(agent), + ["--subscription", *user_args], obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, ) assert result.exit_code == 0, result.output assert captured["subscription"] is True - assert captured["command"] == ["claude", "--resume"] - assert "keeping your claude.ai login" in result.output + assert captured["command"] == [agent, *user_args] + assert f"keeping your {login} login" in result.output + @pytest.mark.parametrize(("agent", "login"), [("claude", "claude.ai"), ("codex", "ChatGPT")]) @pytest.mark.parametrize( - ("configured", "args", "expected"), + ("saved", "args", "expected"), [ - (False, [], False), - (True, [], True), - (True, ["--no-subscription"], False), - (False, ["--subscription"], True), + (None, [], False), + ("false", [], False), + ("true", [], True), + ("true", ["--no-subscription"], False), + ("false", ["--subscription"], True), ], ) - def test_subscription_comes_from_the_flag_else_the_saved_setting(self, configured, args, expected): + def test_subscription_comes_from_the_flag_else_the_saved_setting(self, agent, login, saved, args, expected): + if saved is not None: + save_config({f"{agent}_subscription": saved}) captured = {} - with ( - patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=configured), - patch( - f"{AGENTS_MODULE}.run_agent", - side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)), - ), + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription"), command=list(c)), ): result = self.runner.invoke( - _agent_command("claude"), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} + _agent_command(agent), args, obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} ) assert result.exit_code == 0, result.output assert captured["sub"] is expected - assert captured["command"] == ["claude"] - assert ("keeping your claude.ai login" in result.output) is expected + assert captured["command"] == [agent] + assert (f"keeping your {login} login" in result.output) is expected - @pytest.mark.parametrize("agent", ["codex", "opencode"]) - def test_the_saved_claude_setting_does_not_change_other_agents(self, agent): + @pytest.mark.parametrize( + ("saved_key", "agent"), + [ + ("claude_subscription", "codex"), + ("claude_subscription", "opencode"), + ("codex_subscription", "claude"), + ("codex_subscription", "opencode"), + ], + ) + def test_a_saved_subscription_setting_only_changes_its_own_agent(self, saved_key, agent): + save_config({saved_key: "true"}) captured = {} - with ( - patch(f"{AGENTS_MODULE}.claude_subscription_default", return_value=True), - patch( - f"{AGENTS_MODULE}.run_agent", - side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")), - ), + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(sub=kw.get("subscription")), ): result = self.runner.invoke( _agent_command(agent), [], obj={"base_url": "http://localhost:4000", "api_key": "sk-key"} ) assert result.exit_code == 0, result.output assert captured["sub"] is False + assert "keeping your" not in result.output - @pytest.mark.parametrize("agent", ["codex", "opencode", "pi"]) - def test_only_claude_offers_the_subscription_flag(self, agent): - assert "--subscription" not in {opt for param in _agent_command(agent).params for opt in param.opts} + def test_only_claude_and_codex_offer_the_subscription_flag(self): + offering = { + command.name + for command in agent_commands() + if "--subscription" in {opt for param in command.params for opt in param.opts} + } + assert offering == {"claude", "codex"} + + @pytest.mark.parametrize( + ("agent", "login", "own_credential"), + [("claude", "claude.ai", "Claude Code's"), ("codex", "ChatGPT", "Codex's")], + ) + def test_subscription_help_names_the_agents_own_login_and_its_setting(self, agent, login, own_credential): + result = self.runner.invoke(lite, [agent, "--help"], prog_name="lite") + assert result.exit_code == 0, result.output + text = " ".join(result.output.split()) + assert f"Keep using your {login} login" in text + assert f"instead of replacing {own_credential} own credential" in text + assert f"lite config set {agent}_subscription true" in text def test_skip_verify_is_consumed_not_forwarded(self): captured = {} 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/client/cli/test_config_commands.py b/tests/unit/proxy/client/cli/test_config_commands.py index 2cd7a76ab95..e99a4311caf 100644 --- a/tests/unit/proxy/client/cli/test_config_commands.py +++ b/tests/unit/proxy/client/cli/test_config_commands.py @@ -9,12 +9,12 @@ from click.testing import CliRunner from litellm.litellm_core_utils.private_json import write_private_json from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands.config import ( - claude_subscription_default, get_config_file_path, get_config_value, load_config, pkce_default, save_config, + subscription_default, ) from litellm.proxy.client.cli.interface import show_commands @@ -362,33 +362,48 @@ class TestWritePrivateJson: assert list(tmp_path.glob(".tmp-*")) == [] -class TestClaudeSubscriptionSetting: +@pytest.mark.parametrize(("agent", "key"), [("claude", "claude_subscription"), ("codex", "codex_subscription")]) +class TestSubscriptionSetting: @pytest.mark.parametrize("word", ["true", "TRUE", "on", "1", "yes"]) - def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, word): - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word]) + def test_truthy_words_are_stored_as_true(self, cli_runner, isolated_home, agent, key, word): + result = cli_runner.invoke(cli, ["config", "set", key, word]) assert result.exit_code == 0, result.output - assert load_config()["claude_subscription"] == "true" - assert claude_subscription_default() is True + assert load_config()[key] == "true" + assert subscription_default(agent) is True @pytest.mark.parametrize("word", ["false", "Off", "0", "no"]) - def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, word): - cli_runner.invoke(cli, ["config", "set", "claude_subscription", "true"]) - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", word]) + def test_falsy_words_are_stored_as_false(self, cli_runner, isolated_home, agent, key, word): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + result = cli_runner.invoke(cli, ["config", "set", key, word]) assert result.exit_code == 0, result.output - assert load_config()["claude_subscription"] == "false" - assert claude_subscription_default() is False + assert load_config()[key] == "false" + assert subscription_default(agent) is False - def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home): - result = cli_runner.invoke(cli, ["config", "set", "claude_subscription", "maybe"]) + def test_anything_else_is_rejected_and_nothing_is_saved(self, cli_runner, isolated_home, agent, key): + result = cli_runner.invoke(cli, ["config", "set", key, "maybe"]) assert result.exit_code != 0 - assert "true or false" in result.output - assert "claude_subscription" not in load_config() + assert f"{key} must be true or false" in result.output + assert key not in load_config() - def test_unset_means_off(self, isolated_home): - assert claude_subscription_default() is False + def test_unset_means_off(self, cli_runner, isolated_home, agent, key): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + result = cli_runner.invoke(cli, ["config", "unset", key]) + + assert result.exit_code == 0, result.output + assert subscription_default(agent) is False + + def test_turning_it_on_leaves_every_other_agent_off(self, cli_runner, isolated_home, agent, key): + cli_runner.invoke(cli, ["config", "set", key, "true"]) + + assert {name: subscription_default(name) for name in ("claude", "codex", "opencode", "pi")} == { + "claude": agent == "claude", + "codex": agent == "codex", + "opencode": False, + "pi": False, + } class TestPkceSetting: 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_log.py b/tests/unit/proxy/common_utils/test_token_count_log.py new file mode 100644 index 00000000000..94b976d88cf --- /dev/null +++ b/tests/unit/proxy/common_utils/test_token_count_log.py @@ -0,0 +1,45 @@ +import asyncio +from collections.abc import Callable + +import pytest + +from litellm.proxy.common_utils.token_count_log import logged_token_count +from litellm.types.utils import TokenCountResponse + + +class _TokenizerPanic(BaseException): + pass + + +async def _count_ending_in(error: BaseException) -> TokenCountResponse: + raise error + + +@pytest.mark.asyncio +async def test_a_count_ending_in_a_base_exception_logs_a_500( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + with pytest.raises(_TokenizerPanic): + await logged_token_count(_count_ending_in(_TokenizerPanic()), "count-log-panic", "count-log-gpt", None) + + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-panic", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 500, + }, + ) + + +@pytest.mark.asyncio +async def test_a_cancelled_count_writes_no_line( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + with pytest.raises(asyncio.CancelledError): + await logged_token_count(_count_ending_in(asyncio.CancelledError()), "count-log-cancelled", "count-log-gpt", None) + + assert token_count_log_lines() == () 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/conftest.py b/tests/unit/proxy/conftest.py index 58c9a5e4d76..e13eddef43b 100644 --- a/tests/unit/proxy/conftest.py +++ b/tests/unit/proxy/conftest.py @@ -3,6 +3,7 @@ import asyncio import copy import inspect +import logging import os import tempfile import warnings @@ -17,6 +18,7 @@ from prisma.errors import ClientNotConnectedError import litellm import litellm.proxy.proxy_server +from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes from litellm.integrations.otel.model.payloads import ServiceSpanData from litellm.integrations.otel.model.spans import service_span_name @@ -424,6 +426,31 @@ def fresh_agent_read_through(monkeypatch): return read_through +TOKEN_COUNT_LOG_FIELDS: Final = ( + "litellm_call_id", + "request_model", + "model_id", + "model_used", + "input_tokens", + "count_source", + "status_code", +) + + +@pytest.fixture +def token_count_log_lines( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> Iterator[Callable[[], tuple[dict[str, object], ...]]]: + """The structured fields of every INFO token count line the proxy logger writes while the test runs.""" + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + with caplog.at_level(logging.INFO, logger=verbose_proxy_logger.name): + yield lambda: tuple( + {field: record.__dict__[field] for field in TOKEN_COUNT_LOG_FIELDS} + for record in caplog.records + if record.levelno == logging.INFO and "count_source" in record.__dict__ + ) + + @pytest.fixture def postgres_span_names() -> Iterator[Callable[[], Awaitable[tuple[str, ...]]]]: """The ``postgres.{verb} {table}`` names OTel would render for every DB service event 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_endpoints.py b/tests/unit/proxy/google_endpoints/test_endpoints.py index 17af52b5579..80e45edf897 100644 --- a/tests/unit/proxy/google_endpoints/test_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_endpoints.py @@ -94,6 +94,7 @@ async def test_proxy_gemini_to_openai_like_model_token_counting(): request=Request( scope={ "type": "http", + "headers": [], "parsed_body": ( ["contents"], {"contents": [{"parts": [{"text": "Hello, how are you?"}]}]}, 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_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index 2777d18a399..4867ca9ab63 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -3,9 +3,11 @@ Test to verify the Google GenAI proxy API endpoints """ +from collections.abc import Callable from typing import Final from unittest.mock import AsyncMock, patch +import httpx import pytest from litellm.types.utils import TokenCountResponse @@ -240,3 +242,79 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() assert body["promptTokensDetails"] == [] assert body["litellm_count_source"] == "token_count_model" assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5" + + +@pytest.mark.parametrize("path", ["/v1beta/models/count-log-claude:countTokens", "/models/count-log-claude:countTokens"]) +def test_google_count_tokens_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], path: str +) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm import Router + from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + claude: Final = {"model": "anthropic/claude-zeta-9", "api_key": "test-key", "api_base": "https://count-tokens.test"} + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-claude", + "litellm_params": claude, + "model_info": {"id": "count-log-gemini-deployment"}, + } + ] + ), + ) + monkeypatch.setattr( + anthropic_count_tokens_handler, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, json={"input_tokens": 77}, request=request) + ) + ), + ) + + response: Final = _build_test_client().post( + path, + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + headers={"x-litellm-call-id": "count-log-gemini"}, + ) + + assert (response.status_code, response.json()["totalTokens"]) == (200, 77), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-gemini", + "request_model": "count-log-claude", + "model_id": "count-log-gemini-deployment", + "model_used": "claude-zeta-9", + "input_tokens": 77, + "count_source": "provider", + "status_code": 200, + }, + ) + + +def test_google_count_tokens_rejected_body_logs_the_400_it_answers( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + response: Final = _build_test_client().post( + "/v1beta/models/count-log-claude:countTokens", + json={"contents": "not a list"}, + headers={"x-litellm-call-id": "count-log-gemini-bad-body"}, + ) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-gemini-bad-body", + "request_model": "count-log-claude", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) 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/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 3d1e494e36c..8245c261a12 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -15,28 +15,38 @@ Redis. """ import asyncio +import base64 +import copy +import io import json from datetime import datetime, timedelta from typing import Any, Dict, Final import pytest +import soundfile as sf +from fastapi import HTTPException +from PIL import Image +from litellm import token_counter from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, estimate_image_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _AUDIO_BYTES_PER_TOKEN, + DEFAULT_AUDIO_TOKEN_ESTIMATE, PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, RateLimitedModel, - _AUDIO_BYTES_PER_TOKEN, - PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, -) -from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _call_id_from_callback_kwargs, _request_stash, get_or_create_request_stash, get_request_stash, ) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, +) from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.llms.openai import ( InputTokensDetails, ResponseAPIUsage, @@ -3775,3 +3785,284 @@ def test_reservation_counts_a_compaction_item_by_its_summary(rate_limiter): ) == handler._estimate_input_and_output_tokens( data={"model": "gpt-4o", "input": _COMPACTED_INPUT[2:], "max_output_tokens": 1}, call_type="aresponses" ) + + +_QUOTA_TEXT: Final = "a" * 400 +_QUOTA_JSON_TEXT: Final = '{"type":"input_image","image_url":"data:image/png;base64,AAAA"}' + + +@pytest.mark.parametrize( + ("call_type", "request_data", "expected"), + [ + ("acompletion", {"messages": [{"role": "user", "content": _QUOTA_TEXT}], "max_tokens": 16}, (100, 16)), + ("acompletion", {"messages": [{"role": "user", "content": _QUOTA_JSON_TEXT}], "max_tokens": 16}, (15, 16)), + ( + "anthropic_messages", + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t", + "content": [{"type": "text", "text": _QUOTA_TEXT}], + } + ], + } + ], + "max_tokens": 16, + }, + (0, 16), + ), + ( + "acompletion", + { + "messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "t", "type": "function", "function": {"name": "f", "arguments": _QUOTA_TEXT}} + ], + } + ], + "max_tokens": 16, + }, + (0, 16), + ), + ("aresponses", {"input": _QUOTA_TEXT, "max_output_tokens": 16}, (100, 16)), + ("aresponses", {"input": _QUOTA_TEXT, "instructions": "b" * 400, "max_output_tokens": 16}, (100, 16)), + ( + "aresponses", + {"input": [{"role": "user", "content": _QUOTA_TEXT}], "instructions": "b" * 400, "max_output_tokens": 16}, + (107, 16), + ), + ( + "aresponses", + { + "input": [ + { + "type": "file_search_call", + "id": "f", + "status": "completed", + "queries": ["q"], + "results": [ + { + "file_id": "file_1", + "filename": "result.txt", + "attributes": {}, + "score": 1, + "text": _QUOTA_TEXT, + } + ], + } + ], + "max_output_tokens": 16, + }, + (146, 16), + ), + ( + "aresponses", + { + "input": [{"type": "custom_tool_call_output", "call_id": "t", "output": _QUOTA_JSON_TEXT}], + "max_output_tokens": 16, + }, + (32, 16), + ), + ("aresponses", {"input": "", "instructions": "b" * 400}, (0, 0)), + ("aembedding", {"input": [_QUOTA_TEXT]}, (100, 0)), + ("atext_completion", {"prompt": [_QUOTA_TEXT], "max_tokens": 16}, (100, 16)), + ], +) +def test_media_correction_preserves_exact_non_media_reservations( + rate_limiter: tuple[RateLimitHandler, DualCache], + call_type: str, + request_data: dict[str, object], + expected: tuple[int, int], +) -> None: + handler, _ = rate_limiter + assert handler._estimate_input_and_output_tokens(request_data, call_type=call_type) == expected + + +def _quota_media_payload(kind: str, compressed: bool) -> str: + output: Final = io.BytesIO() + if kind == "image": + Image.new("RGB", (1024, 1024)).save(output, format="PNG", compress_level=9 if compressed else 0) + else: + sample_rate: Final = 8000 if compressed else 48000 + with sf.SoundFile(output, "w", samplerate=sample_rate, channels=1, format="WAV") as audio: + audio.buffer_write(b"\0" * (sample_rate * 2), dtype="int16") + return base64.b64encode(output.getvalue()).decode() + + +def _quota_media_request(kind: str, payload: str, call_type: str) -> dict[str, object]: + block: Final = ( + {"type": "input_audio", "input_audio": {"data": payload, "format": "wav"}} + if kind == "audio" + else {"type": "input_image", "image_url": "data:image/png;base64," + payload} + if call_type == "aresponses" + else {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": payload}} + if call_type == "anthropic_messages" + else {"type": "image_url", "image_url": {"url": "data:image/png;base64," + payload}} + ) + content: Final = ( + [{"type": "tool_result", "tool_use_id": "t", "content": [block]}] + if call_type == "anthropic_messages" + else [block] + ) + field: Final = "input" if call_type == "aresponses" else "messages" + return { + "model": "gpt-6-astra", + field: [{"role": "user", "content": content}], + "max_tokens": 16, + "max_output_tokens": 16, + } + + +@pytest.mark.parametrize( + ("kind", "call_type"), + [ + ("image", "acompletion"), + ("image", "aresponses"), + ("image", "anthropic_messages"), + ("audio", "acompletion"), + ("audio", "aresponses"), + ], +) +def test_media_reservation_is_independent_of_encoding_and_preserves_request( + rate_limiter: tuple[RateLimitHandler, DualCache], kind: str, call_type: str +) -> None: + handler, _ = rate_limiter + compressed: Final = _quota_media_request(kind, _quota_media_payload(kind, True), call_type) + uncompressed: Final = _quota_media_request(kind, _quota_media_payload(kind, False), call_type) + original: Final = copy.deepcopy(uncompressed) + estimate: Final = handler._estimate_input_and_output_tokens(uncompressed, call_type=call_type) + assert estimate == handler._estimate_input_and_output_tokens(compressed, call_type=call_type) + assert estimate[0] > 0 and estimate[1] == 16 + assert uncompressed == original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) +@pytest.mark.parametrize("limit", [1000, 20000]) +async def test_media_pre_call_rejects_tight_quota_and_reserves_sufficient_quota( + rate_limiter: tuple[RateLimitHandler, DualCache], call_type: str, limit: int +) -> None: + handler, cache = rate_limiter + request: Final = { + key: value + for key, value in _quota_media_request("image", _quota_media_payload("image", True), call_type).items() + if key not in ("max_tokens", "max_output_tokens") + } + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-media-quota"), tpm_limit=limit) + if limit == 1000: + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type=call_type) + assert error.value.status_code == 429 + return + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type=call_type) + stash: Final = get_request_stash() + assert stash is not None + expected: Final = handler._estimate_input_and_output_tokens( + request, call_type=call_type, min_configured_tpm_limit=limit + ) + assert stash.reserved_tokens == sum(expected) + assert expected[0] > 1000 and expected[1] >= expected[0] + + +@pytest.mark.asyncio +async def test_ptu_pre_call_weights_media_input_and_output() -> None: + cache: Final = DualCache() + ceiling: Final = PTUTeamCeiling( + model_group="gpt-6-astra", tpm_limit=10000, output_to_input_ratio=4, cached_input_ratio=0 + ) + + def resolve_ceiling(team: str, model: str) -> PTUTeamCeiling | None: + return ceiling if (team, model) == ("test-team", "gpt-6-astra") else None + + handler: Final = RateLimitHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve_ceiling + ) + request: Final = _quota_media_request("image", _quota_media_payload("image", True), "acompletion") + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-media-ptu"), team_id="test-team"), + cache=cache, + data=request, + call_type="acompletion", + ) + stash: Final = get_request_stash() + assert stash is not None + assert ( + stash.ptu_reserved_tokens + == estimate_image_tokens("gpt-6-astra", "data:image/png;base64," + _quota_media_payload("image", True)) + 4 * 16 + ) + + +@pytest.mark.asyncio +async def test_project_fallback_keeps_legacy_audio_estimate_and_output_policy( + rate_limiter: tuple[RateLimitHandler, DualCache], +) -> None: + handler, cache = rate_limiter + unsupported: Final = {"type": "unsupported_content"} + with pytest.raises(UnsupportedContentTypeError): + token_counter(model="gpt-6-astra", messages=[{"role": "user", "content": [unsupported]}]) + request: Final = { + "model": "gpt-6-astra", + "messages": [ + { + "role": "user", + "content": [ + unsupported, + { + "type": "input_audio", + "input_audio": {"data": _quota_media_payload("audio", True), "format": "wav"}, + }, + ], + } + ], + } + auth: Final = UserAPIKeyAuth( + api_key=hash_token("sk-media-project"), + project_id="media-project", + project_metadata={"model_itpm_limit": {"gpt-6-astra": 1000}, "model_otpm_limit": {"gpt-6-astra": 1000}}, + ) + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type="acompletion") + stash: Final = get_request_stash() + assert stash is not None + assert stash.itpm_reserved_tokens == DEFAULT_AUDIO_TOKEN_ESTIMATE + assert stash.otpm_reserved_tokens == 1 + + +@pytest.mark.asyncio +async def test_single_anthropic_image_exceeds_tight_quota_with_explicit_output_cap( + rate_limiter: tuple[RateLimitHandler, DualCache], +) -> None: + handler, cache = rate_limiter + request: Final = { + "model": "anthropic/claude-haiku-5-5", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": _quota_media_payload("image", True), + }, + } + ], + } + ], + "max_tokens": 64, + } + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-claude-media-quota"), tpm_limit=500), + cache=cache, + data=request, + call_type="anthropic_messages", + ) + assert error.value.status_code == 429 diff --git a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py index d5a8abaca9d..748b291d5ff 100644 --- a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py @@ -2769,15 +2769,22 @@ async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typ @pytest.mark.asyncio @pytest.mark.parametrize( - "case", ["allowed", "credential-free", "missing", "blocked", "key", "budget", "team", "not-router"] + "case", + ["allowed", "credential-free", "team-admin-environment", "missing", "blocked", "key", "budget", "team", "not-router"], ) async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: pytest.MonkeyPatch, case: str) -> None: router: Final = RecordingRouter("SIMPLE") stored_key: Final = "synthetic-server-jev-key" + monkeypatch.setenv("SAVED_JEV_KEY", stored_key) + monkeypatch.setenv("SAVED_JEV_BASE", "https://saved-jev.test") stored_config: Final = { "classifier_type": "jev", "tiers": TIERS, - "jev_classifier_config": {"api_key": stored_key, "api_base": "https://saved-jev.test"}, + "jev_classifier_config": ( + {"api_key": "os.environ/SAVED_JEV_KEY", "api_base": "os.environ/SAVED_JEV_BASE"} + if case == "team-admin-environment" + else {"api_key": stored_key, "api_base": "https://saved-jev.test"} + ), } router.add_deployment( Deployment.model_validate( @@ -2790,7 +2797,10 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: "model_info": { "id": "saved-jev-id", "blocked": case == "blocked", - "team_id": "owner-team" if case == "team" else None, + "team_id": ( + "owner-team" if case == "team" + else "admin-preview-team" if case == "team-admin-environment" else None + ), }, } ) @@ -2799,6 +2809,8 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: actor: Final = ( _configure_member_preview(monkeypatch) if case == "team" + else _configure_team_admin_preview(monkeypatch) + if case == "team-admin-environment" else UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-probe", @@ -2812,12 +2824,17 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: { "prompt": "what is 2+2", "saved_model_id": "missing-id" if case == "missing" else "saved-jev-id", - "team_id": "member-preview-team" if case == "team" else None, + "team_id": ( + "member-preview-team" if case == "team" + else "admin-preview-team" if case == "team-admin-environment" else None + ), }, classifier_type="jev", jev_classifier_config=( {"model": "jev-latest", "timeout_ms": 3000} if case == "credential-free" + else {"api_key": "os.environ/IGNORED_BROWSER_KEY", "api_base": "https://browser-override.test"} + if case == "team-admin-environment" else {"api_key": "masked-key", "api_base": "https://browser-override.test"} ), ) @@ -2856,7 +2873,7 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: assert result.routed_model == "cheap-model" assert evaluation.calls.last.request.headers["authorization"] == f"Bearer {stored_key}" assert stored_key not in result.model_dump_json() - assert evaluation.call_count == (1 if case in ("allowed", "credential-free") else 0) + assert evaluation.call_count == (1 if case in ("allowed", "credential-free", "team-admin-environment") else 0) assert router.recorded_calls == [] await handler.client.aclose() @@ -3378,6 +3395,54 @@ async def test_validate_config_gates_like_the_write_it_rehearses(monkeypatch: py assert not_their_team.value.status_code == 403 +def _configure_team_admin_preview(monkeypatch: pytest.MonkeyPatch) -> UserAPIKeyAuth: + from litellm.proxy._types import LiteLLM_TeamTable + + team: Final = LiteLLM_TeamTable( + team_id="admin-preview-team", + models=["saved-jev", "typesafe/jev-latest", *[TIERS[name][0] for name in TIERS]], + members_with_roles=[{"role": "admin", "user_id": "preview-admin"}], + ) + database: Final = MagicMock() + database.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "premium_user", True) + return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="preview-admin", api_key="sk-preview") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("field", ["api_key", "api_base"]) +async def test_team_admin_cannot_preview_or_validate_environment_references( + monkeypatch: pytest.MonkeyPatch, legacy: bool, field: str +) -> None: + from litellm.types.management_endpoints.auto_router_endpoints import ComplexityRouterConfigValidationRequest + + actor: Final = _configure_team_admin_preview(monkeypatch) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + config: Final = { + "tiers": TIERS, + "classifier_type": "jev" if legacy else "oss_classifier", + "jev_classifier_config" if legacy else "opensource_classifier_config": { + "api_key": "caller-key", "api_base": "https://denied-jev.test", field: "os.environ/PRIVATE_JEV_VALUE", + }, + } + validation: Final = ComplexityRouterConfigValidationRequest( + team_id="admin-preview-team", complexity_router_config=config + ) + with pytest.raises(HTTPException) as denied_validation: + await auto_router_endpoints.validate_complexity_router_config(validation, actor) + assert denied_validation.value.status_code == 403 + with respx.mock(assert_all_called=False) as http, pytest.raises(HTTPException) as denied_preview: + await preview_auto_router_routing( + _request_from({"prompt": "what is 2+2", "team_id": "admin-preview-team"}, **config), + actor, + ROUTING_HTTP_REQUEST, + ) + assert denied_preview.value.status_code == 403 + assert len(http.calls) == 0 + + def _configure_member_preview(monkeypatch: pytest.MonkeyPatch, *, allowed: bool = True) -> UserAPIKeyAuth: from litellm.proxy import proxy_server from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTable diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 9ff2a6de9dd..fc6749aed0f 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -15,6 +15,7 @@ from fastapi.testclient import TestClient from litellm._uuid import uuid from litellm.models.credentials import CredentialItem +from litellm.models.project import LiteLLM_ProjectTable from litellm.proxy._types import ( LiteLLM_ModelTable, @@ -31,6 +32,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( KubernetesPodsResponse, ModelManagementAuthChecks, _get_team_deployments, + _raise_if_model_update_adds_missing_rate_limits, _raise_if_rate_limits_required_but_missing, clear_cache, delete_team_models, @@ -38,6 +40,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( patch_model, update_model, ) +from litellm.proxy.management_helpers.model_quota_policy import REQUIRE_BOTH_LIMITS, ModelQuotaPolicy from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams @@ -7369,12 +7372,19 @@ class TestAutoRouterClassifierDefaultPrompt: assert response.system_prompt == classification_system_prompt(5) +_CUSTOMER_MODEL_ADD_POLICY: Final = ModelQuotaPolicy.model_validate( + { + "providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, + "models": {"veo-*": ["rpm"], "jev-latest": []}, + } +) + + class TestEnforceRpmTpmOnModelAdd: def test_passes_when_disabled_even_without_limits(self): assert ( _raise_if_rate_limits_required_but_missing( - litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), - enforced=False, + model_name="gpt-5.2", team_model=False, litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), policy=None ) is None ) @@ -7382,8 +7392,9 @@ class TestEnforceRpmTpmOnModelAdd: def test_passes_when_enabled_and_both_set(self): assert ( _raise_if_rate_limits_required_but_missing( + model_name="gpt-5.2", team_model=False, litellm_params=LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=1000), - enforced=True, + policy=REQUIRE_BOTH_LIMITS, ) is None ) @@ -7398,14 +7409,618 @@ class TestEnforceRpmTpmOnModelAdd: (LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=-1), "tpm"), ], ) - def test_raises_when_enabled_and_missing(self, params, expected_missing): - from litellm.proxy._types import ProxyException - + def test_raises_when_enabled_and_missing(self, params: LiteLLM_Params, expected_missing: str): with pytest.raises(ProxyException) as exc_info: - _raise_if_rate_limits_required_but_missing(litellm_params=params, enforced=True) - assert expected_missing in str(exc_info.value.message) + _raise_if_rate_limits_required_but_missing( + model_name="gpt-5.2", team_model=False, litellm_params=params, policy=REQUIRE_BOTH_LIMITS + ) + assert exc_info.value.message.startswith(f"{expected_missing} must be set") assert exc_info.value.code == "400" + @pytest.mark.parametrize( + ("model_name", "params"), + [ + ("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo")), + ("my-runway", LiteLLM_Params(model="gen4_turbo", custom_llm_provider="runwayml")), + ("jev-latest", LiteLLM_Params(model="openai/jev-latest")), + ("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=10)), + ("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10, tpm=1000)), + ], + ) + def test_policy_asks_each_model_only_for_the_limits_its_entry_requires( + self, model_name: str, params: LiteLLM_Params + ): + assert ( + _raise_if_rate_limits_required_but_missing( + model_name=model_name, team_model=False, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY + ) + is None + ) + + @pytest.mark.parametrize( + ("model_name", "params", "expected_missing"), + [ + ("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", tpm=1000), "rpm"), + ("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10), "tpm"), + ("openai-gpt", LiteLLM_Params(model="openai/gpt-5.2"), "rpm and tpm"), + ("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo", rpm=0), "rpm"), + ], + ) + def test_policy_rejects_a_model_missing_a_required_or_carrying_a_non_positive_limit( + self, model_name: str, params: LiteLLM_Params, expected_missing: str + ): + with pytest.raises(ProxyException) as exc_info: + _raise_if_rate_limits_required_but_missing( + model_name=model_name, team_model=False, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY + ) + assert exc_info.value.message.startswith(f"{expected_missing} must be set to a positive value for '{model_name}'") + assert exc_info.value.code == "400" + + @pytest.mark.parametrize("api_base", ["https://api.groq.com/openai/v1", "os.environ/QA_GROQ_API_BASE"]) + def test_policy_reads_the_provider_behind_a_known_api_base(self, monkeypatch: pytest.MonkeyPatch, api_base: str): + monkeypatch.setenv("QA_GROQ_API_BASE", "https://api.groq.com/openai/v1") + policy: Final = ModelQuotaPolicy.model_validate({"default": [], "providers": {"groq": ["rpm", "tpm"]}}) + with pytest.raises(ProxyException) as exc_info: + _raise_if_rate_limits_required_but_missing( + model_name="qa-llama", team_model=False, + litellm_params=LiteLLM_Params(model="qa-llama", api_base=api_base), + policy=policy, + ) + assert exc_info.value.message.startswith("rpm and tpm must be set to a positive value for 'qa-llama'") + + @pytest.mark.parametrize( + ("model_name", "params", "shared_missing", "team_missing"), + [ + pytest.param( + "jev-latest", LiteLLM_Params(model="vertex_ai/gemini-3.8-flash"), None, "rpm and tpm", + id="team-cannot-borrow-an-exempt-public-name", + ), + pytest.param( + "video-team", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=5), "tpm", None, + id="team-matches-the-model-it-calls", + ), + ], + ) + def test_policy_matches_a_team_model_by_the_model_it_calls( + self, model_name: str, params: LiteLLM_Params, shared_missing: str | None, team_missing: str | None + ): + assert _model_add_error(model_name, params, team_model=False) == _missing_limits_message(shared_missing, model_name) + assert _model_add_error(model_name, params, team_model=True) == _missing_limits_message(team_missing, model_name) + + +def _missing_limits_message(missing: str | None, model_name: str) -> str | None: + if missing is None: + return None + return ( + f"{missing} must be set to a positive value for '{model_name}' when " + "'enforce_rpm_tpm_on_model_add' is enabled in general_settings" + ) + + +def _model_add_error(model_name: str, params: LiteLLM_Params, team_model: bool) -> str | None: + try: + _raise_if_rate_limits_required_but_missing( + model_name=model_name, team_model=team_model, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY + ) + except ProxyException as error: + return error.message + return None + + +_STORED_RUNWAY: Final = {"model": "runwayml/gen4_turbo"} + + +def _stored_model(litellm_params: Mapping[str, object], team_id: str | None) -> Deployment: + return Deployment( + model_name="media-latest" if team_id is None else "model_name_media-team_team-row", + litellm_params=dict(litellm_params), + model_info={ + "id": "team-row", + "team_id": team_id, + "team_public_model_name": None if team_id is None else "media-latest", + }, + ) + + +@pytest.mark.parametrize( + ("stored", "team_id", "incoming", "missing"), + [ + pytest.param( + _STORED_RUNWAY, "media-team", updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash"), "rpm and tpm", + id="team-model-moves-to-a-provider-that-needs-limits", + ), + pytest.param( + _STORED_RUNWAY, None, updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", rpm=10), "tpm", + id="proxy-model-moves-with-only-one-limit", + ), + pytest.param( + {**_STORED_RUNWAY, "tpm": 1000}, None, updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", rpm=10), + None, id="stored-limit-counts-toward-the-move", + ), + pytest.param( + _STORED_RUNWAY, "media-team", updateLiteLLMParams(model="vertex_ai/veo-3.1-generate-001", rpm=5), None, + id="team-model-moves-to-a-model-entry-it-meets", + ), + pytest.param( + {"model": "vertex_ai/gemini-3.8-flash"}, None, updateLiteLLMParams(api_key="qa-rotated-key"), None, + id="unchanged-route-adds-no-limits", + ), + pytest.param( + {"model": "vertex_ai/gemini-3.8-flash", "rpm": 10, "tpm": 1000}, None, updateLiteLLMParams(rpm=0), None, + id="unchanged-route-can-zero-a-limit-as-on-main", + ), + pytest.param( + {"model": "vertex_ai/gemini-3.8-flash"}, None, updateLiteLLMParams(rpm=10), None, + id="legacy-model-without-tpm-stays-editable", + ), + ], +) +def test_model_update_requires_only_the_limits_its_new_route_adds( + stored: Mapping[str, object], team_id: str | None, incoming: updateLiteLLMParams, missing: str | None +) -> None: + assert _model_update_error(_stored_model(stored, team_id), incoming) == _missing_limits_message( + missing, "media-latest" + ) + + +def _model_update_error( + existing: Deployment, + incoming: updateLiteLLMParams, + policy: ModelQuotaPolicy = _CUSTOMER_MODEL_ADD_POLICY, + remaining: LiteLLM_Params | None = None, +) -> str | None: + try: + _raise_if_model_update_adds_missing_rate_limits( + existing=existing, + written_name="media-latest", + written_team_model=existing.model_info.team_id is not None, + incoming_params=incoming, + remaining_params=existing.litellm_params if remaining is None else remaining, + policy=policy, + ) + except ProxyException as error: + return error.message + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("setting", "expected_code", "expected_message"), + [ + ({"default": [], "models": {"veo-*": ["rpm"]}}, "400", "rpm must be set to a positive value for 'veo-3.1'"), + ("yes", "400", "rpm and tpm must be set to a positive value for 'veo-3.1'"), + ({"enabled": True}, "400", "rpm and tpm must be set to a positive value for 'veo-3.1'"), + ( + {"providers": {"vertex_ai": ["qps"]}}, + "500", + "general_settings.enforce_rpm_tpm_on_model_add is invalid", + ), + ], +) +async def test_model_add_resolves_the_quota_policy_for_the_public_model_name( + monkeypatch: pytest.MonkeyPatch, setting: object, expected_code: str, expected_message: str +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + prisma: Final = MagicMock() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_rpm_tpm_on_model_add": setting}) + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment(model_name="veo-3.1", litellm_params={"model": "vertex_ai/veo-3.1"}), + UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert error.value.code == expected_code + assert error.value.message.startswith(expected_message) + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_model_add_cannot_borrow_a_model_entry_by_its_public_name(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + prisma: Final = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="media-team")) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr( + proxy_server, + "general_settings", + {"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"jev-latest": []}}}, + ) + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment( + model_name="jev-latest", + litellm_params={"model": "vertex_ai/gemini-3.8-flash"}, + model_info={"team_id": "media-team"}, + ), + UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert error.value.code == "400" + assert error.value.message.startswith("rpm and tpm must be set to a positive value for 'jev-latest'") + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + +_PROJECT_QUOTA_SETTINGS: Final = { + "enforce_project_model_quota": { + "providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}, + "models": {"jev-latest": []}, + } +} +_TEAM_VERTEX: Final = {"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "qa-project"} +_TEAM_RUNWAY: Final = {"model": "gen4_turbo", "custom_llm_provider": "runwayml"} + + +def _media_team_db(public_name: str, project_metadata: Mapping[str, object]) -> MagicMock: + prisma: Final = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="media-team")) + prisma.db.litellm_projecttable.find_many = AsyncMock( + return_value=[ + LiteLLM_ProjectTable( + project_id="p-1", + project_alias="video-gen", + team_id="media-team", + models=[public_name], + metadata=dict(project_metadata), + ), + LiteLLM_ProjectTable( + project_id="p-2", project_alias="chat-only", team_id="media-team", models=["chat-only-model"] + ), + ] + ) + return prisma + + +def _team_model_row(public_name: str, litellm_params: Mapping[str, object]) -> LiteLLM_ProxyModelTable: + return LiteLLM_ProxyModelTable( + model_id="team-row", + model_name="model_name_media-team_team-row", + litellm_params=dict(litellm_params), + model_info={"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name}, + created_by="admin", + updated_by="admin", + ) + + +def _quota_router(public_name: str, proxy_params: Mapping[str, object], *team_params: Mapping[str, object]) -> Router: + return Router( + model_list=[ + {"model_name": public_name, "litellm_params": dict(proxy_params)}, + *( + { + "model_name": "model_name_media-team_team-row", + "litellm_params": dict(params), + "model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name}, + } + for params in team_params + ), + ] + ) + + +def _use_quota_proxy(monkeypatch: pytest.MonkeyPatch, prisma: MagicMock, router: Router) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "general_settings", _PROJECT_QUOTA_SETTINGS) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "store_model_in_db", False) + monkeypatch.setattr(proxy_server, "premium_user", True) + + +_ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +@pytest.mark.asyncio +async def test_team_model_add_rejects_a_route_its_projects_lack_the_quota_for(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + prisma: Final = _media_team_db("jev-latest", {}) + _use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"})) + + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment( + model_name="jev-latest", + litellm_params={**_TEAM_VERTEX, "rpm": 10, "tpm": 1000}, + model_info={"team_id": "media-team"}, + ), + _ADMIN, + ) + + assert error.value.code == "400" + assert error.value.message.startswith( + "projects video-gen list 'jev-latest' without a positive rpm and tpm quota, " + "which 'enforce_project_model_quota' requires once this team model change applies" + ) + prisma.db.litellm_projecttable.find_many.assert_awaited_once_with(where={"team_id": "media-team"}) + + +@pytest.mark.parametrize( + ("litellm_params", "project_metadata"), + [ + pytest.param( + _TEAM_VERTEX, + {"model_rpm_limit": {"jev-latest": 5}, "model_tpm_limit": {"jev-latest": 500}}, + id="projects-carry-the-added-limits", + ), + pytest.param(_TEAM_RUNWAY, {}, id="new-route-is-exempt"), + ], +) +@pytest.mark.asyncio +async def test_team_model_add_goes_through_when_its_projects_meet_the_new_route( + monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], project_metadata: Mapping[str, object] +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + _use_quota_proxy( + monkeypatch, _media_team_db("jev-latest", project_metadata), _quota_router("jev-latest", {"model": "openai/jev-latest"}) + ) + + with pytest.raises(ProxyException) as error: + await add_new_model( + Deployment(model_name="jev-latest", litellm_params=dict(litellm_params), model_info={"team_id": "media-team"}), + _ADMIN, + ) + + assert "STORE_MODEL_IN_DB" in error.value.message + + +@pytest.mark.parametrize("endpoint", ["patch", "put"]) +@pytest.mark.asyncio +async def test_team_model_update_rejects_moving_its_route_where_its_projects_lack_quota( + monkeypatch: pytest.MonkeyPatch, endpoint: str +) -> None: + from litellm.proxy import proxy_server + + runway_by_prefix: Final = {"model": "runwayml/gen4_turbo"} + prisma: Final = _media_team_db("jev-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=_team_model_row("jev-latest", runway_by_prefix) + ) + _use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}, runway_by_prefix)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + moved_to_vertex: Final = updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", vertex_project="qa-project") + + write: Final = ( + patch_model( + model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to_vertex), user_api_key_dict=_ADMIN + ) + if endpoint == "patch" + else update_model( + model_params=updateDeployment(litellm_params=moved_to_vertex, model_info=ModelInfo(id="team-row")), + user_api_key_dict=_ADMIN, + ) + ) + + with pytest.raises(ProxyException) as error: + await write + + assert error.value.code == "400" + assert error.value.message.startswith("projects video-gen list 'jev-latest' without a positive rpm and tpm quota") + + +def _team_route_router(public_name: str, team_params: Mapping[str, object]) -> Router: + return Router( + model_list=[ + {"model_name": "jev-latest", "litellm_params": {"model": "openai/jev-latest"}}, + { + "model_name": "model_name_media-team_team-row", + "litellm_params": dict(team_params), + "model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": public_name}, + }, + ], + model_group_alias={"media": "jev-latest"}, + ) + + +@pytest.mark.parametrize( + ("public_name", "project_model", "team_params", "moved_to"), + [ + pytest.param( + "video-*", + "video-gen-1", + {"model": "runwayml/*"}, + updateLiteLLMParams(model="vertex_ai/*", vertex_project="qa-project"), + id="a-name-the-team-wildcard-serves", + ), + pytest.param( + "jev-latest", + "media", + {"model": "runwayml/gen4_turbo"}, + updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", vertex_project="qa-project"), + id="an-alias-of-the-team-name", + ), + ], +) +@pytest.mark.asyncio +async def test_team_model_patch_rechecks_projects_listing_any_name_it_serves( + monkeypatch: pytest.MonkeyPatch, + public_name: str, + project_model: str, + team_params: Mapping[str, object], + moved_to: updateLiteLLMParams, +) -> None: + from litellm.proxy import proxy_server + + prisma: Final = _media_team_db(project_model, {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row(public_name, team_params)) + _use_quota_proxy(monkeypatch, prisma, _team_route_router(public_name, team_params)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + with pytest.raises(ProxyException) as error: + await patch_model( + model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to), user_api_key_dict=_ADMIN + ) + + assert error.value.code == "400" + assert error.value.message.startswith( + f"projects video-gen list '{project_model}' without a positive rpm and tpm quota" + ) + prisma.db.litellm_proxymodeltable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_update_checks_the_limits_of_the_name_it_renames_to(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=LiteLLM_ProxyModelTable( + model_id="vertex-row", + model_name="jev-latest", + litellm_params={"model": "vertex_ai/gemini-3.8-flash"}, + model_info={"id": "vertex-row"}, + created_by="admin", + updated_by="admin", + ) + ) + _use_quota_proxy(monkeypatch, prisma, Router(model_list=[])) + monkeypatch.setattr( + proxy_server, + "general_settings", + {"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"]}, "models": {"jev-latest": []}}}, + ) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + with pytest.raises(ProxyException) as error: + await update_model( + model_params=updateDeployment( + model_name="gemini-3.8-flash", litellm_params=updateLiteLLMParams(), model_info=ModelInfo(id="vertex-row") + ), + user_api_key_dict=_ADMIN, + ) + + assert error.value.code == "400" + assert error.value.message.startswith("rpm and tpm must be set to a positive value for 'gemini-3.8-flash'") + prisma.db.litellm_proxymodeltable.update.assert_not_called() + + +@pytest.mark.parametrize( + ("incoming", "remaining"), + [ + pytest.param(updateLiteLLMParams(rpm=0), {"model": "vertex_ai/gemini-3.8-flash", "tpm": 1000}, id="zeroes-rpm"), + pytest.param(updateLiteLLMParams(), {"model": "vertex_ai/gemini-3.8-flash", "tpm": 1000}, id="removes-rpm"), + pytest.param(updateLiteLLMParams(model="runwayml/gen4_turbo"), {"model": "runwayml/gen4_turbo"}, id="moves"), + ], +) +def test_model_update_under_the_boolean_setting_never_rejects_as_on_main( + incoming: updateLiteLLMParams, remaining: Mapping[str, object] +) -> None: + from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting + + stored: Final = _stored_model({"model": "vertex_ai/gemini-3.8-flash", "rpm": 10, "tpm": 1000}, None) + policy: Final = parse_model_quota_setting(True) + + assert policy is not None + assert _model_update_error(stored, incoming, policy, LiteLLM_Params(**remaining)) is None + + +@pytest.mark.asyncio +async def test_team_model_patch_rejects_a_removal_that_moves_its_route_where_its_projects_lack_quota( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest + + prisma: Final = _media_team_db("jev-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("jev-latest", _TEAM_RUNWAY)) + _use_quota_proxy(monkeypatch, prisma, _quota_router("jev-latest", {"model": "openai/jev-latest"}, _TEAM_RUNWAY)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + with pytest.raises(ProxyException) as error: + await patch_model( + model_id="team-row", + patch_data=PatchModelRequest(litellm_params_to_remove=("custom_llm_provider",)), + user_api_key_dict=_ADMIN, + ) + + assert error.value.code == "400" + assert error.value.message.startswith("projects video-gen list 'jev-latest' without a positive rpm and tpm quota") + + +@pytest.mark.asyncio +async def test_team_model_delete_rejects_falling_back_to_a_route_its_projects_lack_quota_for( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model + + prisma: Final = _media_team_db("media-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY)) + _use_quota_proxy(monkeypatch, prisma, _quota_router("media-latest", _TEAM_VERTEX, _TEAM_RUNWAY)) + + with pytest.raises(ProxyException) as error: + await delete_model(model_info=ModelInfoDelete(id="team-row"), user_api_key_dict=_ADMIN) + + assert error.value.code == "400" + assert error.value.message.startswith("projects video-gen list 'media-latest' without a positive rpm and tpm quota") + prisma.db.litellm_proxymodeltable.delete.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_model_delete_goes_through_when_nothing_else_serves_its_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model + + prisma: Final = _media_team_db("media-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY)) + only_the_team_model: Final = Router( + model_list=[ + { + "model_name": "model_name_media-team_team-row", + "litellm_params": dict(_TEAM_RUNWAY), + "model_info": {"id": "team-row", "team_id": "media-team", "team_public_model_name": "media-latest"}, + } + ] + ) + _use_quota_proxy(monkeypatch, prisma, only_the_team_model) + + with pytest.raises(ProxyException) as error: + await delete_model(model_info=ModelInfoDelete(id="team-row"), user_api_key_dict=_ADMIN) + + assert "STORE_MODEL_IN_DB" in str(error.value.message) + + +@pytest.mark.parametrize("endpoint", ["patch", "put"]) +@pytest.mark.asyncio +async def test_team_model_update_rejects_moving_to_a_provider_that_needs_limits_it_lacks( + monkeypatch: pytest.MonkeyPatch, endpoint: str +) -> None: + from litellm.proxy import proxy_server + + prisma: Final = _media_team_db("media-latest", {}) + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=_team_model_row("media-latest", _TEAM_RUNWAY)) + _use_quota_proxy(monkeypatch, prisma, _quota_router("media-latest", _TEAM_RUNWAY, _TEAM_RUNWAY)) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr( + proxy_server, + "general_settings", + {"enforce_rpm_tpm_on_model_add": {"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []}}}, + ) + moved_to_vertex: Final = updateLiteLLMParams(model="vertex_ai/gemini-3.8-flash", custom_llm_provider="vertex_ai") + + write: Final = ( + patch_model( + model_id="team-row", patch_data=updateDeployment(litellm_params=moved_to_vertex), user_api_key_dict=_ADMIN + ) + if endpoint == "patch" + else update_model( + model_params=updateDeployment(litellm_params=moved_to_vertex, model_info=ModelInfo(id="team-row")), + user_api_key_dict=_ADMIN, + ) + ) + + with pytest.raises(ProxyException) as error: + await write + + assert error.value.code == "400" + assert error.value.message == _missing_limits_message("rpm and tpm", "media-latest") + prisma.db.litellm_proxymodeltable.update.assert_not_called() + class TestBlockModelResponseSerialization: @pytest.mark.parametrize( @@ -7860,6 +8475,76 @@ class TestTeamMemberAutoRouterWrites: "jev_classifier_config" if legacy else "opensource_classifier_config": classifier, } + @pytest.mark.asyncio + @pytest.mark.parametrize("endpoint", ["create", "patch", "legacy"]) + @pytest.mark.parametrize("legacy", [False, True]) + @pytest.mark.parametrize("field", ["api_key", "api_base"]) + async def test_team_admin_classifier_writes_reject_environment_references( + self, endpoint: str, legacy: bool, field: str + ) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + stored_config: Final = self._classifier_config( + {"api_key": "os.environ/PRIVATE_JEV_VALUE", "api_base": "https://trusted-jev.test"}, legacy + ) + row: Final = self._row().model_copy(update={"litellm_params": { + "model": "auto_router/complexity_router", "complexity_router_config": stored_config, + }}) + team: Final = self._team().model_copy(update={"members_with_roles": [Member(user_id="owner", role="admin")]}) + database: Final = self._database(team, row) + actor: Final = UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER) + config: Final = self._classifier_config( + {"api_key": "caller-key", "api_base": "https://untrusted-jev.test", field: "os.environ/PRIVATE_JEV_VALUE"}, + legacy, + ) + request: Final = updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=config), + model_info=ModelInfo(id=row.model_id, team_id="member-team"), + ) + operation: Final = ( + add_new_model(Deployment( + model_name="new-jev-router", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=config), + model_info=ModelInfo(id=row.model_id, team_id="member-team"), + ), actor) + if endpoint == "create" + else patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor) + ) + with self._environment(database, row), pytest.raises((HTTPException, ProxyException)) as denied: + await operation + assert (denied.value.status_code if isinstance(denied.value, HTTPException) else int(denied.value.code)) == 403 + database.db.litellm_proxymodeltable.create.assert_not_awaited() + database.db.litellm_proxymodeltable.update.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) + @pytest.mark.parametrize("change_config", [False, True]) + async def test_team_admin_unrelated_updates_preserve_stored_environment_references( + self, endpoint: str, change_config: bool + ) -> None: + connection: Final = {"api_key": "os.environ/SAVED_JEV_KEY", "api_base": "os.environ/SAVED_JEV_BASE"} + stored_config: Final = self._classifier_config(connection, True) + row: Final = self._row().model_copy(update={"litellm_params": { + "model": "auto_router/complexity_router", "complexity_router_config": stored_config, + }}) + team: Final = self._team().model_copy(update={"members_with_roles": [Member(user_id="owner", role="admin")]}) + database: Final = self._database(team, row) + actor: Final = UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER) + params: Final = ( + {"complexity_router_config": self._classifier_config({"timeout_ms": 8100}, False)} + if change_config else {"max_tokens": 100} + ) + request: Final = updateDeployment( + litellm_params=updateLiteLLMParams.model_validate(params), model_info=ModelInfo(id=row.model_id), + ) + with self._environment(database, row): + await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor)) + written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"] + saved: Final = json.loads(written["litellm_params"])["complexity_router_config"] + assert saved == ( + self._classifier_config({**connection, "timeout_ms": 8100}, False) if change_config else stored_config + ) + @pytest.mark.asyncio @pytest.mark.parametrize( "stored_params,removed,denied_param", diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index ccd5f7f2fea..78c37184d6c 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -9209,6 +9209,119 @@ class TestMCPMetaTraceCarrier: assert _mcp_meta_trace_carrier(SimpleNamespace(meta=only_progress)) is None +@pytest.mark.asyncio +@pytest.mark.parametrize("initialized_excluded", [False, True]) +async def test_stateful_mcp_tracing_follows_each_message_exclusion( + monkeypatch: pytest.MonkeyPatch, initialized_excluded: bool +) -> None: + import httpx2 + from fastapi import FastAPI + from mcp.server.context import HandlerResult + from mcp.server.lowlevel import Server + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + from mcp.shared import _otel + from mcp.types import CallToolRequestParams, NotificationParams + from mcp_types import CLIENT_CAPABILITIES_META_KEY, PROTOCOL_VERSION_META_KEY + from opentelemetry.instrumentation.propagators import get_global_response_propagator, set_global_response_propagator + + from litellm.integrations.otel.model.config import is_otel_v2_enabled + from litellm.integrations.otel.mount import instrument_fastapi_app + from litellm.integrations.otel.plumbing.providers import in_memory_provider + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server import server as gateway + + original_propagator: Final = get_global_response_propagator() + provider, exporter = in_memory_provider() + tracer: Final = provider.get_tracer("mcp-exclusion-test") + monkeypatch.setattr(_otel, "_tracer", tracer) + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS", "excluded.example") + monkeypatch.setitem(proxy_server.general_settings, "mcp_advertised_versions", ["2025-11-25", "2026-07-28"]) + is_otel_v2_enabled.cache_clear() + notified: Final = asyncio.Event() + + async def call_tool( + ctx: ServerRequestContext[object, object], params: CallToolRequestParams + ) -> CallToolResult: + with tracer.start_as_current_span("tool-handler"): + return CallToolResult(content=[TextContent(type="text", text="done")]) + + async def notification( + ctx: ServerRequestContext[object, object], params: NotificationParams + ) -> HandlerResult: + with tracer.start_as_current_span("notification-handler"): + notified.set() + return None + + protocol: Final = Server[object]("trace-exclusion", on_call_tool=call_tool) + protocol.middleware[:] = gateway.server.middleware + protocol.add_notification_handler("notifications/exclusion-probe", NotificationParams, notification) + manager: Final = StreamableHTTPSessionManager(protocol, stateless=False) + app: Final = FastAPI() + app.mount("/mcp", manager.handle_request) + instrument_fastapi_app(app, tracer_provider=provider) + initial_url: Final = f"http://{'excluded' if initialized_excluded else 'included'}.example/mcp/" + current_url: Final = f"http://{'included' if initialized_excluded else 'excluded'}.example/mcp/" + headers: Final = {"Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2025-11-25"} + try: + async with manager.run(): + async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client: + initialized: Final = await client.post(initial_url, headers=headers, json={ + "jsonrpc": "2.0", "id": 0, "method": "initialize", "params": { + "protocolVersion": "2025-11-25", "capabilities": {}, + "clientInfo": {"name": "trace-exclusion", "version": "1"}, + }, + }) + assert initialized.status_code == 200, initialized.text + assert any( + span.name == "initialize" for span in exporter.get_finished_spans() + ) is not initialized_excluded + exporter.clear() + session_headers: Final = {**headers, "MCP-Session-Id": initialized.headers["mcp-session-id"]} + ready: Final = await client.post(initial_url, headers=session_headers, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}) + assert ready.status_code == 202, ready.text + response: Final = await client.post(current_url, headers=session_headers, json={ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": { + "name": "probe", "arguments": {}, "_meta": { + "traceparent": "00-11111111111111111111111111111111-2222222222222222-01", + }, + }, + }) + assert response.status_code == 200 and '"text":"done"' in response.text, response.text + names: Final = {span.name for span in exporter.get_finished_spans()} + assert ("tools/call probe" in names) is initialized_excluded + assert ("tool-handler" in names) is initialized_excluded + exporter.clear() + notification_response: Final = await client.post(current_url, headers=session_headers, + json={"jsonrpc": "2.0", "method": "notifications/exclusion-probe"}) + assert notification_response.status_code == 202, notification_response.text + await notified.wait() + assert any( + span.name == "notification-handler" for span in exporter.get_finished_spans() + ) is initialized_excluded + exporter.clear() + modern_response: Final = await client.post(current_url, headers={ + "Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2026-07-28", + "MCP-Method": "tools/call", "MCP-Name": "probe", + }, json={ + "jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": { + "name": "probe", "arguments": {}, "_meta": { + PROTOCOL_VERSION_META_KEY: "2026-07-28", CLIENT_CAPABILITIES_META_KEY: {}, + }, + }, + }) + assert modern_response.status_code == 200, modern_response.text + assert '"text":"done"' in modern_response.text, modern_response.text + modern_names: Final = {span.name for span in exporter.get_finished_spans()} + assert ("tools/call probe" in modern_names) is initialized_excluded + assert ("tool-handler" in modern_names) is initialized_excluded + finally: + provider.shutdown() + set_global_response_propagator(original_propagator) + is_otel_v2_enabled.cache_clear() + + @pytest.mark.asyncio async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations(_mcp_request_ctx) -> None: from types import SimpleNamespace diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 60043daa402..43fe5d4e786 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2142,7 +2142,6 @@ class TestBedrockPassthroughAwsErrors: assert exc_info.value.aws_error_type == "ValidationException" assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) - @pytest.mark.asyncio async def test_bedrock_request_check_is_raised_with_the_aws_error_type_it_carries( self, monkeypatch: pytest.MonkeyPatch @@ -2183,6 +2182,7 @@ class TestBedrockPassthroughAwsErrors: await self._answer(error) assert (exc_info.value.code, exc_info.value.headers.get("x-model-group")) == ("403", "modèle") + @pytest.mark.asyncio @pytest.mark.parametrize( ("error", "message", "code"), @@ -7250,6 +7250,189 @@ async def test_bedrock_count_tokens_error_forwards_provider_headers(): assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" +_BEDROCK_COUNT_ALIAS: Final = "count-log-bedrock" +_BEDROCK_COUNT_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + + +async def _bedrock_native_count( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response +) -> dict[str, object]: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + litellm.Router( + model_list=[ + { + "model_name": _BEDROCK_COUNT_ALIAS, + "litellm_params": { + "model": _BEDROCK_COUNT_MODEL, + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"id": "count-log-bedrock-deployment"}, + } + ] + ), + ) + respx_mock.post(host="bedrock-runtime.us-east-1.amazonaws.com", path__regex=r"/count-tokens$").mock( + return_value=aws_reply + ) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/bedrock/model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + "headers": [(b"x-litellm-call-id", b"bedrock-native-count")], + } + ) + return await handle_bedrock_count_tokens( + endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]}, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_logs_the_count_against_the_router_model( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + result: Final = await _bedrock_native_count( + monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}) + ) + + assert result == {"input_tokens": 12} + assert len(respx_mock.calls) == 1 + assert respx_mock.calls.last.request.url.path == "/model/anthropic.claude-haiku-4-5-20251001-v1:0/count-tokens" + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-count", + "request_model": _BEDROCK_COUNT_ALIAS, + "model_id": "count-log-bedrock-deployment", + "model_used": "anthropic.claude-haiku-4-5-20251001-v1:0", + "input_tokens": 12, + "count_source": "provider", + "status_code": 200, + }, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_count_fails( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + rejected: Final = httpx.Response(400, json={"message": "The provided model doesn't support counting tokens."}) + with pytest.raises(HTTPException) as exc_info: + await _bedrock_native_count(monkeypatch, respx_mock, rejected) + + assert exc_info.value.status_code == 400 + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-count", + "request_model": _BEDROCK_COUNT_ALIAS, + "model_id": "count-log-bedrock-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_without_a_model_logs_the_400_it_answers( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/bedrock/model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + "headers": [(b"x-litellm-call-id", b"bedrock-native-no-model")], + } + ) + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"messages": [{"role": "user", "content": "hi"}]}, + ) + + assert exc_info.value.status_code == 400 + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-no-model", + "request_model": None, + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_with_a_region_only_model_answers_the_validation_error_and_logs_it( + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/bedrock/model/us/count-tokens", + "headers": [(b"x-litellm-call-id", b"bedrock-native-region-only")], + } + ) + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint="model/us/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"model": "us"}, + ) + + assert exc_info.value.detail == { + "error": "CountTokens processing error: Request must contain content to count tokens" + } + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-region-only", + "request_model": "us", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": exc_info.value.status_code, + }, + ) + class _AzureGroupRouter: def __init__(self, captured: list[dict]) -> None: self.captured = captured diff --git a/tests/unit/proxy/pass_through_endpoints/test_upstream_usage_headers.py b/tests/unit/proxy/pass_through_endpoints/test_upstream_usage_headers.py index 9dd0233ad8b..33328e418cd 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_upstream_usage_headers.py +++ b/tests/unit/proxy/pass_through_endpoints/test_upstream_usage_headers.py @@ -4,10 +4,12 @@ import httpx import pytest from litellm.constants import MAX_SPEND_LOG_MODEL_NAME_LENGTH +from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.upstream_usage_headers import ( UpstreamReportedUsage, apply_upstream_reported_model, + apply_upstream_reported_operation, apply_upstream_reported_usage, has_upstream_reported_usage, parse_upstream_reported_model, @@ -211,3 +213,24 @@ def test_model_header_alone_does_not_make_the_upstream_own_the_cost(): assert has_upstream_reported_usage(logging_obj) is False assert "response_cost" not in logging_obj.model_call_details + + +@pytest.mark.parametrize("operation", [GenAIOperation.RETRIEVAL, GenAIOperation.EMBEDDINGS]) +def test_apply_operation_records_the_operation_the_upstream_named(operation: GenAIOperation) -> None: + logging_obj: Final = _logging_obj() + + apply_upstream_reported_operation(logging_obj, _headers(**{"x-litellm-operation-name": operation.value})) + + assert logging_obj.model_call_details[PASSTHROUGH_OPERATION_KEY] is operation + + +@pytest.mark.parametrize( + "headers", + [{}, {"x-litellm-operation-name": "summarize"}, {"x-litellm-operation-name": "Retrieval"}], +) +def test_apply_operation_ignores_a_missing_or_unknown_name(headers: dict[str, str]) -> None: + logging_obj: Final = _logging_obj() + + apply_upstream_reported_operation(logging_obj, _headers(**headers)) + + assert PASSTHROUGH_OPERATION_KEY not in logging_obj.model_call_details diff --git a/tests/unit/proxy/proxy_server/test_exception_handlers.py b/tests/unit/proxy/proxy_server/test_exception_handlers.py index f6efdecc884..b34c2d2fb08 100644 --- a/tests/unit/proxy/proxy_server/test_exception_handlers.py +++ b/tests/unit/proxy/proxy_server/test_exception_handlers.py @@ -488,22 +488,17 @@ async def test_otel_unhandled_exception_handler_reraises_http_exception_invalid( @pytest.mark.asyncio @pytest.mark.parametrize("media_type", ["application/json", "application/x-protobuf"]) @pytest.mark.parametrize("root_path", ["", "/tenant-a"]) -@pytest.mark.parametrize("native_available", [True, False]) @pytest.mark.parametrize("error", [ ProxyException("database credentials: secret", "auth_error", None, 401), HTTPException(403, "database credentials: secret"), ]) -async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native( - media_type: str, root_path: str, native_available: bool, - error: ProxyException | HTTPException, monkeypatch: pytest.MonkeyPatch, +async def test_otlp_auth_errors_hide_internal_details( + media_type: str, root_path: str, error: ProxyException | HTTPException, ) -> None: from google.rpc.status_pb2 import Status from litellm.proxy.proxy_server import otlp_http_exception_handler - from litellm.rust_bridge import loader - if not native_available: - monkeypatch.setattr(loader, "_cached_bridge", None) request: Final = Request({ "type": "http", "method": "POST", "path": root_path + "/v1/traces", "root_path": root_path, "headers": [(b"content-type", media_type.encode())], @@ -522,7 +517,7 @@ async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native else Status.FromString(response.body).message ) expected: Final = "Unauthorized" if isinstance(error, ProxyException) else "Forbidden" - assert message == (expected if native_available or media_type == "application/json" else "") + assert message == expected def _websocket(sent: list[Message]) -> WebSocket: diff --git a/tests/unit/proxy/proxy_server/test_routes_config.py b/tests/unit/proxy/proxy_server/test_routes_config.py index 7d0a7f3f261..e3b8aff8fcb 100644 --- a/tests/unit/proxy/proxy_server/test_routes_config.py +++ b/tests/unit/proxy/proxy_server/test_routes_config.py @@ -716,12 +716,16 @@ def test_config_list_exposes_config_reload_interval(client, auth_as, mock_prisma assert entry["field_default_value"] == 30 -@pytest.mark.parametrize("enabled", (True, False)) +@pytest.mark.parametrize( + ("sent", "enabled"), + ((True, True), (False, False), ("true", True), ("yes", True), (1, True), ("off", False)), +) def test_config_list_exposes_and_saves_wildcard_routes_toggle( client: TestClient, auth_as: Callable[..., contextlib.AbstractContextManager[object]], mock_prisma: MagicMock, monkeypatch: pytest.MonkeyPatch, + sent: object, enabled: bool, ) -> None: from pydantic import TypeAdapter @@ -747,13 +751,11 @@ def test_config_list_exposes_and_saves_wildcard_routes_toggle( saved: Final = client.post( "/config/field/update", - json={"field_name": "return_wildcard_routes", "field_value": enabled, "config_type": "general_settings"}, + json={"field_name": "return_wildcard_routes", "field_value": sent, "config_type": "general_settings"}, ) assert saved.status_code == 200, saved.text - persisted: Final = TypeAdapter(dict[str, bool]).validate_json( - table.upsert.call_args.kwargs["data"]["create"]["param_value"] - ) - assert persisted == {"return_wildcard_routes": enabled} + persisted: Final = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted["return_wildcard_routes"] is enabled after: Final = client.get("/config/list", params={"config_type": "general_settings"}) assert after.status_code == 200, after.text diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index f60c2254b41..d55ce25d878 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -363,6 +363,26 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_zerogpu_add_model_form_defaults_resolve_to_a_priced_zerogpu_model(): + import litellm + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + app_instance: Final = FastAPI() + app_instance.include_router(router) + providers: Final = TestClient(app_instance).get("/public/providers/fields").json() + + zerogpu: Final = next(p for p in providers if p["litellm_provider"] == "zerogpu") + fields_by_key: Final = {f["key"]: f for f in zerogpu["credential_fields"]} + placeholder: Final = zerogpu["default_model_placeholder"] + _, provider, _, api_base = get_llm_provider(model=placeholder) + + assert provider == "zerogpu" + assert api_base == fields_by_key["api_base"]["placeholder"] + assert litellm.model_cost[placeholder]["litellm_provider"] == "zerogpu" + assert fields_by_key["api_key"]["required"] is True + assert fields_by_key["api_key"]["field_type"] == "password" + + def test_qwen_mainland_provider_fields_carry_the_qianwen_brand(): app_instance = FastAPI() app_instance.include_router(router) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 985aee061bc..7760b6b7fcc 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -7,7 +7,7 @@ import contextlib import dataclasses import json import unittest -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call, patch @@ -19,6 +19,7 @@ from fastapi.testclient import TestClient from httpx import Response import litellm +from litellm.constants import EMPTY_MAPPING from litellm.proxy._types import ProxyException from litellm.proxy.proxy_server import app from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse @@ -4275,6 +4276,7 @@ class TestResponsesInputTokens: llm_router: litellm.Router | None = None, response_id_security: object | None = None, managed_files: object | None = None, + headers: Mapping[str, str] = EMPTY_MAPPING, ) -> tuple[Response, AsyncMock]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -4307,7 +4309,7 @@ class TestResponsesInputTokens: app.dependency_overrides[_proxy_managed_files_hook] = lambda: managed_files try: client = TestClient(app) - response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876"}) + response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876", **headers}) return response, token_counter_mock finally: app.dependency_overrides.pop(user_api_key_auth, None) @@ -4552,6 +4554,84 @@ class TestResponsesInputTokens: assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text assert sent == [{"model": "gpt-4o-mini", "input": "hi"}] + @pytest.mark.parametrize( + "path", ["/v1/responses/input_tokens", "/responses/input_tokens", "/openai/v1/responses/input_tokens"] + ) + def test_count_logs_one_line_with_the_callers_call_id_model_and_count( + self, + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], + path: str, + ) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.llms.openai.responses.count_tokens import handler as openai_count_tokens_http + from litellm.proxy.proxy_server import count_request_tokens + + gpt: Final = {"model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://count.test/v1"} + router: Final = litellm.Router( + model_list=[ + { + "model_name": "count-log-gpt", + "litellm_params": gpt, + "model_info": {"id": "count-log-responses-deployment"}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr( + openai_count_tokens_http, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, json={"object": "response.input_tokens", "input_tokens": 7}, request=request + ) + ) + ), + ) + + response, _ = self._post_input_tokens( + {"model": "count-log-gpt", "input": "hi"}, + path=path, + counter=AsyncMock(side_effect=count_request_tokens), + llm_router=router, + headers={"x-litellm-call-id": "count-log-responses"}, + ) + + assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-responses", + "request_model": "count-log-gpt", + "model_id": "count-log-responses-deployment", + "model_used": "gpt-4o-mini", + "input_tokens": 7, + "count_source": "provider", + "status_code": 200, + }, + ) + + def test_count_rejected_before_counting_logs_the_400_it_answers( + self, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] + ) -> None: + response, counter = self._post_input_tokens( + {"model": "count-log-gpt"}, headers={"x-litellm-call-id": "count-log-no-input"} + ) + + assert response.status_code == 400, response.text + counter.assert_not_awaited() + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-no-input", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + def test_a_managed_file_the_key_cannot_read_answers_403_without_a_count(self) -> None: from litellm.integrations.custom_logger import CustomLogger 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_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index a32c40e5a43..81a196af676 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -19,6 +19,7 @@ from starlette.datastructures import Headers import litellm from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, + COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER, @@ -3684,6 +3685,110 @@ async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypa assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]["requester_metadata"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "surface,headers,body,child_id,trace_id", + [ + ( + "chat/completions", {"X-Session-Id": "child-01"}, {"metadata": {"session_id": "body-001"}}, + "child-01", "parent-01", + ), + ("responses", {"x-vendor-session-id": "child-01"}, {}, "child-01", "parent-01"), + ("messages", {"session-id": "child-01", "user-agent": "codex_cli_rs/1.0"}, {}, "child-01", "parent-01"), + ("chat/completions", {"x-litellm-session-id": "short", "x-session-id": "child-01"}, {}, "short", "short"), + ("responses", {"x-litellm-trace-id": "trace-01", "x-session-id": "child-01"}, {}, "child-01", "trace-01"), + ( + "messages", {"x-litellm-trace-id": "trace-01", "x-litellm-session-id": "child-01"}, + {}, "child-01", "trace-01", + ), + ("responses", {}, {"metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), + ("messages", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), + ("chat/completions", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), + ("messages", {"baggage": "session.id=child-01"}, {}, "child-01", "parent-01"), + ( + "responses", {"baggage": "session.id=child-01"}, {"metadata": {"session_id": "body-001"}}, + "body-001", "parent-01", + ), + ( + "chat/completions", {}, {"litellm_session_id": "root-001", "metadata": {"session_id": "body-001"}}, + "root-001", "parent-01", + ), + ( + "responses", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": "active-1"}}, + "active-1", "parent-01", + ), + ( + "chat/completions", {}, + {"metadata": {"session_id": "active-1"}, "litellm_metadata": {"session_id": "body-001"}}, + "active-1", "parent-01", + ), + ( + "chat/completions", {}, {"metadata": {"session_id": ""}, "litellm_metadata": {"session_id": "body-001"}}, + None, "parent-01", + ), + ( + "messages", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": ""}}, + None, "parent-01", + ), + ( + "messages", {}, {"litellm_session_id": "root-001", "litellm_metadata": {"session_id": "active-1"}, + "metadata": {"session_id": "body-001", "user_id": "user_abc_account__session_anthro-1"}}, + "anthro-1", "parent-01", + ), + ("messages", {"x-session-id": "child-01"}, {"metadata": {"user_id": {"session_id": "body-001"}}}, + "child-01", "parent-01"), + ("chat/completions", {}, {}, None, "parent-01"), + ("responses", {"x-session-id": "bad"}, {}, None, "parent-01"), + ("messages", {}, {"metadata": {"session_id": 123}}, None, "parent-01"), + ("chat/completions", {"x-session-id": "parent-01"}, {}, None, "parent-01"), + ], +) +async def test_child_tier_identity_preserves_trace_and_body_session_precedence( + surface: str, headers: dict[str, str], body: dict[str, object], child_id: str | None, trace_id: str +) -> None: + request: Final = _make_request_mock(f"/v1/{surface}", {"x-parent-session-id": "parent-01", **headers}) + updated: Final = await add_litellm_data_to_request( + data={"model": "test-router", "messages": [], **copy.deepcopy(body)}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + metadata: Final = updated[get_metadata_variable_name(request)] + assert metadata[COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY] == child_id + assert metadata["session_id"] == metadata["trace_id"] == trace_id + assert updated["litellm_session_id"] == updated["litellm_trace_id"] == trace_id + expected_caller: Final = ( + headers.get("x-litellm-session-id", "parent-01") if "x-litellm-trace-id" in headers else None + ) + assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == expected_caller + assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata.get("requester_metadata", {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface,parent_id", [("chat/completions", None), ("responses", None), ("messages", "bad")]) +async def test_non_child_requests_cannot_forge_a_child_tier_identity(surface: str, parent_id: str | None) -> None: + headers: Final = {"x-session-id": "child-01", **({"x-parent-session-id": parent_id} if parent_id else {})} + request: Final = _make_request_mock(f"/v1/{surface}", headers) + updated: Final = await add_litellm_data_to_request( + data={ + "model": "test-router", + "messages": [], + "metadata": {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: "forged-1"}, + "litellm_metadata": {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: "forged-2"}, + }, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + for bucket in ("metadata", "litellm_metadata"): + metadata: Final = updated.get(bucket, {}) + assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata + assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata.get("requester_metadata", {}) + assert updated[get_metadata_variable_name(request)]["session_id"] == "child-01" + + def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id(): """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining.""" headers = {"x-litellm-session-id": "bar"} diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index ebe3c7407d8..0e8fc6823bd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -10,7 +10,9 @@ import json import logging 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 @@ -40,7 +42,7 @@ from starlette.types import Message import litellm from litellm import Router -from litellm._logging import verbose_proxy_logger +from litellm._logging import JsonFormatter, verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError @@ -208,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 @@ -364,6 +365,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" @@ -445,6 +447,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" @@ -1264,6 +1267,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -1322,6 +1326,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -1380,6 +1385,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -1679,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, ) @@ -1742,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, ) @@ -1780,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") @@ -1834,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 @@ -2831,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", @@ -2882,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" @@ -3647,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"}, } ] @@ -3827,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) @@ -3897,3 +3987,149 @@ async def test_responses_count_on_claude_counts_the_turns_previous_response_id_r ] assert [sorted(message) for message in messages if isinstance(message, dict)] == [["content", "role"]] * 3 spend_logs.assert_awaited_once_with("chatcmpl-stored", ("chatcmpl-stored", "resp_client_form")) + + +_COUNTED_HELLO: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"model": "count-log-gpt", "messages": [{"role": "user", "content": "Hello, how many tokens is this?"}]} +) + + +def _utils_token_counter( + monkeypatch: pytest.MonkeyPatch, body: Mapping[str, JsonValue], headers: Mapping[str, str] +) -> httpx.Response: + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-gpt", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + "model_info": {"id": "count-log-utils-deployment"}, + } + ] + ), + ) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-test")) + return TestClient(app).post("/utils/token_counter", json=dict(body), headers=dict(headers)) + + +def test_utils_token_counter_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter(monkeypatch, _COUNTED_HELLO, {"x-litellm-call-id": "count-log-utils"}) + + assert response.status_code == 200, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils", + "request_model": "count-log-gpt", + "model_id": "count-log-utils-deployment", + "model_used": "gpt-4o-mini", + "input_tokens": response.json()["total_tokens"], + "count_source": "local_estimate", + "status_code": 200, + }, + ) + + +def test_utils_token_counter_logs_a_generated_call_id_when_the_client_sends_none( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter(monkeypatch, _COUNTED_HELLO, {}) + + assert response.status_code == 200, response.text + (line,) = token_count_log_lines() + assert str(uuid.UUID(str(line["litellm_call_id"]))) == line["litellm_call_id"] + + +def test_utils_token_counter_without_messages_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter( + monkeypatch, {"model": "count-log-gpt"}, {"x-litellm-call-id": "count-log-utils-empty"} + ) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils-empty", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def test_utils_token_counter_failed_count_logs_its_status_in_place_of_a_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter( + monkeypatch, {**_COUNTED_HELLO, "prompt": "and a prompt"}, {"x-litellm-call-id": "count-log-utils-error"} + ) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": "prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", + "type": "invalid_request_error", + "param": "prompt", + "code": "400", + } + } + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils-error", + "request_model": "count-log-gpt", + "model_id": "count-log-utils-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def test_utils_token_counter_json_line_keeps_its_own_fields_when_the_call_id_is_a_dict( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + forged_call_id: Final = "{'component':'forged-comp','logger':'forged.py:1','injected_key':'injected'}" + + _utils_token_counter(monkeypatch, _COUNTED_HELLO, {"x-litellm-call-id": forged_call_id}) + + (line,) = token_count_log_lines() + (record,) = (record for record in caplog.records if "count_source" in record.__dict__) + rendered: Final = json.loads(JsonFormatter().format(record)) + assert line["litellm_call_id"] == forged_call_id + assert {key: rendered.get(key) for key in ("litellm_call_id", "component", "logger", "injected_key")} == { + "litellm_call_id": forged_call_id, + "component": verbose_proxy_logger.name, + "logger": f"{record.filename}:{record.lineno}", + "injected_key": None, + } + + +def test_utils_token_counter_line_stays_one_line_when_the_model_carries_a_newline( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + forged_model: Final = "count-log-gpt\ntoken count - litellm_call_id=forged input_tokens=1" + + _utils_token_counter(monkeypatch, {**_COUNTED_HELLO, "model": forged_model}, {"x-litellm-call-id": "count log"}) + + (line,) = token_count_log_lines() + (message,) = (record.getMessage() for record in caplog.records if "count_source" in record.__dict__) + assert line["request_model"] == forged_model + assert "\n" not in message + assert message.startswith(f'token count - litellm_call_id="count log" request_model={json.dumps(forged_model)} ') 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_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index bf571c74e22..b926b824fc6 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -12,7 +12,7 @@ import math from collections.abc import Awaitable, Callable, Mapping from pathlib import Path from importlib import import_module -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, patch import httpx @@ -254,6 +254,84 @@ async def test_aresponses_forwards_non_enum_reasoning_effort( assert response.output[0].content[0].text == "Done." +@pytest.mark.asyncio +async def test_aresponses_forwards_moderation( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_moderation", "gpt-5.4")) + ) + moderation: Final = {"model": "omni-moderation-latest"} + + await litellm.aresponses(model="openai/gpt-5.4", input="hi", moderation=moderation) + + request_body: Final = TypeAdapter(dict[str, object]).validate_python( + cast(object, json.loads(upstream.calls[0].request.read())) + ) + assert upstream.call_count == 1 + assert request_body.get("moderation") == moderation + + +@pytest.mark.asyncio +async def test_aresponses_omits_null_moderation( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_null_moderation", "gpt-5.4")) + ) + + await litellm.aresponses(model="openai/gpt-5.4", input="hi", moderation=None) + + request_body: Final = TypeAdapter(dict[str, object]).validate_python( + cast(object, json.loads(upstream.calls[0].request.read())) + ) + assert upstream.call_count == 1 + assert "moderation" not in request_body + + +@pytest.mark.asyncio +async def test_aresponses_stream_forwards_moderation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + payload: Final = _minimal_responses_api_payload("resp_stream_moderation", "gpt-5.4") + moderation: Final = {"model": "omni-moderation-latest"} + + def respond(request: httpx.Request) -> httpx.Response: + received.put_nowait(request) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join( + f"event: {kind}\ndata: {json.dumps({'type': kind, 'response': payload})}\n\n" + for kind in ("response.created", "response.completed") + ), + ) + + async_client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + response: Final = await litellm.aresponses( + model="openai/gpt-5.4", input="hi", moderation=moderation, stream=True, client=async_client + ) + events: Final = tuple([event.type async for event in response]) + assert "response.completed" in events + finally: + await async_client.close() + + request_body: Final = TypeAdapter(dict[str, object]).validate_python( + cast(object, json.loads(received.get_nowait().content)) + ) + assert request_body.get("moderation") == moderation + assert received.empty() + + @pytest.mark.asyncio async def test_acompletion_with_tools_forwards_non_enum_reasoning_effort_over_the_bridge( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter 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_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py index 10a8f7929d9..77d465526a4 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -732,6 +732,85 @@ def test_jev_api_base_without_its_own_key_is_rejected_so_the_environment_key_sta assert JevClassifierConfig(api_key="sk-own").api_base is None +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("connection", ["literal", "key-reference", "base-reference", "references", "defaults"]) +async def test_jev_connection_resolves_only_explicit_environment_references( + monkeypatch: pytest.MonkeyPatch, legacy: bool, connection: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("JEV_CONNECTION_KEY", "resolved-key") + monkeypatch.setenv("JEV_CONNECTION_BASE", "https://jev.test") + monkeypatch.setenv("TYPESAFE_API_KEY", "default-key") + monkeypatch.setenv("TYPESAFE_API_BASE", "https://default-jev.test") + monkeypatch.setenv("literal-key", "must-not-replace-literal") + fields: Final = ( + {} + if connection == "defaults" + else { + "api_key": "os.environ/JEV_CONNECTION_KEY" + if connection in ("key-reference", "references") + else "literal-key", + "api_base": "os.environ/JEV_CONNECTION_BASE" + if connection in ("base-reference", "references") + else "https://jev.test", + } + ) + config: Final = { + "classifier_type": "jev" if legacy else "oss_classifier", + "jev_classifier_config" if legacy else "opensource_classifier_config": {"model": "openjev-latest", **fields}, + "tiers": {"SIMPLE": "cheap", "COMPLEX": "strong"}, + } + classifier: Final = ComplexityRouter("env-router", litellm.Router(), config, derive_savings_baseline=False) + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post( + "https://default-jev.test/v1/systemone" if connection == "defaults" else "https://jev.test/v1/systemone" + ).respond(json={"answers": {"tier": _answer("SIMPLE").model_dump()}}) + outcome: Final = await classifier.aclassify("What is two plus two?") + expected_key: Final = ( + "default-key" + if connection == "defaults" + else "resolved-key" if connection in ("key-reference", "references") else "literal-key" + ) + assert outcome.cause == "jev_classifier" + assert route.calls[0].request.headers["authorization"] == f"Bearer {expected_key}" + assert classifier.config.opensource_classifier_config is not None + assert classifier.config.opensource_classifier_config.api_key == fields.get("api_key") + assert classifier.config.opensource_classifier_config.api_base == fields.get("api_base") + + +@pytest.mark.parametrize("field", ["api_key", "api_base"]) +@pytest.mark.parametrize("value", [None, "", " "]) +def test_jev_missing_explicit_environment_reference_never_uses_provider_defaults( + monkeypatch: pytest.MonkeyPatch, field: str, value: str | None +) -> None: + monkeypatch.setenv("TYPESAFE_API_KEY", "never-send-default-key") + monkeypatch.setenv("TYPESAFE_API_BASE", "https://never-use-default.test") + if value is None: + monkeypatch.delenv("JEV_MISSING_CONNECTION", raising=False) + else: + monkeypatch.setenv("JEV_MISSING_CONNECTION", value) + with ( + respx.mock(assert_all_called=False) as upstream, + pytest.raises(ValueError, match=rf"opensource_classifier_config\.{field}"), + ): + ComplexityRouter( + "env-router", + litellm.Router(), + { + "classifier_type": "jev", + "jev_classifier_config": { + "api_key": "literal-key", + "api_base": "https://jev.test", + field: "os.environ/JEV_MISSING_CONNECTION", + }, + "tiers": {"SIMPLE": "cheap"}, + }, + derive_savings_baseline=False, + ) + assert len(upstream.calls) == 0 + + def test_databricks_api_base_without_its_own_key_is_rejected_so_the_workspace_token_stays_home() -> None: with pytest.raises(ValueError, match="DATABRICKS_API_KEY or DATABRICKS_TOKEN is only sent to DATABRICKS_API_BASE"): JevClassifierConfig.model_validate( diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index cb3faf89358..ca39c7f5c8f 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -7714,6 +7714,176 @@ class _DummyPlugin: return context +class TestChildSessionTierPins: + @staticmethod + def _router(session_affinity: bool = False) -> Router: + return Router( + model_list=[ + { + "model_name": "child-tier-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "simple", "REASONING": ["reasoning", "reasoning-peer"]}, + "classification_mode": "user_turn", + "session_affinity": session_affinity, + "keyword_tier_rules": [{"keywords": ["reason carefully"], "tier": "REASONING"}], + }, + }, + }, + *[ + { + "model_name": group, + "litellm_params": {"model": "openai/test-model", "api_key": "test-only"}, + "model_info": {"id": f"{group}-{index}"}, + } + for group, index in ( + ("simple", 0), ("simple", 1), ("reasoning", 0), + ("reasoning", 1), ("reasoning-peer", 0), ("reasoning-peer", 1), + ) + ], + ], + num_retries=0, + ) + + @staticmethod + def _normalized_request( + surface: str, session_id: str | None, caller: str = "caller-a", child: bool = True + ) -> tuple[dict[str, object], str]: + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + metadata_variable_name_for_route, + ) + + route: Final = { + "chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages" + }[surface] + bucket: Final = metadata_variable_name_for_route(route) + data: Final[dict[str, object]] = {bucket: {"user_api_key_hash": caller, "user_api_key_request_route": route}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={ + **({"X-Session-Id": session_id} if session_id is not None else {}), + **({"x-parent-session-id": "parent-session"} if child else {}), + }, + data=data, + _metadata_variable_name=bucket, + ) + return data, bucket + + async def _route( + self, router: Router, surface: str, session_id: str | None, prompt: str | None, + expected_tier: str, expected_cause: str, caller: str = "caller-a", child: bool = True, + expected_deployment: str | None = None, + ) -> str: + from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck + + data, bucket = self._normalized_request(surface, session_id, caller, child) + turn: Final[dict[str, object]] = ( + {"role": "user", "content": prompt} if prompt is not None else + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]} + if surface == "messages" else + {"type": "function_call_output", "call_id": "call_1", "output": "done"} + if surface == "responses" else + {"role": "tool", "tool_call_id": "call_1", "content": "done"} + ) + tool_call: Final = ( + {"type": "function_call", "call_id": "call_1", "name": "read_file", "arguments": "{}"} + if surface == "responses" else + {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}}]} + if surface == "messages" else TestClassificationMode.TOOL_CALL_1 + ) + turns: Final = [turn] if prompt is not None else [TestClassificationMode.SIMPLE_ASK, tool_call, turn] + if surface == "responses": + data["input"] = turns + selected: Final = await router.async_get_available_deployment( + model="child-tier-router", request_kwargs=data, + messages=None if surface == "responses" else turns, + input=turns if surface == "responses" else None, + ) + metadata: Final = data[bucket] + assert isinstance(metadata, dict) + assert (metadata["session_id"], metadata["trace_id"], data["litellm_session_id"], data["litellm_trace_id"]) == ( + "parent-session", "parent-session", "parent-session", "parent-session" + ) + assert (metadata["routing_decision"]["tier"], metadata["routing_decision"]["cause"]) == ( + expected_tier, expected_cause + ) + affinity: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + if expected_deployment is not None: + filtered: Final = await affinity.async_filter_deployments( + model=selected["model_name"], healthy_deployments=router.get_model_list(selected["model_name"]) or [], + messages=None, request_kwargs=data, + ) + assert [deployment["model_info"]["id"] for deployment in filtered] == [expected_deployment] + assert selected["model_info"]["id"] == expected_deployment + await affinity.async_pre_call_deployment_hook( + kwargs={ + **data, "model_info": selected["model_info"], + bucket: {**metadata, "deployment_model_name": selected["model_name"]}, + }, + call_type=None, + ) + return str(selected["model_info"]["id"]) + + @pytest.mark.asyncio + @pytest.mark.parametrize("surface", ["chat", "responses", "messages"]) + async def test_children_keep_their_tier_without_replacing_the_parent(self, surface: str) -> None: + router: Final = self._router() + try: + parent: Final = await self._route( + router, surface, "parent-session", "reason carefully", "REASONING", "literal_keyword_match", child=False + ) + for session, prompt, tier, cause, caller, child in ( + ("child-session-a", "Hello!", "SIMPLE", "heuristic_scorer", "caller-a", True), + ("child-session-a", None, "SIMPLE", "user_turn_continuation", "caller-a", True), + ("parent-session", None, "REASONING", "user_turn_continuation", "caller-a", False), + ("child-session-b", None, "SIMPLE", "heuristic_scorer", "caller-a", True), + ("child-session-b", "reason carefully", "REASONING", "literal_keyword_match", "caller-a", True), + ("child-session-a", None, "SIMPLE", "user_turn_continuation", "caller-a", True), + ("child-session-a", "reason carefully", "REASONING", "literal_keyword_match", "caller-b", True), + ("child-session-a", None, "SIMPLE", "user_turn_continuation", "caller-a", True), + (None, "Hello!", "SIMPLE", "heuristic_scorer", "caller-a", True), + (None, None, "SIMPLE", "heuristic_scorer", "caller-a", True), + ): + await self._route( + router, surface, session, prompt, tier, cause, caller, child, + expected_deployment=parent if tier == "REASONING" and caller == "caller-a" else None, + ) + assert await self._route( + router, surface, "parent-session", None, "REASONING", "user_turn_continuation", + child=False, expected_deployment=parent, + ) == parent + finally: + router.discard() + + @pytest.mark.asyncio + async def test_session_affinity_still_shares_the_parent_tier(self) -> None: + router: Final = self._router(session_affinity=True) + try: + parent: Final = await self._route( + router, "chat", "parent-session", "reason carefully", "REASONING", "literal_keyword_match", child=False + ) + assert await self._route( + router, "chat", "child-session-a", "Hello!", "REASONING", "session_affinity_pin", + expected_deployment=parent, + ) == parent + finally: + router.discard() + + @pytest.mark.asyncio + async def test_child_tier_model_affinity_keeps_using_the_parent_identity(self) -> None: + router: Final = self._router() + try: + strategy: Final = router.complexity_routers["child-tier-router"][0].strategy + parent, _ = self._normalized_request("chat", "parent-session", child=False) + child, _ = self._normalized_request("chat", "child-session-a") + candidates: Final = ("reasoning", "reasoning-peer") + assert await strategy._pin_model_for_tier("REASONING", "reasoning", candidates, parent) == "reasoning" + assert await strategy._pin_model_for_tier("REASONING", "reasoning-peer", candidates, child) == "reasoning" + finally: + router.discard() + + class TestClassificationMode: """Test classification_mode='user_turn': classify only requests whose newest turn is a new human ask; tool-loop continuation turns replay the session's held routing decision.""" diff --git a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py index 8fa1c799863..f79c8a6fadb 100644 --- a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -4,13 +4,24 @@ import litellm from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_KWARG, MID_STREAM_CONTINUATION_MARKER, + RESPONSES_CONTINUATION_KWARG, ContinuationOutput, ContinuationPrefillDeploymentCheck, + ResponsesContinuation, + ResponsesContinuationSeam, + constrains_continuation, + constrains_responses_continuation, continuation_output_ceilings, deployment_supports_prefill, + deployment_takes_responses_prefill, + responses_continuation_input, + responses_continuation_input_for, ) +from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents PREFILL_MODEL = "anthropic/prefill-capable-test-model" +CLAUDE_PREFILL_MODEL = "anthropic/claude-prefill-capable-test-model" +CLAUDE_NON_PREFILL_MODEL = "anthropic/claude-prefill-rejecting-test-model" NON_PREFILL_MODEL = "openai/prefill-unknown-test-model" UNMAPPED_MODEL = "openai/unmapped-test-model" @@ -21,6 +32,12 @@ def _synthetic_cost_map_entries() -> None: { PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}, NON_PREFILL_MODEL: {"litellm_provider": "openai", "mode": "chat"}, + CLAUDE_PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}, + CLAUDE_NON_PREFILL_MODEL: { + "litellm_provider": "anthropic", + "mode": "chat", + "supports_assistant_prefill": False, + }, }, persist_across_reloads=False, ) @@ -191,7 +208,11 @@ def _text_chunk(text: str) -> litellm.ModelResponseStream: ("Hello", (" world",), (" world",)), ("Hello", ("Hel", "p me"), ("Hel", "p me")), ("Hello", ("Hel",), ("Hel",)), - ([{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}], ("Sure: Hello", " world"), ("", " world")), + ( + [{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}], + ("Sure: Hello", " world"), + ("", " world"), + ), ], ids=[ "echo_split_across_chunks_is_dropped", @@ -206,9 +227,7 @@ def _text_chunk(text: str) -> litellm.ModelResponseStream: async def test_continuation_output_drops_only_a_full_echo_of_the_prefill( prefill_content, fallback_texts, expected_texts ): - output = ContinuationOutput( - "Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", () - ) + output = ContinuationOutput("Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", ()) async def fallback_stream(): for text in fallback_texts: @@ -266,7 +285,14 @@ async def test_continuation_output_drops_only_whitespace_the_client_already_has( ({"extra_body": {"max_completion_tokens": 45}}, {"extra_body": {"max_completion_tokens": 15}}), ({"extra_body": {"top_k": 3}}, {}), ], - ids=["max_tokens", "both_ceilings", "no_ceiling", "extra_body_max_tokens", "extra_body_only", "extra_body_no_ceiling"], + ids=[ + "max_tokens", + "both_ceilings", + "no_ceiling", + "extra_body_max_tokens", + "extra_body_only", + "extra_body_no_ceiling", + ], ) def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_the_request_sends( request_kwargs, expected @@ -286,3 +312,217 @@ def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_ ) def test_continuation_output_ceilings_decline_once_a_ceiling_has_no_tokens_left(request_kwargs): assert continuation_output_ceilings(request_kwargs, 30) is None + + +@pytest.mark.parametrize( + "sent,expected", + [ + ({"text": {"format": {"type": "json_schema", "name": "n", "schema": {"type": "object"}}}}, True), + ({"text": {"format": {"type": "json_object"}}}, True), + ({"text": {"format": {"type": "text"}}}, False), + ({"text": {"verbosity": "low"}}, False), + ({"tool_choice": "required"}, True), + ({"tool_choice": "auto"}, False), + ({"tool_choice": {"type": "function", "name": "lookup"}}, True), + ({"tool_choice": {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "f"}]}}, True), + ({"tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "f"}]}}, True), + ({"tool_choice": {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": []}}}, True), + ({"tool_choice": {"type": "web_search_preview"}}, True), + ], + ids=[ + "json_schema", + "json_object", + "text_format", + "verbosity_only", + "required_tool", + "auto_tool", + "forced_function", + "allowed_tools_required", + "allowed_tools_auto", + "chat_allowed_tools_auto", + "hosted_tool", + ], +) +def test_constrains_continuation_reads_the_responses_text_format(sent, expected): + assert constrains_continuation(sent) is expected + + +@pytest.mark.parametrize( + "sent,expected", + [ + ({"tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "f"}]}}, False), + ({"tool_choice": {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "f"}]}}, True), + ({"tool_choice": {"type": "function", "name": "lookup"}}, True), + ({"tool_choice": {"type": "web_search_preview"}}, True), + ({"tool_choice": "required"}, True), + ({"tool_choice": "auto"}, False), + ( + { + "tool_choice": {"type": "allowed_tools", "mode": "auto", "tools": []}, + "text": {"format": {"type": "json_object"}}, + }, + True, + ), + ], + ids=[ + "allowed_tools_auto", + "allowed_tools_required", + "forced_function", + "hosted_tool", + "required_tool", + "auto_tool", + "allowed_tools_auto_with_json", + ], +) +def test_constrains_responses_continuation_lets_allowed_tools_in_auto_mode_answer_in_text(sent, expected): + assert constrains_responses_continuation(sent) is expected + + +@pytest.mark.parametrize( + "deployment,expected", + [ + (_deployment(CLAUDE_PREFILL_MODEL, "a"), True), + (_deployment(CLAUDE_NON_PREFILL_MODEL, "b"), False), + (_deployment(PREFILL_MODEL, "c"), False), + ( + {"litellm_params": {"model": "anthropic/support-model-v1"}, "model_info": {"id": "d", "base_model": CLAUDE_PREFILL_MODEL}}, + True, + ), + ], + ids=["claude_with_prefill", "claude_without_prefill", "prefill_capable_but_not_claude", "claude_named_by_base_model"], +) +def test_deployment_takes_responses_prefill_only_on_claude_with_prefill(deployment, expected): + assert deployment_takes_responses_prefill(deployment) is expected + + +def _continuation_kwargs() -> dict[str, object]: + continuation = ResponsesContinuation("Count to 5.", "1, 2, ") + return {RESPONSES_CONTINUATION_KWARG: continuation, "input": responses_continuation_input(continuation)} + + +@pytest.mark.parametrize( + "prior_input,expected_prior", + [ + ( + "Count to 5.", + [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Count to 5."}]}], + ), + ( + [ + {"role": "user", "content": "Count to 5."}, + {"type": "function_call_output", "call_id": "c", "output": "x"}, + ], + [ + {"role": "user", "content": "Count to 5."}, + {"type": "function_call_output", "call_id": "c", "output": "x"}, + ], + ), + (None, []), + ], + ids=["string_input", "item_list_input", "no_input"], +) +def test_responses_continuation_input_keeps_the_prior_input_then_the_partial_then_a_request_to_continue( + prior_input, expected_prior +): + built = responses_continuation_input(ResponsesContinuation(prior_input, "1, 2, ")) + + partial = {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "1, 2, "}]} + assert built[:-1] == [*expected_prior, partial] + assert built[-1]["role"] == "user" + assert "already ends with whitespace" in built[-1]["content"][0]["text"] + + +def test_responses_continuation_request_asks_for_a_leading_space_only_after_a_word(): + built = responses_continuation_input(ResponsesContinuation("Hi", "Transformers use")) + + assert "begin with a space" in built[-1]["content"][0]["text"] + + +def test_responses_continuation_accumulates_across_hops(): + continuation = ResponsesContinuation("Count to 5.", "1, ").extended("2, ") + + assert continuation == ResponsesContinuation("Count to 5.", "1, 2, ") + + +@pytest.mark.parametrize( + "deployment,expected_roles", + [ + (_deployment(CLAUDE_PREFILL_MODEL, "a"), ["user", "assistant"]), + (_deployment(CLAUDE_NON_PREFILL_MODEL, "b"), ["user", "assistant", "user"]), + (_deployment(NON_PREFILL_MODEL, "c"), ["user", "assistant", "user"]), + ], + ids=["claude_prefill", "claude_without_prefill", "openai"], +) +def test_responses_continuation_input_for_shapes_the_input_for_the_deployment_that_serves_it( + deployment, expected_roles +): + shaped = responses_continuation_input_for(_continuation_kwargs(), deployment) + + assert [item["role"] for item in shaped] == expected_roles + + +def test_responses_continuation_input_for_keeps_what_the_router_stripped_from_the_input(): + continuation = ResponsesContinuation( + [ + {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "opaque"}, + {"role": "user", "content": "Count to 5."}, + ], + "1, 2, ", + ) + stripped_input = responses_continuation_input(continuation)[1:] + + shaped = responses_continuation_input_for( + {RESPONSES_CONTINUATION_KWARG: continuation, "input": stripped_input}, _deployment(CLAUDE_PREFILL_MODEL, "a") + ) + + assert shaped == stripped_input[:-1] + + +def test_responses_continuation_input_for_leaves_an_input_without_its_request_to_continue_alone(): + kwargs = {**_continuation_kwargs(), "input": [{"role": "user", "content": "Count to 5."}]} + + assert responses_continuation_input_for(kwargs, _deployment(CLAUDE_PREFILL_MODEL, "a")) == kwargs["input"] + + +def test_responses_continuation_input_for_leaves_a_request_without_a_continuation_alone(): + kwargs = {"input": "Count to 5."} + + assert responses_continuation_input_for(kwargs, _deployment(CLAUDE_PREFILL_MODEL, "a")) == "Count to 5." + + +def _delta(text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="m", output_index=0, content_index=0, delta=text + ) + + +@pytest.mark.parametrize( + "emitted,fallback_texts,expected_texts", + [ + ("1, 2, ", (" 3, 4", ", 5"), ("3, 4", ", 5")), + ("1, 2, ", (" ", " 3"), ("", " 3")), + ("1, 2,", (" 3",), (" 3",)), + ("1, 2, ", ("3",), ("3",)), + ("A.\n\n", ("\n Hi",), (" Hi",)), + ("", (" Hi",), (" Hi",)), + ], + ids=[ + "re_sent_space_is_dropped", + "seam_split_across_events", + "no_trailing_whitespace_keeps_the_space", + "no_re_sent_space_passes_through", + "seam_ends_at_the_first_divergence", + "fresh_start_passes_through", + ], +) +@pytest.mark.asyncio +async def test_responses_continuation_seam_drops_only_whitespace_the_client_already_has( + emitted, fallback_texts, expected_texts +): + async def fallback_stream(): + for text in fallback_texts: + yield _delta(text) + + released = [event async for event in ResponsesContinuationSeam(emitted).splice(fallback_stream())] + + assert tuple(event.delta for event in released) == expected_texts 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 70b268bf2a2..f38bbcf4404 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from datetime import datetime from importlib import import_module from pathlib import Path -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -44,12 +44,13 @@ from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kw from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.litellm_params import ControlOptions from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.utils import Delta, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, Usage from tests.unit.llms.bedrock.slow_upstream import ( + CONVERSE_ANSWER, STREAM_TIMEOUT_SECONDS, UPSTREAM_FIRST_BYTE_SECONDS, slow_upstream_sync_client, @@ -2126,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", @@ -9779,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" @@ -9950,3 +10306,126 @@ async def test_converse_fallback_prices_the_current_attempt( else: client.close() configuration.reset_rust_configuration() + + +_LOOKUP_FUNCTION: Final = { + "name": "lookup_record", + "description": "Look up a record", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, +} +_STRICT_LOOKUP_TOOL: Final = {"type": "function", "function": {**_LOOKUP_FUNCTION, "strict": True}} +# Bedrock's 400 for a request past Claude's strict tool limits, seen live on us.anthropic.claude-sonnet-4-6 on 2026-10-20 +_TOO_MANY_STRICT_TOOLS: Final = ( + "The model returned the following errors: Too many strict tools (21). The maximum number of strict tools " + "supported is 20. Try reducing the number of tools marked as strict." +) + + +@dataclass(slots=True) +class _StrictToolsLimitRejectionFirst: + stream: Literal["off", "fake", "real"] + sent: tuple[tuple[str, dict[str, object]], ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.sent = (*self.sent, (request.url.path.rsplit("/", 1)[-1], json.loads(request.content))) + if len(self.sent) == 1: + return httpx.Response(400, json={"message": _TOO_MANY_STRICT_TOOLS}) + return httpx.Response(200, content=b"" if self.stream == "real" else CONVERSE_ANSWER) + + +async def _converse_with_a_strict_tools_limit_rejection_first( + asynchronous: bool, + stream: Literal["off", "fake", "real"], + drop_params: bool | None, + tool: Mapping[str, object] = _STRICT_LOOKUP_TOOL, +) -> tuple[tuple[tuple[str, dict[str, object]], ...], litellm.BadRequestError | None]: + respond: Final = _StrictToolsLimitRejectionFirst(stream) + kwargs: Final = { + "model": "bedrock/us.anthropic.claude-sonnet-4-6", + "messages": [{"role": "user", "content": "Find invoice 42"}], + "tools": [tool], + "aws_access_key_id": "AKIASTRICTTOOLS", + "aws_secret_access_key": "strict-tools-secret", + "aws_region_name": "us-east-1", + "num_retries": 0, + "drop_params": drop_params, + "stream": stream != "off", + "fake_stream": stream == "fake", + } + try: + chunks: Final = ( + await _acompletion_chunks(AsyncHTTPHandler(transport=httpx.MockTransport(respond)), kwargs, stream) + if asynchronous + else _completion_chunks( + HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), kwargs, stream + ) + ) + except litellm.BadRequestError as error: + return respond.sent, error + assert chunks + return respond.sent, None + + +async def _acompletion_chunks( + client: AsyncHTTPHandler, kwargs: Mapping[str, object], stream: Literal["off", "fake", "real"] +) -> tuple[object, ...]: + response: Final = await litellm.acompletion(client=client, **kwargs) + return tuple([chunk async for chunk in response]) if stream == "fake" else (response,) + + +def _completion_chunks( + client: HTTPHandler, kwargs: Mapping[str, object], stream: Literal["off", "fake", "real"] +) -> tuple[object, ...]: + response: Final = litellm.completion(client=client, **kwargs) + return tuple(response) if stream == "fake" else (response,) + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("stream", ["off", "fake", "real"], ids=["completion", "fake_stream", "stream"]) +async def test_bedrock_strict_tools_limit_rejection_is_sent_again_without_strict_under_drop_params( + asynchronous: bool, stream: Literal["off", "fake", "real"], local_model_cost_map: None +) -> None: + sent, error = await _converse_with_a_strict_tools_limit_rejection_first(asynchronous, stream, drop_params=True) + + assert error is None + (first_path, first_body), (second_path, second_body) = sent + first_tool_spec: Final = first_body["toolConfig"]["tools"][0]["toolSpec"] + assert first_tool_spec["strict"] is True + assert first_path == second_path == ("converse-stream" if stream == "real" else "converse") + assert second_body == { + **first_body, + "toolConfig": { + **first_body["toolConfig"], + "tools": [{"toolSpec": {key: value for key, value in first_tool_spec.items() if key != "strict"}}], + }, + } + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_bedrock_strict_tools_limit_rejection_reaches_the_caller_without_drop_params( + asynchronous: bool, local_model_cost_map: None +) -> None: + sent, error = await _converse_with_a_strict_tools_limit_rejection_first(asynchronous, stream="off", drop_params=None) + + assert error is not None and "Too many strict tools" in str(error) + assert len(sent) == 1 + + +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_bedrock_strict_tools_limit_rejection_is_not_sent_again_when_no_tool_sent_strict( + asynchronous: bool, local_model_cost_map: None +) -> None: + sent, error = await _converse_with_a_strict_tools_limit_rejection_first( + asynchronous, + stream="off", + drop_params=True, + tool={"type": "function", "function": {**_LOOKUP_FUNCTION, "strict": False}}, + ) + + assert error is not None and "Too many strict tools" in str(error) + assert len(sent) == 1 diff --git a/tests/unit/test_mistral_large_4_model_metadata.py b/tests/unit/test_mistral_large_4_model_metadata.py index f45c53920dd..bfcd6690890 100644 --- a/tests/unit/test_mistral_large_4_model_metadata.py +++ b/tests/unit/test_mistral_large_4_model_metadata.py @@ -2,6 +2,7 @@ import json from pathlib import Path +from typing import Final import pytest @@ -32,9 +33,7 @@ def test_metadata_and_backup(model): assert info["litellm_provider"] == "mistral" assert info["mode"] == "chat" assert info["max_input_tokens"] == 1048576 - assert info["input_cost_per_token"] == 6.8e-07 - assert info["output_cost_per_token"] == 2.09e-06 - assert info["cache_read_input_token_cost"] == 6.8e-08 + assert 0 < info["cache_read_input_token_cost"] < info["input_cost_per_token"] < info["output_cost_per_token"] assert info["reasoning_effort_levels"] == ["none", "high"] for capability in ( "supports_vision", @@ -59,12 +58,16 @@ def test_reasoning_effort(model, effort): @pytest.mark.parametrize("model", MODELS) def test_completion_cost(model): + info: Final = litellm.get_model_info(model) + input_price: Final[float] = info["input_cost_per_token"] + output_price: Final[float] = info["output_cost_per_token"] + cache_read_price: Final[float] = info["cache_read_input_token_cost"] uncached_response = litellm.ModelResponse( model=model.split("/", 1)[1], usage={"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, ) assert litellm.completion_cost(completion_response=uncached_response, model=model) == pytest.approx( - 1000 * 6.8e-07 + 100 * 2.09e-06 + 1000 * input_price + 100 * output_price ) response = litellm.ModelResponse( model=model.split("/", 1)[1], @@ -76,7 +79,7 @@ def test_completion_cost(model): }, ) assert litellm.completion_cost(completion_response=response, model=model) == pytest.approx( - 400 * 6.8e-07 + 600 * 6.8e-08 + 100 * 2.09e-06 + 400 * input_price + 600 * cache_read_price + 100 * output_price ) diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 08a39b94c12..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 { @@ -224,6 +241,41 @@ def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict): ) +# Azure OpenAI web_search runs on Grounding with Bing and bills each Bing request reported in +# tool_usage.web_search.num_requests: https://learn.microsoft.com/en-us/azure/foundry/openai/how-to/web-search +# Grounding with Bing lists $14 per 1,000 transactions: https://www.microsoft.com/en-us/bing/apis/grounding-pricing +AZURE_WEB_SEARCH_COST_PER_REQUEST: Final = 0.014 +_AZURE_DEPLOYMENT_PREFIX: Final = re.compile(r"^azure/(us/|eu/|global/|global-standard/)?") + + +def _azure_rows(prices: Mapping[str, object]) -> Mapping[str, dict]: + return MappingProxyType( + {key: row for key, row in prices.items() if key.startswith("azure/") and isinstance(row, dict)} + ) + + +def test_azure_web_search_rows_charge_the_grounding_with_bing_price(prices: dict): + wrong = { + key: row["search_context_cost_per_query"] + for key, row in _azure_rows(prices).items() + if row.get("search_context_cost_per_query") + and set(row["search_context_cost_per_query"].values()) != {AZURE_WEB_SEARCH_COST_PER_REQUEST} + } + assert wrong == {} + + +def test_azure_rows_price_web_search_when_their_openai_model_does(prices: dict): + missing = sorted( + key + for key, row in _azure_rows(prices).items() + if not row.get("search_context_cost_per_query") + and row.get("mode") in ("chat", "responses", None) + and isinstance(twin := prices.get(_AZURE_DEPLOYMENT_PREFIX.sub("", key)), dict) + and twin.get("search_context_cost_per_query") + ) + assert missing == [] + + PRIORITY_SUFFIX: Final = "_priority" MIN_PRIORITY_MARKUP: Final = 1.0 MAX_PRIORITY_MARKUP: Final = 5.0 @@ -845,7 +897,10 @@ def test_volcengine_doubao_seed_2_1_pro_260915_row_matches_published_metadata(pa assert entry["cache_read_input_token_cost"] == pytest.approx(1.2 / 6.9565 / 1_000_000, rel=1e-4) sibling = rows["volcengine/doubao-seed-2-1-pro-260628"] assert isinstance(sibling, dict) - assert all(entry[key] == sibling[key] for key in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost")) + assert all( + entry[key] == sibling[key] + for key in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost") + ) def _bedrock_gpt_oss_rows(prices: Mapping[str, JsonValue]) -> tuple[str, ...]: diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 30d675e130d..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,10 +7163,310 @@ 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: input is rewritten to include user prompt + - developer instruction + prior assistant message with partial output.""" + """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" from litellm.exceptions import MidStreamFallbackError router = _make_router_with_fallback() @@ -7076,11 +7502,11 @@ async def test_aresponses_streaming_iterator_partial_content_injects_continuatio assert isinstance(new_input, list) assert new_input[0]["role"] == "user" assert new_input[0]["content"][0]["text"] == "What's the capital of France?" - assert new_input[1]["role"] == "developer" - assert "do not repeat" in new_input[1]["content"][0]["text"].lower() - assert new_input[2]["role"] == "assistant" - assert new_input[2]["content"][0]["type"] == "output_text" - assert new_input[2]["content"][0]["text"] == "The capital of France is" + assert new_input[1]["role"] == "assistant" + assert new_input[1]["content"][0]["type"] == "output_text" + assert new_input[1]["content"][0]["text"] == "The capital of France is" + assert new_input[2]["role"] == "user" + assert "do not repeat" in new_input[2]["content"][0]["text"].lower() @pytest.mark.asyncio @@ -17297,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( @@ -17309,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=[ { @@ -17320,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( @@ -17336,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=[ { @@ -17352,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 @@ -20330,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 @@ -21617,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", @@ -27634,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( @@ -29495,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_router_per_deployment_num_retries.py b/tests/unit/test_router_per_deployment_num_retries.py index 99ad7c224f8..b08ae25c438 100644 --- a/tests/unit/test_router_per_deployment_num_retries.py +++ b/tests/unit/test_router_per_deployment_num_retries.py @@ -6,7 +6,10 @@ GitHub Issue: #18968 - Per-deployment max_retries/num_retries in litellm_params import httpx import pytest import pytest_asyncio -from unittest.mock import patch +from queue import SimpleQueue +from typing import Final, Literal +from unittest.mock import AsyncMock, patch +from pydantic import JsonValue, TypeAdapter import litellm from litellm import Router @@ -14,6 +17,70 @@ from litellm.types.router import RetryPolicy from litellm.integrations.custom_logger import CustomLogger +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("source", ("deployment", "default", "request")) +@pytest.mark.parametrize(("sdk_retries", "router_retries"), ((0, 0), (0, 1), (5, 1))) +async def test_routed_images_do_not_multiply_sdk_and_router_retries( + monkeypatch: pytest.MonkeyPatch, + is_async: bool, + source: Literal["deployment", "default", "request"], + sdk_retries: int, + router_retries: int, +) -> None: + requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue() + + def respond(request: httpx.Request) -> httpx.Response: + requests.put(request) + assert request.url.path == "/v1/images/generations" + body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content) + assert body == {"model": "gpt-image-1-mini", "prompt": "a red circle"} + return httpx.Response( + 503, + headers={"retry-after": "0"}, + json={"error": {"message": "scripted unavailable", "type": "server_error"}}, + ) + + transport: Final = httpx.MockTransport(respond) + litellm.in_memory_llm_clients_cache.flush_cache() + with httpx.Client(transport=transport) as sync_session: + async with httpx.AsyncClient(transport=transport) as async_session: + monkeypatch.setattr(litellm, "client_session", sync_session) + monkeypatch.setattr(litellm, "aclient_session", async_session) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "num_retries", None) + router: Final = Router( + model_list=[ + { + "model_name": "image-retries", + "litellm_params": { + "model": "openai/gpt-image-1-mini", + "api_key": "synthetic-image-key", + "api_base": "https://router-image.test/v1", + **({"max_retries": sdk_retries} if source == "deployment" else {}), + }, + } + ], + default_litellm_params={"max_retries": sdk_retries} if source == "default" else {}, + num_retries=router_retries, + disable_cooldowns=True, + ) + call_params: Final = {"max_retries": sdk_retries} if source == "request" else {} + try: + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + if is_async: + with pytest.raises(litellm.ServiceUnavailableError): + await router.aimage_generation(model="image-retries", prompt="a red circle", **call_params) + else: + with pytest.raises(litellm.ServiceUnavailableError): + router.image_generation(model="image-retries", prompt="a red circle", **call_params) + finally: + litellm.in_memory_llm_clients_cache.flush_cache() + assert requests.qsize() == router_retries + 1 + + class TestPerDeploymentNumRetries: """Test that per-deployment num_retries in litellm_params is correctly used.""" 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_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/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index b59aa051e38..ac13af14a6b 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -37,6 +37,7 @@ import { } from "./tier_rows"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; +import { routingModelGroups } from "./auto_setup"; import ClassificationMethodConfig, { InactiveHeuristicV2Threshold } from "./ClassificationMethodConfig"; import ComplexityRouterAdvancedSections from "./ComplexityRouterAdvancedSections"; import { type TierSetAction, applyTierSetAction, setFallbackTier } from "./tier_set_actions"; @@ -657,13 +658,10 @@ const ComplexityRouterConfig: React.FC = ({ ); const classifierEffortOptionsByModel = classifierEffortOptionsForModels(modelInfo); - // Embedding models can't serve a chat-completion role, so they're excluded here. - const modelOptions = modelInfo - .filter((model) => model.mode !== "embedding") - .map((model) => ({ - value: model.model_group, - label: model.model_group, - })); + const modelOptions = routingModelGroups(modelInfo).map((model) => ({ + value: model.model_group, + label: model.model_group, + })); const handleTierModelParamChange = (tier: string, model: string, change: TierModelParamChange) => onChange({ diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx index 52c4e9093a4..78782e755ac 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx @@ -319,6 +319,27 @@ describe("AddAutoRouterTab", () => { }, ); + it("carries Responses model choices from simple setup into Advanced tiers", async () => { + const models = ["gpt-6-luna", "gpt-6-astra"]; + mockFetchAvailableModels.mockResolvedValue(models.map((model_group) => ({ model_group, mode: "responses" }))); + mockFetchAllModelDeployments.mockResolvedValue( + models.map((model_name) => ({ + model_name, + litellm_params: { model: `chatgpt/${model_name}` }, + model_info: { mode: "responses" }, + })), + ); + renderWithProviders(); + await waitFor(() => + expect(screen.getByRole("combobox", { name: "Everyday model" })).toHaveTextContent("gpt-6-luna"), + ); + expect(screen.getByRole("combobox", { name: "Hard tasks model" })).toHaveTextContent("gpt-6-astra"); + expect(screen.queryByText(/Two deployed models are needed/)).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Advanced configuration", exact: true })); + expectTierModel("Simple", "gpt-6-luna"); + expectTierModel("Reasoning", "gpt-6-astra"); + }); + it("keeps manual solver choices on goal changes and carries them into Advanced", async () => { mockFetchAvailableModels.mockResolvedValue( ["claude-haiku-5-5", "custom-fast", "claude-opus-5-5"].map((model_group) => ({ diff --git a/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts b/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts index 4097afc17fe..c0d4931873f 100644 --- a/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts @@ -7,6 +7,7 @@ import { solverPairConfig, nextRouterName, simpleSetupError, + routingModelGroups, } from "./auto_setup"; const deployment = (name: string, model: string, extra: Partial = {}): AutoRouterDeployment => ({ @@ -76,7 +77,7 @@ describe("deployedRouterModels", () => { ]); }); - it("uses only caller-visible chat models and excludes blocked and recursive targets", () => { + it("matches Advanced choices while keeping blocked and recursive targets out of suggestions", () => { const rows = [ deployment("hidden", "anthropic/claude-opus-6"), deployment("blocked", "anthropic/claude-opus-5", { model_info: { blocked: true } }), @@ -94,9 +95,40 @@ describe("deployedRouterModels", () => { ], rows, ).map((model) => model.name), - ).toEqual(["ready"]); + ).toEqual(["blocked", "ready", "recursive"]); + expect( + recommendSolverPair( + deployedRouterModels( + rows + .filter((row) => row.model_name !== "hidden" && row.model_name !== "embedding") + .map((row) => ({ model_group: row.model_name! })), + rows, + ), + "balanced", + "auto", + ), + ).toBeNull(); }); + it.each(["chat", "responses", "completion", "unknown", undefined])( + "offers the same routing choices as Advanced for mode %s", + (mode) => { + const groups = [ + { model_group: "everyday", mode }, + { model_group: "hard", mode }, + { model_group: "embedding", mode: "embedding" }, + ]; + const rows = [ + deployment("everyday", "chatgpt/gpt-6-luna", { model_info: { mode } }), + deployment("hard", "chatgpt/gpt-6-astra", { model_info: { mode } }), + ]; + const models = deployedRouterModels(groups, rows); + expect(models.map((model) => model.name)).toEqual(routingModelGroups(groups).map((group) => group.model_group)); + expect(models.map((model) => model.name)).toEqual(["everyday", "hard"]); + expect(names(recommendSolverPair(models, "balanced", "openai"))).toEqual(["everyday", "hard"]); + }, + ); + it("does not claim a mixed-version or unknown model group is a known version", () => { const rows = [ deployment("mixed", "anthropic/claude-opus-5"), @@ -112,7 +144,7 @@ describe("deployedRouterModels", () => { ], rows, ); - expect(models.map((model) => model.name)).toEqual(["claude-sonnet-5", "mixed"]); + expect(models.map((model) => model.name)).toEqual(["claude-sonnet-5", "mixed", "wildcard"]); expect(models.every((model) => model.family === undefined)).toBe(true); expect(recommendSolverPair(models, "balanced", "auto")).toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/add_model/auto_setup.ts b/ui/litellm-dashboard/src/components/add_model/auto_setup.ts index 39476470928..eb21e17c265 100644 --- a/ui/litellm-dashboard/src/components/add_model/auto_setup.ts +++ b/ui/litellm-dashboard/src/components/add_model/auto_setup.ts @@ -146,24 +146,21 @@ const compareVersion = (left: readonly number[], right: readonly number[]): numb const compareRelease = (left: ModelIdentity, right: ModelIdentity): number => compareVersion(left.version, right.version) || left.snapshot - right.snapshot; +export const routingModelGroups = (groups: readonly ModelGroup[]): ModelGroup[] => + groups.filter((group) => group.mode !== "embedding"); + export const deployedRouterModels = ( groups: readonly ModelGroup[], deployments: readonly AutoRouterDeployment[], ): SolverModel[] => - groups - .flatMap((group) => { - if (group.mode && group.mode !== "chat") return []; + routingModelGroups(groups) + .map((group) => { const rows = deployments.filter((row) => row.model_name === group.model_group); - if ( - rows.some(isAutoRouterDeployment) || - group.model_group.includes("*") || - group.model_group.startsWith("auto_router/") - ) - return []; const active = rows.filter((row) => !row.model_info?.blocked); - if (rows.length > 0 && active.length === 0) return []; - if (active.some((row) => row.model_info?.mode && row.model_info.mode !== "chat")) return []; - if (active.some((row) => row.litellm_params?.model?.includes("*"))) return []; + const isRouter = rows.some(isAutoRouterDeployment) || group.model_group.startsWith("auto_router/"); + const hasWildcard = + group.model_group.includes("*") || active.some((row) => row.litellm_params?.model?.includes("*")); + const manualOnly = isRouter || hasWildcard || (rows.length > 0 && active.length === 0); const identities = active.length > 0 ? active.map(identityOf) : [parseModel(group.model_group)]; const first = identities[0]; const homogeneous = @@ -173,19 +170,17 @@ export const deployedRouterModels = ( const sameLine = identity.family === first.family && identity.line === first.line; return sameLine && compareRelease(identity, first) === 0; }); - const identity = homogeneous ? first : undefined; + const identity = homogeneous && !manualOnly ? first : undefined; const providers = active.length > 0 ? active.map((row) => providerOf(row, group, identity?.family)) : [providerOf({}, group, identity?.family)]; - return [ - { - ...identity, - name: group.model_group, - model: identity?.model ?? group.model_group, - provider: [...new Set(providers)].toSorted().join(", "), - }, - ]; + return { + ...identity, + name: group.model_group, + model: identity?.model ?? group.model_group, + provider: [...new Set(providers)].toSorted().join(", "), + }; }) .toSorted((a, b) => a.name.localeCompare(b.name)); 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/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index 47a1e3608c5..f9c5e3b0181 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -490,6 +490,137 @@ describe("responses_api", () => { }); }); +describe("responses_api reasoning", () => { + type Run = { reasoning: string; answer: string; timeToFirstToken: number[]; reasoningBeforeFirstToken: string }; + + const run = async (streamingEnabled: boolean): Promise => { + const reasoningChunks: string[] = []; + const answerChunks: string[] = []; + const timeToFirstToken: number[] = []; + const reasoningAtFirstToken: string[] = []; + await makeOpenAIResponsesRequest( + [{ role: "user", content: "What is 3+4?" }], + (_role, delta) => answerChunks.push(delta), + "qwen3", + "test-token", + undefined, + undefined, + (content) => reasoningChunks.push(content), + (ttft) => { + timeToFirstToken.push(ttft); + reasoningAtFirstToken.push(reasoningChunks.join("") + "|" + answerChunks.join("")); + }, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + streamingEnabled, + ); + return { + reasoning: reasoningChunks.join(""), + answer: answerChunks.join(""), + timeToFirstToken, + reasoningBeforeFirstToken: reasoningAtFirstToken[0] ?? "", + }; + }; + + const streamOf = (events: Record[]) => { + async function* mockStream() { + for (const event of events) { + yield event; + } + yield { type: "response.completed", response: { id: "resp_reasoning" } }; + } + mockResponsesCreate.mockResolvedValue(mockStream()); + }; + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("shows streamed reasoning_text deltas and starts the first-token timer on them", async () => { + streamOf([ + { type: "response.reasoning_text.delta", item_id: "rs_1", content_index: 0, delta: "Add " }, + { type: "response.reasoning_text.delta", item_id: "rs_1", content_index: 0, delta: "them" }, + { type: "response.output_text.delta", delta: "7" }, + ]); + + const result = await run(true); + + expect(result.reasoning).toBe("Add them"); + expect(result.answer).toBe("7"); + expect(result.timeToFirstToken).toHaveLength(1); + expect(result.reasoningBeforeFirstToken).toBe("Add |"); + }); + + it("separates streamed reasoning summary parts and items with a blank line", async () => { + streamOf([ + { type: "response.reasoning_summary_text.delta", item_id: "rs_1", summary_index: 0, delta: "First" }, + { type: "response.reasoning_summary_text.delta", item_id: "rs_1", summary_index: 0, delta: " part" }, + { type: "response.reasoning_summary_text.delta", item_id: "rs_1", summary_index: 1, delta: "Second" }, + { type: "response.reasoning_summary_text.delta", item_id: "rs_2", summary_index: 0, delta: "Third" }, + ]); + + await expect(run(true)).resolves.toMatchObject({ reasoning: "First part\n\nSecond\n\nThird" }); + }); + + it("still shows the legacy reasoning.delta event", async () => { + streamOf([{ type: "response.reasoning.delta", delta: "Legacy thought" }]); + + await expect(run(true)).resolves.toMatchObject({ reasoning: "Legacy thought" }); + }); + + it("shows non-streamed reasoning content when the summary is empty", async () => { + mockResponsesCreate.mockReturnValueOnce( + nonStreamingResponse({ + id: "resp_content", + output: [ + { + type: "reasoning", + summary: [], + content: [ + { type: "reasoning_text", text: "Add" }, + { type: "reasoning_text", text: "them" }, + ], + }, + { type: "message", content: [{ type: "output_text", text: "7" }] }, + ], + }), + ); + + await expect(run(false)).resolves.toMatchObject({ reasoning: "Add\n\nthem", answer: "7" }); + }); + + it("prefers the non-streamed reasoning summary over its content", async () => { + mockResponsesCreate.mockReturnValueOnce( + nonStreamingResponse({ + id: "resp_summary", + output: [ + { + type: "reasoning", + summary: [{ type: "summary_text", text: "Short summary" }], + content: [{ type: "reasoning_text", text: "Raw chain" }], + }, + { type: "reasoning", content: [{ type: "output_text", text: "Second item" }] }, + ], + }), + ); + + await expect(run(false)).resolves.toMatchObject({ reasoning: "Short summary\n\nSecond item" }); + }); +}); + describe("responses_api prompt cache usage", () => { const captureUsage = async (usage: Record): Promise => { async function* mockStream() { diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index 426abd4e172..ce1a8aab0c5 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -44,6 +44,22 @@ type SynthesizedResponseEvent = | { type: "response.output_text.delta"; delta: string } | { type: "response.completed"; response: NonStreamedResponse }; +const REASONING_PART_SEPARATOR = "\n\n"; + +const REASONING_DELTA_EVENT_TYPES: ReadonlySet = new Set([ + "response.reasoning_text.delta", + "response.reasoning_summary_text.delta", + "response.reasoning.delta", +]); + +const nonEmptyTexts = (parts: ResponseOutputPart[] | undefined): string[] => + (parts ?? []).map((part) => part.text ?? "").filter((text) => text.length > 0); + +const reasoningItemText = (item: ResponseOutputItem): string => { + const summary = nonEmptyTexts(item.summary); + return (summary.length > 0 ? summary : nonEmptyTexts(item.content)).join(REASONING_PART_SEPARATOR); +}; + const responseAsEvents = (response: NonStreamedResponse): SynthesizedResponseEvent[] => { const outputItems = response.output ?? []; const outputText = outputItems @@ -54,9 +70,9 @@ const responseAsEvents = (response: NonStreamedResponse): SynthesizedResponseEve .join(""); const reasoningText = outputItems .filter((item) => item.type === "reasoning") - .flatMap((item) => item.summary ?? []) - .map((part) => part.text ?? "") - .join(""); + .map(reasoningItemText) + .filter((text) => text.length > 0) + .join(REASONING_PART_SEPARATOR); return [ ...outputItems.map((item) => ({ type: "response.output_item.done" as const, item })), @@ -124,6 +140,7 @@ export async function makeOpenAIResponsesRequest( try { const startTime = Date.now(); let firstTokenReceived = false; + let reasoningPartKey: string | null = null; let servedFromResponseCache = false; const formattedInput = messages.flatMap((message) => { @@ -298,10 +315,16 @@ export async function makeOpenAIResponsesRequest( } // Handle reasoning content - if (event.type === "response.reasoning.delta" && "delta" in event) { - const delta = event.delta; - if (typeof delta === "string" && onReasoningContent) { - onReasoningContent(delta); + if (REASONING_DELTA_EVENT_TYPES.has(event.type) && typeof event.delta === "string" && event.delta.length > 0) { + const partKey = `${event.item_id ?? ""}:${event.summary_index ?? event.content_index ?? 0}`; + const separator = reasoningPartKey !== null && reasoningPartKey !== partKey ? REASONING_PART_SEPARATOR : ""; + reasoningPartKey = partKey; + onReasoningContent?.(separator + event.delta); + + const isFirstToken = !firstTokenReceived; + firstTokenReceived = true; + if (isFirstToken && onTimingData && streamingEnabled) { + onTimingData(Date.now() - startTime); } } diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 0ebd2cdff17..83e80911cff 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -224,6 +224,7 @@ describe("provider_info_helpers", () => { Providers.Sail, Providers.WANDB, Providers.ZAI, + Providers.ZeroGPU, ]; const logolessProviders = Object.values(Providers).filter((provider) => !providerLogoMap[provider]); expect([...logolessProviders].sort()).toEqual([...knownLogolessProviders].sort()); @@ -357,6 +358,10 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.Cognition)).toBe("cognition/swe-1.7"); }); + it("should return zerogpu/gpt-oss-120b placeholder for ZeroGPU provider", () => { + expect(getPlaceholder("ZeroGPU")).toBe("zerogpu/gpt-oss-120b"); + }); + it("should return a chatgpt/ placeholder for the CHATGPT dropdown key", () => { expect(getPlaceholder("CHATGPT")).toBe("chatgpt/gpt-5.4"); }); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 2541aa45050..bfb578b218a 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -197,6 +197,7 @@ export enum Providers { xAI = "xAI", XINFERENCE = "Xinference", ZAI = "Z.AI (Zhipu AI)", + ZeroGPU = "ZeroGPU", } export const provider_map: Record = { @@ -325,6 +326,7 @@ export const provider_map: Record = { xAI: "xai", XINFERENCE: "xinference", ZAI: "zai", + ZeroGPU: "zerogpu", }; const standaloneSubproviderSlugs = new Set(["bedrock_mantle"]); @@ -497,6 +499,7 @@ const providerPlaceholderMap: Partial> = { [Providers.Voyage]: "voyage/", [Providers.WATSONX]: "watsonx/ibm/granite-3-3-8b-instruct", [Providers.ZAI]: "zai/glm-4.5", + [Providers.ZeroGPU]: "zerogpu/gpt-oss-120b", }; export const getPlaceholder = (selectedProvider: string): string => { diff --git a/ui/litellm-dashboard/src/components/shared/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 8af76e4d16f..45f9630983f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -19058,6 +19058,7 @@ export interface paths { * @description Args: * request: TokenCountRequest * call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs. + * litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs * * Returns: * TokenCountResponse @@ -23365,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": [ @@ -44163,6 +44165,11 @@ export interface components { */ type: "tool_search_output"; }; + /** ResponsesAPIModeration */ + ResponsesAPIModeration: { + /** Model */ + model: string; + }; /** * ResponsesAPIRequestParams * @description TypedDict for request parameters supported by the responses API. @@ -44188,6 +44195,7 @@ export interface components { } | null; /** Model */ model: string; + moderation?: components["schemas"]["ResponsesAPIModeration"] | null; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Partial Images */ @@ -66975,6 +66983,7 @@ export interface operations { } | null; /** Model */ model: string; + moderation?: components["schemas"]["ResponsesAPIModeration"] | null; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Partial Images */ @@ -70439,6 +70448,7 @@ export interface operations { } | null; /** Model */ model: string; + moderation?: components["schemas"]["ResponsesAPIModeration"] | null; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Partial Images */ @@ -76765,7 +76775,9 @@ export interface operations { query?: { call_endpoint?: boolean; }; - header?: never; + header?: { + "x-litellm-call-id"?: string | null; + }; path?: never; cookie?: never; }; @@ -81237,6 +81249,7 @@ export interface operations { } | null; /** Model */ model: string; + moderation?: components["schemas"]["ResponsesAPIModeration"] | null; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Partial Images */ diff --git a/uv.lock b/uv.lock index d38e96882d0..b40490112cf 100644 --- a/uv.lock +++ b/uv.lock @@ -1077,11 +1077,30 @@ wheels = [ name = "bytecode" version = "0.17.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.14.*'", + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", + "python_full_version < '3.11'", +] sdist = { url = "https://files.pythonhosted.org/packages/98/c4/4818b392104bd426171fc2ce9c79c8edb4019ba6505747626d0f7107766c/bytecode-0.17.0.tar.gz", hash = "sha256:0c37efa5bd158b1b873f530cceea2c645611d55bd2dc2a4758b09f185749b6fd", size = 105863, upload-time = "2025-09-03T19:55:45.703Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ce/80/379e685099841f8501a19fb58b496512ef432331fed38276c3938ab09d8e/bytecode-0.17.0-py3-none-any.whl", hash = "sha256:64fb10cde1db7ef5cc39bd414ecebd54ba3b40e1c4cf8121ca5e72f170916ff8", size = 43045, upload-time = "2025-09-03T19:55:43.879Z" }, ] +[[package]] +name = "bytecode" +version = "0.19.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.15'", +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/a0/f4499eea31339ba709fad8ff778c6ad9afa2d87434b9771e0b37a1ab6fb1/bytecode-0.19.1.tar.gz", hash = "sha256:f59a6a51a980037f64196448504d12fb04937ae8dcb980943265e799af2943ae", size = 106830, upload-time = "2026-10-07T12:01:11.178Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/d2/7b937ec5bcc2fabca017e32541a0b1ef081bb5a1bb3531d382300f7a1ccc/bytecode-0.19.1-py3-none-any.whl", hash = "sha256:211a5091abd6f16f5bb3d34f1b88fc77f8f8539801a02384b51a97bdcee45f0f", size = 43473, upload-time = "2026-10-07T12:01:09.866Z" }, +] + [[package]] name = "cachetools" version = "7.1.1" @@ -1817,64 +1836,71 @@ wheels = [ [[package]] name = "ddtrace" -version = "4.11.0" +version = "4.15.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "bytecode" }, + { name = "bytecode", version = "0.17.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.15'" }, + { name = "bytecode", version = "0.19.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.15'" }, + { name = "ddtrace-internal" }, { name = "envier" }, { name = "opentelemetry-api" }, { name = "wrapt" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/81/51/a628f0177274bab5b67c93a1558fd222babb15286a427e4e8c1d65a265b2/ddtrace-4.11.0.tar.gz", hash = "sha256:260c5b46e80565f4fd08cec2650f707627fb57cd6f8951a2d8c9e6a02b490074", size = 2422158, upload-time = "2026-07-10T08:57:53.396Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/94/f0f005257dc49a7062c6fed98a6dbcca7691e9d1db5c64e8a2d14f204cac/ddtrace-4.15.6.tar.gz", hash = "sha256:2914f646d4ae338c644b814c4ad80c963e7aadfb39d1e67b0367a15d5c4bc0db", size = 2655968, upload-time = "2026-10-08T09:46:58.392Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/16/1f/05dc12819ca3a3000028a522cab0025c5dd969e97dc472ac38a7d2f1ea71/ddtrace-4.11.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:0bc55e48da62b6a287b36e01f5876489023178c09b60e8e01608e2e895ccdc4b", size = 7103767, upload-time = "2026-07-10T08:55:10.348Z" }, - { url = "https://files.pythonhosted.org/packages/e8/8a/efcea448bc15b4d95c48c9b476f2ddc5862a53cd4baf58fe50c64763f9d8/ddtrace-4.11.0-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:5df47873c625993e81fabd8f6ae2ab61f66c6805cfd1a8c551a509d5b5d91062", size = 7447104, upload-time = "2026-07-10T08:55:15.794Z" }, - { url = "https://files.pythonhosted.org/packages/39/28/89aa239e47d28900ab394a0b8e0a5549cab631c0e5977016ad9d578eaee7/ddtrace-4.11.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e857428cc86fe541f5b3768d7d8331e2f1cbde1a2492deedff4619109cbfd5fc", size = 8511839, upload-time = "2026-07-10T08:55:17.609Z" }, - { url = "https://files.pythonhosted.org/packages/30/43/37a5475ade911abae592e121e852922745d842e16294c549cba4ad0f3dbe/ddtrace-4.11.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc60575f8f47af58e6523f8a802d1f92a4912aa7d4ba71024958421044cba61a", size = 8733381, upload-time = "2026-07-10T08:55:19.66Z" }, - { url = "https://files.pythonhosted.org/packages/10/df/d61f71236a405480c8809522219b184cc486056c32f9cc8bf07cc91e2be9/ddtrace-4.11.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:d65ab60bfb0b2edde6c564bd1f97b99a8ea9bb082c93b265488ba214c606d240", size = 9513521, upload-time = "2026-07-10T08:55:21.921Z" }, - { url = "https://files.pythonhosted.org/packages/79/4f/bff08af986df419179b86fc9c73732836531c2d08ab8dcfe982a6f2e8cb8/ddtrace-4.11.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:7ebfdee1342b5b1cdad8505e040bbfe5bdf6976e820803eaa4a1814417d6bb0e", size = 9788008, upload-time = "2026-07-10T08:55:24.111Z" }, - { url = "https://files.pythonhosted.org/packages/af/b8/2b1fd5b0374edf33ad31adb746d1d1358a44c39a45b1c2a0c4b09651d81e/ddtrace-4.11.0-cp310-cp310-win32.whl", hash = "sha256:6247a431ae3a2622a0e835357a644b1bac81719e28f21748fd7d4034c79c3ed7", size = 5739013, upload-time = "2026-07-10T08:55:27.028Z" }, - { url = "https://files.pythonhosted.org/packages/f2/f5/9bdff59bdce1f80592e70633dc26678c762fff79188c9fb6004c06967708/ddtrace-4.11.0-cp310-cp310-win_amd64.whl", hash = "sha256:9197110982756a97243906ffce7b98e6007c8258a6bbc22451d7b26e8f1b0103", size = 6326099, upload-time = "2026-07-10T08:55:28.962Z" }, - { url = "https://files.pythonhosted.org/packages/76/36/ad0abcc97b63d821ada82cc7452c3070a71f9546f0aa81b50eda1f8ffd41/ddtrace-4.11.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:b4074e538e708b0e9d7f812be15db306abbe3a6cb70046f9f786d66daa862344", size = 7105089, upload-time = "2026-07-10T08:55:31.006Z" }, - { url = "https://files.pythonhosted.org/packages/a6/3a/c0f8f70883a70ec9d4050900823ff7b52ad54dc6b2a3d29561727825410f/ddtrace-4.11.0-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:05f701229233ad8a4ad5d52b8ff2bfd6e2f99601a3bc04367615cd1ce4a01a80", size = 7448237, upload-time = "2026-07-10T08:55:32.962Z" }, - { url = "https://files.pythonhosted.org/packages/4e/46/18a42afab21cd718467e7d7ae3e80898569d9969d731b771b5a8ada10294/ddtrace-4.11.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c3a557891e91a443e6a3a9b7c31dc96b9963acada0986dab5a3d65507efa16a1", size = 8516433, upload-time = "2026-07-10T08:55:35.011Z" }, - { url = "https://files.pythonhosted.org/packages/da/b2/6367a75518f843872e13c062ac3c329f880b2f4d62ff4de0624bdcf64719/ddtrace-4.11.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b0328bfeb4a6bdc00168a4184aea670eedc073734f6fbc6ca1abfeb6547415d2", size = 8740643, upload-time = "2026-07-10T08:55:37.743Z" }, - { url = "https://files.pythonhosted.org/packages/ae/e8/04f1d7ee4737444d45f81545ce21b01bd9d653a994997ab1e7e04b6ee947/ddtrace-4.11.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:ad6f17111746568a58da89b6db5e5d18eb07db7b8ce3be3396631f37391cff50", size = 9518626, upload-time = "2026-07-10T08:55:40.248Z" }, - { url = "https://files.pythonhosted.org/packages/ed/8b/314ba8a81e562a31655328061afb30d26d46ba04aca4dabd7da85c93a9e7/ddtrace-4.11.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:eb30dd327789ec1851bc42194c8db0b93a7ad708c2963d694e3622b684a085f7", size = 9795392, upload-time = "2026-07-10T08:55:42.609Z" }, - { url = "https://files.pythonhosted.org/packages/9e/6e/a580d53c05a771bc53d33a88174e983146fdf879fe5d4f52e51ceb56ee99/ddtrace-4.11.0-cp311-cp311-win32.whl", hash = "sha256:dfe73130778ca22652ae398fc77de3643d50687deb127bf4cd2566b10778125a", size = 5737732, upload-time = "2026-07-10T08:55:45.605Z" }, - { url = "https://files.pythonhosted.org/packages/12/83/094f5b88690f7801129b0d52a7b6a58a95b34ecbfcb337555d323c639faa/ddtrace-4.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:1241f38d52b40b516d620ed68ecd04afbc715f2c1c9f3ddae8a7b86ecf848c91", size = 6329473, upload-time = "2026-07-10T08:55:47.603Z" }, - { url = "https://files.pythonhosted.org/packages/60/72/036d11ceca6fbdc4ef95a1aba210cb3fc72bc912703cba49cf2bc8b610c5/ddtrace-4.11.0-cp311-cp311-win_arm64.whl", hash = "sha256:3013b67e1a172eb0edd8c37559b81760a9db205a671541925ebd0127f496409b", size = 5990873, upload-time = "2026-07-10T08:55:49.705Z" }, - { url = "https://files.pythonhosted.org/packages/f7/e9/e50295e520f7d14ce2b5d900f7a2fe8f096d8f90e9040f8d7f43108b72ce/ddtrace-4.11.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:d0eadcf7df3143bf0611f09a351040e475b5df7f512410698cd481a076bca056", size = 7101986, upload-time = "2026-07-10T08:55:56.51Z" }, - { url = "https://files.pythonhosted.org/packages/cb/c6/e8d13546a4b2cd72880d8b0823dd487956d9eb291473d7d03b6e742e5bd9/ddtrace-4.11.0-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:57eb0616dc5c8ff04b5c560ba6d3def69d2bc618e69ce65cc015f103fc68649e", size = 7455362, upload-time = "2026-07-10T08:55:58.664Z" }, - { url = "https://files.pythonhosted.org/packages/50/c6/2c22d71290df36b87393794d6d803ec8ea96ee54d42286836a7ddf0dbae3/ddtrace-4.11.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c311109d1e0f3482b7a44126c855eeb796b186f01604ce74c54800aa5892863d", size = 8502274, upload-time = "2026-07-10T08:56:00.869Z" }, - { url = "https://files.pythonhosted.org/packages/19/84/5afef2c9696b289afa1c608e862930c0e447489776a0133da9f2345edd27/ddtrace-4.11.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:695ec915640072e2917b13d317fe1696086a38924ba02cbe1dcb3918ffb0fa65", size = 8730138, upload-time = "2026-07-10T08:56:03.711Z" }, - { url = "https://files.pythonhosted.org/packages/12/68/a2aff5999fa2b1053b3f50cd78297cbbe7a427f63388fb5f6c500dfb859a/ddtrace-4.11.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:355a288756a7a5d6b73221d28ba4536cb0271d2ea7909adab3e2ee37d62a569f", size = 9507981, upload-time = "2026-07-10T08:56:06.317Z" }, - { url = "https://files.pythonhosted.org/packages/79/12/53bb54b3451f85dbf3c752c2d60b22f672fd55bb08d7495613a70dc87102/ddtrace-4.11.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c19888af09f6ad3340f2935ea5c9eb80d9298141e196fd33e544369920aefba", size = 9793730, upload-time = "2026-07-10T08:56:08.893Z" }, - { url = "https://files.pythonhosted.org/packages/9c/7e/d9a52dd0d877ebdce5eb105d9f3d4e3579391ee0106c8335d59aa97006dd/ddtrace-4.11.0-cp312-cp312-win32.whl", hash = "sha256:cd524b299bd15ed14192cb041a35e992003e0d3a9f8831062f9f28bcfacb63e2", size = 5734745, upload-time = "2026-07-10T08:56:11.764Z" }, - { url = "https://files.pythonhosted.org/packages/6d/da/bb7c609ebc368cabb01db38d7022f42a39d30207f484a5e18a306bd9b333/ddtrace-4.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:af663120d3b4116965bebbd79ce8e73a85644816a0917999e8164b6204678cb6", size = 6317182, upload-time = "2026-07-10T08:56:14.156Z" }, - { url = "https://files.pythonhosted.org/packages/30/f3/e3cf49ce0a4c377fdea8203b127984e26ac2c9cad0521fab27d61371d568/ddtrace-4.11.0-cp312-cp312-win_arm64.whl", hash = "sha256:f0b50a1739573318b0a81fb56a8b22ed6ae6dac017b89400cbc76c17b439d9ed", size = 5980198, upload-time = "2026-07-10T08:56:16.458Z" }, - { url = "https://files.pythonhosted.org/packages/d3/f8/89f2b9cb3d413c55a3d1f19a90fda76ec45abdc0ce13a7f6d4aaf10bed1a/ddtrace-4.11.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:16b9f06c841858debb3de2c4a7ce6acbd6967eaa6ce264dbd0e2a58e376b8583", size = 7094539, upload-time = "2026-07-10T08:56:18.943Z" }, - { url = "https://files.pythonhosted.org/packages/ed/98/f262ad401cdd1a8a2394299d12f31bcc037d0e2dc61d3a843f1c855c293f/ddtrace-4.11.0-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:1cb352d7df16b864f258d4602addf23b8d29d2d416ccc1d0af8d88c614f70a01", size = 7448347, upload-time = "2026-07-10T08:56:21.359Z" }, - { url = "https://files.pythonhosted.org/packages/1f/17/c9123fc263b89934cc1c831b81d90af299d28512740abf008fd963bd006e/ddtrace-4.11.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:54f4cf656aa452bc37b62fa7939c5bddca751ffd4d315b3c925236d4794767ea", size = 8497865, upload-time = "2026-07-10T08:56:24.143Z" }, - { url = "https://files.pythonhosted.org/packages/39/31/00d529b9ca80c77bf3c6aa7370615e2186fc9ead573ef56cef345e6a4dbf/ddtrace-4.11.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:776d99ccb01abecf7dbb4910f4a570b6ab625ed47e00bb6c0159d88d4e2ef1c8", size = 8721481, upload-time = "2026-07-10T08:56:26.983Z" }, - { url = "https://files.pythonhosted.org/packages/41/56/759eab73938820054faa6db6eda607753d821b6322ce7679a2b9f728f8c7/ddtrace-4.11.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5533ce94e50412565a36ce0c8fc6ff6f7490e69b4ecff10c550482e49d5d065e", size = 9505774, upload-time = "2026-07-10T08:56:29.901Z" }, - { url = "https://files.pythonhosted.org/packages/5d/a4/6ab67386244c08340c1c30affffc644b14095e59eb8d2e8d8ee1d4d46d84/ddtrace-4.11.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:81d147bd3d6fa067f2c72c7e32aba90d5d61ed379850896f1cc7fd1ea3437bb1", size = 9787192, upload-time = "2026-07-10T08:56:32.824Z" }, - { url = "https://files.pythonhosted.org/packages/e3/8d/44e074f927b35007e8bf3dc2b5aeb5ba5e4c97e76d4b8e994aaabddfd648/ddtrace-4.11.0-cp313-cp313-win32.whl", hash = "sha256:bcf1d3ba8c3771b8c75f0b9ecd9038ae0a44c8bb0c01d21458c405a29341518a", size = 5732179, upload-time = "2026-07-10T08:56:35.958Z" }, - { url = "https://files.pythonhosted.org/packages/7e/ec/ff034010a1fd598863a15606aa7867ee6abaf4a5989b1a22e61753229432/ddtrace-4.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:d465b9f6f049fc32278cbfd74e1e593b0c3460aacb622bef0cee71e50ffc4b05", size = 6314643, upload-time = "2026-07-10T08:56:38.565Z" }, - { url = "https://files.pythonhosted.org/packages/bc/dd/ad7da7ca071c76815b16a0523875c0c061902cfd8e9ae6595577464e923e/ddtrace-4.11.0-cp313-cp313-win_arm64.whl", hash = "sha256:00cc0fbfa6b8741477528c3fe7bc70a4a5ea439221f188e14da6907f03ed0125", size = 5976789, upload-time = "2026-07-10T08:56:41.172Z" }, - { url = "https://files.pythonhosted.org/packages/12/1b/37b9a84cc4a95d74d4358a0f78346376bc03625b6a9fa0d568a1587665ca/ddtrace-4.11.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:e6f1ca84237e9f52fd71e3458262cfc407dd43337a56e4a6ba9beea88f0958c5", size = 7099136, upload-time = "2026-07-10T08:56:43.766Z" }, - { url = "https://files.pythonhosted.org/packages/f1/42/f227eeb9c5bbb519c67fe11a5326b531f4c1bbd90257e74ab8b72b5ded75/ddtrace-4.11.0-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:c3ddbf9a5a0cd51b5ab9bca898c85cfce73072054b9f5b577164a95065cb8331", size = 7452151, upload-time = "2026-07-10T08:56:46.822Z" }, - { url = "https://files.pythonhosted.org/packages/eb/c4/35fbc0e8ef585d93f82dadae38246eb3643621f08efe11ca0c20e77e7d42/ddtrace-4.11.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7c435b708f4956208141c22d941ed60a1551813353f6c1cfc1a5283ec3d80cc5", size = 8507170, upload-time = "2026-07-10T08:56:49.646Z" }, - { url = "https://files.pythonhosted.org/packages/70/99/a9681a53e9b0557497b24a4e0b1495acf2bcb74368e86129e549d6249a4a/ddtrace-4.11.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:26b3fd7ee97cf7fa1296eecacd393bbbef343cf8bb2d0b1ee39d40cb1d4818fa", size = 8725097, upload-time = "2026-07-10T08:56:53.106Z" }, - { url = "https://files.pythonhosted.org/packages/e0/c6/b18f02120c34abf55e28bc180e872ddbe27614559c01bbcfe5be35b5494a/ddtrace-4.11.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:8d23d66ba9525fd9b9a61c0db8c43e64461066b1d49c75c85f639ea65776367b", size = 9516895, upload-time = "2026-07-10T08:56:56.513Z" }, - { url = "https://files.pythonhosted.org/packages/c9/ed/dc73e8317c229300b0ff8513f903d16ff140eede8084b3b6f76c8646a548/ddtrace-4.11.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:85e081e2295a77c9b38e5792a097f1df8169e1dd9b3339f820d00183e10e8b1b", size = 9792739, upload-time = "2026-07-10T08:57:00.333Z" }, - { url = "https://files.pythonhosted.org/packages/7c/8e/b97984d9f43e2f36f5317034711140c3bd1fbb5d96880b1c11b9e1e8c3d3/ddtrace-4.11.0-cp314-cp314-win32.whl", hash = "sha256:c267fe4efceb71019f5e6833fdb4fbdff18a51641bdca9fceabbe579984461e4", size = 5831072, upload-time = "2026-07-10T08:57:03.578Z" }, - { url = "https://files.pythonhosted.org/packages/e8/6a/7fb18925588f645d6d777c90d388f9a640adf7b986be7457943717b74170/ddtrace-4.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:3bebf3cb12ac51c422ad1ed74873746b8398cce845075b234115c2f518a7a652", size = 6455126, upload-time = "2026-07-10T08:57:06.362Z" }, - { url = "https://files.pythonhosted.org/packages/a0/b2/44d7f07ba94c70bd1827669d446a2e979f7efd6e4a409c469a59684d9a7b/ddtrace-4.11.0-cp314-cp314-win_arm64.whl", hash = "sha256:f596cb2da582600e091c8174ab4f8e44bf4f05c50c03cc450e0378c516e0c6e3", size = 6130096, upload-time = "2026-07-10T08:57:09.413Z" }, - { url = "https://files.pythonhosted.org/packages/99/22/dadad1a1213a846499f360344c4be0854635bbd9c5726d2902bc8e92b194/ddtrace-4.11.0-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9d9a7c4644e4ffc858e951eeab37ab86e9b717ef1bc2f2a47de69f1ec373bccb", size = 6454958, upload-time = "2026-07-10T08:57:12.714Z" }, - { url = "https://files.pythonhosted.org/packages/d4/45/4a5bb647d47f01d135c918ad33678d1dba5b0210cc688da5d2f69534830e/ddtrace-4.11.0-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b49433f0e85e16c34f578e01ec4e263ab5181d9636635add8d6194bd47dd163c", size = 6638706, upload-time = "2026-07-10T08:57:15.791Z" }, - { url = "https://files.pythonhosted.org/packages/80/2b/175096efe3784f36d383e09d7bfa1defdba85689008417008f122a3e8d0c/ddtrace-4.11.0-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:2affbac04583808cba374976d07e59a4fe156d1de122f8def30550bf7717b745", size = 7466814, upload-time = "2026-07-10T08:57:18.873Z" }, - { url = "https://files.pythonhosted.org/packages/80/b3/8fb55c9183ed5a162525ae308e0aaa8237ceb2a3ba78d65a8f1336c16b21/ddtrace-4.11.0-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:51a28ab8072c741e104d8c6f47d990eb5989a359b7ea4d91bab7c5c1f5c1a03c", size = 7707994, upload-time = "2026-07-10T08:57:21.918Z" }, + { url = "https://files.pythonhosted.org/packages/9b/c8/6b8de8b78503f54a6c19f5b4c9430825976f1a2ef1b8342f21ac47c70e22/ddtrace-4.15.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:68c18fa37cd718f2360c980bbd2928a5c2bf502a173e61a83f25e42321d5b4b8", size = 7998372, upload-time = "2026-10-08T09:44:43.566Z" }, + { url = "https://files.pythonhosted.org/packages/67/30/b5ec13494b638453c568ced825b863c086fc7237528aa30498b4fcb6928a/ddtrace-4.15.6-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:7d7696cf05e1675847303701c0a2f7733014f1b2cb5938ed2038e27fa32a18b2", size = 8382818, upload-time = "2026-10-08T09:44:45.73Z" }, + { url = "https://files.pythonhosted.org/packages/fa/dc/57c5d150061cf27935a2e630ca5334f3fa815d82f0585ba694821f793445/ddtrace-4.15.6-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:66647b3c9673468e35fe0294d25015ef6f535bb540a9ce63c76c95ff2679e059", size = 9544001, upload-time = "2026-10-08T09:44:47.366Z" }, + { url = "https://files.pythonhosted.org/packages/f5/7b/6651a225df4438761ddd45c21f624d453aca1508d88ed7fd4c880a3a454a/ddtrace-4.15.6-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4fd260a1339a2df4867f5b0f7ff9111a0f23ba86d6b402e4c3dfe2f71853cd17", size = 9794825, upload-time = "2026-10-08T09:44:49.468Z" }, + { url = "https://files.pythonhosted.org/packages/e4/8d/ffd12a1fa866068c397a1fbbebaebcd6e40f21e618650e95d93860afef44/ddtrace-4.15.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:95a04ff00e401caef808689ae6bd70a9fc92d3572647f9f1bbd43dfacfa053f6", size = 10533745, upload-time = "2026-10-08T09:44:51.57Z" }, + { url = "https://files.pythonhosted.org/packages/bf/19/64e52eced9a66cbb61dd40b110dc25297b1185f1027016876ae366fcaa1e/ddtrace-4.15.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:4eb19b57f1cc6ab8e13e57b8257b91f6b0e5bede753320b30e3920a28d4b1a4f", size = 10833139, upload-time = "2026-10-08T09:44:53.716Z" }, + { url = "https://files.pythonhosted.org/packages/9d/ae/0d964c6eb264ec0ab797071db48221acca6990b5837f1a6d804b88b818ec/ddtrace-4.15.6-cp310-cp310-win32.whl", hash = "sha256:740a87a6504617ec7584d5afb2a6169d780571ddc2b44270274aa8291a70a03c", size = 6756230, upload-time = "2026-10-08T09:44:55.961Z" }, + { url = "https://files.pythonhosted.org/packages/0a/2c/35a05709d630b3ad6f5c05dd90d34bf784fcde93e0722f9e1cec111ec5c4/ddtrace-4.15.6-cp310-cp310-win_amd64.whl", hash = "sha256:7aa9f30da1832111de778a17a26bc2bc564f4810d573ac2d8633609c3db48306", size = 7387627, upload-time = "2026-10-08T09:44:57.786Z" }, + { url = "https://files.pythonhosted.org/packages/2d/b5/3ce5d64ed92d87388a83b52ed766866177997585e224d018922b0494a88d/ddtrace-4.15.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:b3f6e66e3b303252edce525ca573e7092e741049556ae80a395eadf0685a036f", size = 7997045, upload-time = "2026-10-08T09:44:59.612Z" }, + { url = "https://files.pythonhosted.org/packages/f7/25/b7f15af30cf446fa95bf6835e515c135cc801ddc9ec18ecd665b7a82063a/ddtrace-4.15.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:6bec8671c1e0bdadd7fad2513f449c07b9bf4f847bc652aef2720f0e2f74aeaa", size = 8381818, upload-time = "2026-10-08T09:45:01.325Z" }, + { url = "https://files.pythonhosted.org/packages/d9/95/1c597a597123eb99fccd0fc0ef64ee25ad46460e158dbf258136126134c3/ddtrace-4.15.6-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3017a4d4fe8375ae4555b80419d5259bdc444af485bcca6ee52dcd1a41890d35", size = 9541087, upload-time = "2026-10-08T09:45:03.274Z" }, + { url = "https://files.pythonhosted.org/packages/55/88/7e8c66d3db4739599a8092197940d211a006edca56b8fbbf9bbd34869a27/ddtrace-4.15.6-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:220fbb49481c6dd85866110477183b70882cb2578020c3bbbb4db2d89d0b41a8", size = 9793899, upload-time = "2026-10-08T09:45:05.384Z" }, + { url = "https://files.pythonhosted.org/packages/84/f8/08c6b74b9fd13df49cfdb95e6a1a86e06c434994db8f4a4a3395f0fee013/ddtrace-4.15.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:2d50b28793f381bb53e2d7cab5ffeb97c1aa0032f436a1fd5984cb3740d75b10", size = 10534489, upload-time = "2026-10-08T09:45:07.549Z" }, + { url = "https://files.pythonhosted.org/packages/44/6c/03cd4a5cc888f7c540fa9282d3c7c3a7675aae2c622c41668d1a765e0699/ddtrace-4.15.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b5a9edbd149d793b134bdaea28f40db5a199d42135f122c2be325519db55fb75", size = 10834145, upload-time = "2026-10-08T09:45:09.661Z" }, + { url = "https://files.pythonhosted.org/packages/25/6d/ab22cf125e6618597c41b55682dc76617cb6e071ca7381aac346840d6404/ddtrace-4.15.6-cp311-cp311-win32.whl", hash = "sha256:6b6a59fc3dfb081fd7708c9431b94dd17c5148625ec84bc425852cdfbf2f214d", size = 6752572, upload-time = "2026-10-08T09:45:11.785Z" }, + { url = "https://files.pythonhosted.org/packages/b6/38/ffc989894b1d7ea7c06fb249185807c574da5b4572adb9c974fcf02ea65d/ddtrace-4.15.6-cp311-cp311-win_amd64.whl", hash = "sha256:fef3cdf314eee555c5f6b831aeeb826a7eec147e3e1ab85591154a0ad96828f9", size = 7388927, upload-time = "2026-10-08T09:45:14.111Z" }, + { url = "https://files.pythonhosted.org/packages/ab/b8/1b8483fac4a4bcf1d7338d6323b2a3bde576553b111c5a288aee71224613/ddtrace-4.15.6-cp311-cp311-win_arm64.whl", hash = "sha256:4f163cc7fd9d76711c9df54273cbc2bb15933d90d1027fbf6ea3e938f042d23a", size = 7391090, upload-time = "2026-10-08T09:45:16.749Z" }, + { url = "https://files.pythonhosted.org/packages/c1/3f/010c67d812aa0950bb44aeabec69c858e5bd7d31d9fe560ae431e2a08d95/ddtrace-4.15.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:c4ca75ae638e105e5c28b7f7b3e1d3762f027ab14a315530e62c71d402b49dc1", size = 7993463, upload-time = "2026-10-08T09:45:18.638Z" }, + { url = "https://files.pythonhosted.org/packages/16/e3/c42f708e8506a75621678e5aac5219b03c8d95014b70e5755287cbe81d2a/ddtrace-4.15.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:2f533ad60bb9855ab47a429eeb8b7e28667647c1a482899ac6bd5c00aa525407", size = 8387610, upload-time = "2026-10-08T09:45:20.667Z" }, + { url = "https://files.pythonhosted.org/packages/a3/c5/ea86a9859f8e9b69a448c1ccbb70c51c7989c8c7bf68288356b9f53883a1/ddtrace-4.15.6-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e00a2943adf9cc886da3cf377b6c32001828565c877a9f0a7905e60eda527dad", size = 9514941, upload-time = "2026-10-08T09:45:22.929Z" }, + { url = "https://files.pythonhosted.org/packages/8d/af/b3bd0fdacb0d2b6baf8bd033ea00e97b06e74e3e504918ecf2f4d176ad6b/ddtrace-4.15.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:dfb419eecf9ab3ad774e0fa612f08814530ebd2e1f0aa10128bb0f98eb7cea2e", size = 9772538, upload-time = "2026-10-08T09:45:25.295Z" }, + { url = "https://files.pythonhosted.org/packages/ce/05/f1cb0e9007a5c0d155948726bfefe97fc5fc9d7e886a951a4d51bc1c0236/ddtrace-4.15.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7f8bf477362e29742a6b282fbbfc9c6985c745ba302e4b6a8a53388911ceacae", size = 10511565, upload-time = "2026-10-08T09:45:27.732Z" }, + { url = "https://files.pythonhosted.org/packages/09/90/9f27058af312b213abfb946ab27b334b58b5482aada8b00ce2c41bee2ee6/ddtrace-4.15.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bd6c6cd99d6edcffab804755f8246f3c1ecf727bdb241ecfe94f5de31fbb484b", size = 10824106, upload-time = "2026-10-08T09:45:30.202Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9a/f4402bdfb563d103ece19bd56a4622accd0113dfb9ac9b118050757823f9/ddtrace-4.15.6-cp312-cp312-win32.whl", hash = "sha256:903a24868a743e035ff0f7bdc7e397a90c00fac77b551765a9ceae5537c40a44", size = 6756938, upload-time = "2026-10-08T09:45:32.864Z" }, + { url = "https://files.pythonhosted.org/packages/55/fb/0b24d9fbcecfeb72ac2fb24db5b4a4684a0f0ecda73bb2b0c56c14282645/ddtrace-4.15.6-cp312-cp312-win_amd64.whl", hash = "sha256:57cd4ca85c5fef665fc2794c37369e39d6b9f51b1ff37ae8e68ac935360d1561", size = 7375054, upload-time = "2026-10-08T09:45:35.014Z" }, + { url = "https://files.pythonhosted.org/packages/b6/34/09763e279fb865196966cc895fb234cf12350f08fde8a284c796f13d6d34/ddtrace-4.15.6-cp312-cp312-win_arm64.whl", hash = "sha256:be11b8c8fce20b6e02d6069a53d43fc92aec9a50741e0f0dadebd913f2917ae5", size = 7376141, upload-time = "2026-10-08T09:45:37.46Z" }, + { url = "https://files.pythonhosted.org/packages/17/0c/f0b7dfb83917ea3f9c0d072cf8cf9454550ecf26372b19be49b8cf16bbcb/ddtrace-4.15.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:8a118c85925d143e88e48e4583500e95a9b88d9c1fabc96df276c086eb515603", size = 7987311, upload-time = "2026-10-08T09:45:40.025Z" }, + { url = "https://files.pythonhosted.org/packages/f5/01/f39c527b10ceef714c16a99d59429ddd8b314c7d5301bff86fab6201dfae/ddtrace-4.15.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:1307ee87674c3d2d93f38bcd258b62b70a41081183e4bc94b03594c855751322", size = 8382130, upload-time = "2026-10-08T09:45:42.316Z" }, + { url = "https://files.pythonhosted.org/packages/e2/29/0c46ee3d6dcce28f8bd82c1dc658e420bed935fc0488a1ffac2853b42b0f/ddtrace-4.15.6-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fabb658b1ec068ab6046fb1ecaa6e6860bd40338235244c3c156c973c4957c00", size = 9507940, upload-time = "2026-10-08T09:45:44.887Z" }, + { url = "https://files.pythonhosted.org/packages/b7/b0/77d18cb89ff7b2652d734ae00e4869ef703945a0393bc465e00878631099/ddtrace-4.15.6-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2185df0b730e4bb78dfe40cf61b7e242d4b81724cea4f496ca5a0ae0f5427309", size = 9764498, upload-time = "2026-10-08T09:45:47.528Z" }, + { url = "https://files.pythonhosted.org/packages/64/06/7435d3534b646ddafaf16fced76b0dec8d239492eef1e119d4c3b3fe1259/ddtrace-4.15.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7fc21e13b52290d33bcdec36644eb63eac03f49adb0b36c6d84dc4bb02b9b935", size = 10508643, upload-time = "2026-10-08T09:45:50.63Z" }, + { url = "https://files.pythonhosted.org/packages/34/d0/65db05d5b72fcb88378edd4bf44e88b2dc0c3e8e1c3070a4630e5ec2dfdd/ddtrace-4.15.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:008a5c93db9489861394190b9a9ecffc464a9a83882792c83c005e938fc69df0", size = 10817707, upload-time = "2026-10-08T09:45:53.417Z" }, + { url = "https://files.pythonhosted.org/packages/df/e0/b2ad40242ce8d17ba4432c4f96a7584d93852c442fcb8ac1285cbd10ae96/ddtrace-4.15.6-cp313-cp313-win32.whl", hash = "sha256:d5ca2d38c37194aee97801775193a034f7f27f2610d5075b6a550ee664b044dd", size = 6754044, upload-time = "2026-10-08T09:45:56.603Z" }, + { url = "https://files.pythonhosted.org/packages/1e/e3/7a549bc05e6b604ef3c5ad8fef19a951178d8d93ee7088660d3743e53b7b/ddtrace-4.15.6-cp313-cp313-win_amd64.whl", hash = "sha256:fc35805f2c5b64bd7a29304a5c54932792a47b1436d2550354c051870cd706b4", size = 7370629, upload-time = "2026-10-08T09:45:58.955Z" }, + { url = "https://files.pythonhosted.org/packages/50/dc/a140b6640dd57f143e9137e8e4d04653ba9047c7ef65ff26ca03721c90ff/ddtrace-4.15.6-cp313-cp313-win_arm64.whl", hash = "sha256:03221f056c03728b01a4295ff4a3854db2dc33bcfab3131757f5eb4ca2a5a41b", size = 7378982, upload-time = "2026-10-08T09:46:01.698Z" }, + { url = "https://files.pythonhosted.org/packages/62/de/448c66ce790d6fcd1abbeb08e8e88ef1a5b7fc990f6ec48b818b4d9a8b29/ddtrace-4.15.6-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:70d5dfc73d970fa1b633e0aad3b9457cb94478958976e139b1c78461a3c48b33", size = 7993200, upload-time = "2026-10-08T09:46:04.284Z" }, + { url = "https://files.pythonhosted.org/packages/2f/87/50bbd2d4fe589aa13038c1fe07defd57e63c1bc83c22b92eac13676b3219/ddtrace-4.15.6-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:26ead4c2956924d0574b6681d063dee6c8741dca98ed3b747c353f6252eae0d7", size = 8388799, upload-time = "2026-10-08T09:46:07.071Z" }, + { url = "https://files.pythonhosted.org/packages/9e/13/d947c88b0fdd7cddb89174f5e9116aead05a935e1406e2921c820cb7bcab/ddtrace-4.15.6-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:03e62ec9c83e9e98122f4544592a6bf7d3338e9ea7ff2e34fe4750ade2b2c3e5", size = 9519810, upload-time = "2026-10-08T09:46:09.975Z" }, + { url = "https://files.pythonhosted.org/packages/48/5a/15711bfd54778a28cca9f7988bcdd03e4d680dfceed7f4db22f6b7a06e73/ddtrace-4.15.6-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3684507243feb45e7438b68989aae9a7df01b7996112d7fe74846e1d5aeb3726", size = 9771503, upload-time = "2026-10-08T09:46:13.009Z" }, + { url = "https://files.pythonhosted.org/packages/7d/d4/d0e81dff83be9370b90c090828330c4b2f06d74644f331d8f6d9518bad44/ddtrace-4.15.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:4882e25808cf6102015d7421d6214aa94a788ce3242fec9a1a726a41dcfd79ac", size = 10521211, upload-time = "2026-10-08T09:46:15.867Z" }, + { url = "https://files.pythonhosted.org/packages/f7/c1/e33b44ebfa4dcb80f66fb2cc99c61d9c7706edd48e5ccb5ad04bcea0d5ad/ddtrace-4.15.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5336630ed8c829333dd4195d817b3f3f8211ddf6074e1c731e884771094fc41c", size = 10826453, upload-time = "2026-10-08T09:46:19.084Z" }, + { url = "https://files.pythonhosted.org/packages/e5/32/2125cf7f03e9e8b68b6a9072e393a3be4af91505ada32bc65811d6965881/ddtrace-4.15.6-cp314-cp314-win32.whl", hash = "sha256:a970bc70c81361c9a27452bfd3335182d722315466109208290a1c3a4c221748", size = 6867770, upload-time = "2026-10-08T09:46:22.66Z" }, + { url = "https://files.pythonhosted.org/packages/b4/a0/c18b1db1cf4843a3254aed03b61c1c9cbd5964bd52c18922f633c7ed2304/ddtrace-4.15.6-cp314-cp314-win_amd64.whl", hash = "sha256:e7e8b948d13737e61a8e0dc1fe082d38431a3d4fb285d23236f1c8f733ee5c67", size = 7538347, upload-time = "2026-10-08T09:46:25.41Z" }, + { url = "https://files.pythonhosted.org/packages/8f/77/6181506d7f7693aceccdc13f5a2dc6b18662ed3eedebc02043f00d502e5e/ddtrace-4.15.6-cp314-cp314-win_arm64.whl", hash = "sha256:11662a7a4e1f51c388450cd1092b57738bb32ef98401c6032754a7e220a5f6d0", size = 7574014, upload-time = "2026-10-08T09:46:28.185Z" }, +] + +[[package]] +name = "ddtrace-internal" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/51/6dc02752dca5920d6627fb92b7e0611f594a823bbd57ed31167571d819bc/ddtrace_internal-0.1.0.tar.gz", hash = "sha256:7390721276ab5f174cb103088be7d0de22d452f2817c949e8bde1182ff5b961e", size = 2504, upload-time = "2026-07-15T18:35:10.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/ad/3cae600fdfc41f54653c14339ac1f125381a6fc2a823d5e685b627ff152e/ddtrace_internal-0.1.0-py3-none-any.whl", hash = "sha256:8eeac4a812e003cebafc8be35ef05103ae2b132893dfbcee00a67f4af81404a2", size = 2495, upload-time = "2026-07-15T18:35:09.295Z" }, ] [[package]] @@ -1948,11 +1974,30 @@ wheels = [ name = "dnspython" version = "2.8.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] sdist = { url = "https://files.pythonhosted.org/packages/8c/8b/57666417c0f90f08bcafa776861060426765fdb422eb10212086fb811d26/dnspython-2.8.0.tar.gz", hash = "sha256:181d3c6996452cb1189c4046c61599b84a5a86e099562ffde77d26984ff26d0f", size = 368251, upload-time = "2025-09-07T18:58:00.022Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" }, ] +[[package]] +name = "dnspython" +version = "2.9.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.15'", + "python_full_version == '3.14.*'", + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", +] +sdist = { url = "https://files.pythonhosted.org/packages/ef/4a/50822184bd67cc6493f0fb6a880749158fcd31ab3fa07409acfd91f9fc85/dnspython-2.9.0.tar.gz", hash = "sha256:b44dc6b18f07a8b1c56676a19fbfdb5209415b046a9cece286baafa87ff3f7f1", size = 423560, upload-time = "2026-10-09T00:07:24.352Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/02/cdcc9b7c051786a103c3b09e1003a82fa0c66bcb91ffbdabcfbf7b4163b9/dnspython-2.9.0-py3-none-any.whl", hash = "sha256:9a4aedb833c3c1b49214d04d44d3032ab7a9135f7c1d29a549b4ff78fd82fda9", size = 354822, upload-time = "2026-10-09T00:07:22.622Z" }, +] + [[package]] name = "docker" version = "7.1.0" @@ -2009,7 +2054,8 @@ name = "email-validator" version = "2.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "dnspython" }, + { name = "dnspython", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "dnspython", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "idna" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f5/22/900cb125c76b7aaa450ce02fd727f452243f2e91a61af068b40adba60ea9/email_validator-2.3.0.tar.gz", hash = "sha256:9fc05c37f2f6cf439ff414f8fc46d917929974a82244c20eb10231ba60c54426", size = 51238, upload-time = "2025-08-26T13:09:06.831Z" } @@ -5329,7 +5375,7 @@ ci = [ { name = "psycopg2-binary", specifier = "==2.9.11" }, { name = "pyarrow", specifier = "==23.0.1" }, { name = "pygithub", specifier = "==2.8.1" }, - { name = "pylint", specifier = "==4.0.5" }, + { name = "pylint", specifier = "==4.0.9" }, { name = "pytest-codspeed", specifier = "==4.3.0" }, { name = "pytest-retry", specifier = "==1.7.0" }, { name = "tenacity", specifier = "==8.5.0" }, @@ -6111,75 +6157,100 @@ wheels = [ [[package]] name = "msgpack" -version = "1.2.1" +version = "1.2.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/31/f9/c0a1c127f9049db9155afc316952ea571720dd01833ff5e4d7e8e6352dbb/msgpack-1.2.1.tar.gz", hash = "sha256:04c721c2c7448767e9e3f2520a475663d8ee0f09c31890f6d2bd70fd636a9647", size = 183960, upload-time = "2026-06-18T16:13:52.594Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0a/e7/bb605a7bab2d8425a64b3fa762b39dc1bf1c7e3f11ba6fb5413d6db0ff8c/msgpack-1.2.3.tar.gz", hash = "sha256:32edb81a2b5eb7cd7c9d941b2bfbbb082fd2cd09e0e725930316af6b708db186", size = 196517, upload-time = "2026-09-29T02:33:52.276Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5b/16/f70100614b69feb3ade7285f08c9c52d6cda0a5c03f3f5e2facd63acb211/msgpack-1.2.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8c7b398c56ff125feae96c2737abfec5595f1fa0aa186df60c56040b8accb95c", size = 82926, upload-time = "2026-06-18T16:12:31.531Z" }, - { url = "https://files.pythonhosted.org/packages/e4/3c/08ecd5cdfe4e2de43aec79062028ad0f7b2d9b1fea5430068c198ba570da/msgpack-1.2.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1548006a91aa93c5da81f3bdcebc1a0d10cea2d25969754fbe848da622b2b895", size = 82730, upload-time = "2026-06-18T16:12:32.894Z" }, - { url = "https://files.pythonhosted.org/packages/19/9f/a70c9cb1a04ecc134005149367dcfe35d167284e8f65035a1e4156ad17b5/msgpack-1.2.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1dabedcd0f23559f3596428c6589c1cd8c6eaed3a0d720795b07b0225d769203", size = 400729, upload-time = "2026-06-18T16:12:34.052Z" }, - { url = "https://files.pythonhosted.org/packages/fa/7f/5ce020168cf0439041526e95aa068c722c016aee21624e331aeabeee2e8e/msgpack-1.2.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:83efa1c898e0fc5380fc0cabbf75164c52e3b5cbb45973710d75821928380c73", size = 407625, upload-time = "2026-06-18T16:12:35.239Z" }, - { url = "https://files.pythonhosted.org/packages/79/70/fb7668ce0386819303047057aef6fc1da73b584291d9cff82b821744e2ef/msgpack-1.2.1-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:01e2dd6c9b19d333a00282330cc8a73d38d8dabc306dc5b42cd668c3ac82e833", size = 377891, upload-time = "2026-06-18T16:12:36.684Z" }, - { url = "https://files.pythonhosted.org/packages/3d/dc/9ebe654a73c3aed2e40aa6b52e3c2a02b5f53ef0085fa235a45d5b367f87/msgpack-1.2.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:350cb813d0af6e65d2f7ef0d729f7ff5be5a8bce03665892f43e5883d4ecc1b8", size = 391987, upload-time = "2026-06-18T16:12:37.839Z" }, - { url = "https://files.pythonhosted.org/packages/42/eb/b67cf64218a2fa25e1c671fe1d3dbb06cbeb973e71bc4b822da079862d0b/msgpack-1.2.1-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:ee1d9ed27d0497b848923746cf762ed2e7db24f4be7eec8e5cbe8c766aa707b7", size = 374603, upload-time = "2026-06-18T16:12:39.221Z" }, - { url = "https://files.pythonhosted.org/packages/a2/2e/9ee200cde32fd1a0101b4006202fde554c1860adfb9bf7bff31ea4c08df8/msgpack-1.2.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:633727297ed063441fd1cda2288865487f33ad14eeb8831afb5f0c396a62cfce", size = 405121, upload-time = "2026-06-18T16:12:40.524Z" }, - { url = "https://files.pythonhosted.org/packages/43/b6/f10117be7ca7a51e8feed699a907b8e663a8cd66e115ae6b4fb30cc7945c/msgpack-1.2.1-cp310-cp310-win32.whl", hash = "sha256:298872ecf9e61950f1c6af4ca969b859ee91783bb920ef6e6172697d0c8aad74", size = 64088, upload-time = "2026-06-18T16:12:41.762Z" }, - { url = "https://files.pythonhosted.org/packages/ba/93/89976c696fb0224662239d952c47b4d1661b34d79a332ef5584facaa8579/msgpack-1.2.1-cp310-cp310-win_amd64.whl", hash = "sha256:2ff164c1b0bcb740b073b99e945234d0212852fa378e44a208c425379140dbeb", size = 70113, upload-time = "2026-06-18T16:12:42.78Z" }, - { url = "https://files.pythonhosted.org/packages/f4/6b/e9b1cdc042c4458801d2545ed782a95f3d6ba8e270cce8745b8603c7f748/msgpack-1.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:29a3f6e9667868429d8240dfd063ea5ffdc1321c13d783aa23827a38de0dcb22", size = 82812, upload-time = "2026-06-18T16:12:45.022Z" }, - { url = "https://files.pythonhosted.org/packages/0c/3a/dd518a1bf78ed1e9ad8afe57307c079a00eafe4b3068932a27ca1ea56b4f/msgpack-1.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:aded5bdf32609dc7987a49bbbd15a8ef096193f96dd8bbeb791de729e650acf5", size = 82739, upload-time = "2026-06-18T16:12:46.025Z" }, - { url = "https://files.pythonhosted.org/packages/70/e0/7ba9e1542bf0771a27b8b37c1316e3f95ae9d748fd765284655c476ad4ef/msgpack-1.2.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:146ee4e9ce80b365c6d4c47073da9da7bcec473e58194ceee5dd7620ace77e06", size = 414233, upload-time = "2026-06-18T16:12:47.029Z" }, - { url = "https://files.pythonhosted.org/packages/03/8d/671d81534ea0e2b0e8a121be100020da09eb78861fe3aa8f3ef7dcd3bed1/msgpack-1.2.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a28d076ca7c82b9c8728ad90b7147489449557038bed50e4241eb832395169b4", size = 423843, upload-time = "2026-06-18T16:12:48.19Z" }, - { url = "https://files.pythonhosted.org/packages/d2/b6/e5c737515ed1f166664b87601b532f58cbb73d8aa6a90b99f7c2c5037e8e/msgpack-1.2.1-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7d31c0ac0c640f877804c67cb2bc9f4e23dc2db97e96c2e67fa27d38283b41f8", size = 390772, upload-time = "2026-06-18T16:12:49.624Z" }, - { url = "https://files.pythonhosted.org/packages/a8/46/62ed8c2e87d7021eab19921594d961ef3aa3794eec76c716dc30f3bfd433/msgpack-1.2.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8ff92d7feeaf5bc26c51495b69e2f99ed97ab79346fb6555f44be7dd2ac6503b", size = 409559, upload-time = "2026-06-18T16:12:50.936Z" }, - { url = "https://files.pythonhosted.org/packages/70/ff/59aa3887b860bbf43532835e192b1c388a17590d6068ae4f8b2bc74c906e/msgpack-1.2.1-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:779197a6513bab3c3632265e3d0f7cb3227e62510841a6f34f1eaa37efbb345e", size = 387838, upload-time = "2026-06-18T16:12:52.161Z" }, - { url = "https://files.pythonhosted.org/packages/09/11/f8563e471093420cf6478cb3271a0175d8402b82d879783d4035d2d03360/msgpack-1.2.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:67f6dd22fa72a93752643f07889796d62739a13415ee630169a8ce764f86cf9f", size = 421732, upload-time = "2026-06-18T16:12:53.556Z" }, - { url = "https://files.pythonhosted.org/packages/57/cf/e673683c4c6c90c1022b24c65af4b03eda72b182a1176ef6449069d66acc/msgpack-1.2.1-cp311-cp311-win32.whl", hash = "sha256:91054a783328e0ea7954b8771095705c8d2243b814743fbaadf14552c9c52c5d", size = 64091, upload-time = "2026-06-18T16:12:54.821Z" }, - { url = "https://files.pythonhosted.org/packages/3f/07/ca212739d179f9083bff2c7c08c24101c3555a334fadc2b876b18768a3ae/msgpack-1.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:2eda0b7ebb1283a98d3e4492ac933c8af6aff59fd3df1c3ed024f536af4b1dc8", size = 70462, upload-time = "2026-06-18T16:12:55.898Z" }, - { url = "https://files.pythonhosted.org/packages/6d/be/6798347b425e26f35db82e69dd83c09716c856a3714e7bffc4c0860fd830/msgpack-1.2.1-cp311-cp311-win_arm64.whl", hash = "sha256:6ee967f7c7e1df2890c671ff2ee51a28ded0efc95da3e507176dee881ce36c66", size = 65059, upload-time = "2026-06-18T16:12:57.053Z" }, - { url = "https://files.pythonhosted.org/packages/bc/dd/9e8cbd8f5582ca4b590336f2b91ee5662f6a6ca562b565abaf696a0f81ff/msgpack-1.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2ef59c659f289eddf8aa6623823f19fa2f40a4029266889eac7a2505dd210c35", size = 83531, upload-time = "2026-06-18T16:12:58.249Z" }, - { url = "https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d3567748a5107cb40cdf66a275430c2f87c07777698f4bfd25c35f44d533258c", size = 82657, upload-time = "2026-06-18T16:12:59.396Z" }, - { url = "https://files.pythonhosted.org/packages/26/aa/753ad8b007b464e1d8aa0c8e650b9c5f4f725e658fc5ac8a7635c55b7f6e/msgpack-1.2.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:60926b75d00c8e816ef98f3034f484a8bc64242d66839cef4cf7e503142316a0", size = 410634, upload-time = "2026-06-18T16:13:00.383Z" }, - { url = "https://files.pythonhosted.org/packages/6a/fd/6adabd4f6d5e686f97dd02ce7fce3fe4cf672cbac36b8f67ff4040e8ad8b/msgpack-1.2.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:020e881a764b20d8d7ca1a54fc01b8175519d108e3c3f194fddc200bda95951a", size = 419989, upload-time = "2026-06-18T16:13:01.776Z" }, - { url = "https://files.pythonhosted.org/packages/5a/cc/85039b7b0eb168aaad7383a23c97e291a11f08351cb45a606ce865e4e3f1/msgpack-1.2.1-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4202c74688ca06591f78cb18988228bd4cca2cc75d57b60008372892d2f1e6e6", size = 377544, upload-time = "2026-06-18T16:13:03.637Z" }, - { url = "https://files.pythonhosted.org/packages/ed/bf/35963899493b32030c85fc513b723ae66144ac70c11ebc52e889e16e3d99/msgpack-1.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8b267ce94efb76fbd1b3373511420074ee3187f0f7811bf394531de13294735a", size = 400842, upload-time = "2026-06-18T16:13:05.012Z" }, - { url = "https://files.pythonhosted.org/packages/a6/df/8e2ac970c8f99264cd9997d1c73df5466bc19da3301d7dc5500862a9b089/msgpack-1.2.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:e4f1d0f8f98ade9634e01fb704a408f9336c0a8f1117b369f5db83dc7551d8b1", size = 374108, upload-time = "2026-06-18T16:13:06.232Z" }, - { url = "https://files.pythonhosted.org/packages/17/dd/fa8bd265110dfa51c20cb529f9e6d240a16fafe7e645004c6af2d01353ba/msgpack-1.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f02cf17a6ca1abe29b5f980644f7551f94d71f2011509b26d8625ce038f0df64", size = 414939, upload-time = "2026-06-18T16:13:07.478Z" }, - { url = "https://files.pythonhosted.org/packages/2e/b9/8377a5ad8953fc0437c70cc98d9ae29f27fe5ac5109fbec0812085865735/msgpack-1.2.1-cp312-cp312-win32.whl", hash = "sha256:0c0d9802354507bcba62af19c17918e3eb437cc25e6f50657d511b5856a77aac", size = 64504, upload-time = "2026-06-18T16:13:08.822Z" }, - { url = "https://files.pythonhosted.org/packages/57/7f/ce1e377df7e62461fefd9eb23bfb93a4a523f40a517b377b8f844d836828/msgpack-1.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:5c24aa15d5963051e1a5c62b12c50cd705992502b5ec1f3bece6046f33c9fc24", size = 71421, upload-time = "2026-06-18T16:13:09.828Z" }, - { url = "https://files.pythonhosted.org/packages/8f/32/ebfe84c9929f08f188d56c7a2fd913406a9ddad76a634697c1c43b8112e6/msgpack-1.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:4227224aaec8f7fbcbfbd4272319347b2bb4030366502600f8c45588c5187b07", size = 64775, upload-time = "2026-06-18T16:13:11.056Z" }, - { url = "https://files.pythonhosted.org/packages/b0/ac/dcddcab6f6c20ecb387ca5e980371cdb3f87ff69aeca388be97eebc4c074/msgpack-1.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0a70e3cf2804a300d921bb0940426e35f4e489a23adfb77a808892241db0a064", size = 83151, upload-time = "2026-06-18T16:13:12.173Z" }, - { url = "https://files.pythonhosted.org/packages/64/71/fbcfa83a1d6a9c6091942d1cfd070962244664b87427a9a49a6897b1b219/msgpack-1.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:491cc39455ca765fad51fb451bf2915eb2cf41192ab5801ce8d67c1d614fe056", size = 82351, upload-time = "2026-06-18T16:13:13.194Z" }, - { url = "https://files.pythonhosted.org/packages/e3/10/ddf7b06db879e8792d13934ddda09ff20bd2a583fd84c9b59aae9b0e650b/msgpack-1.2.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f310233ef7fb9c14e201c93639fe5f5260b005f56f0b29048e999c30935596cc", size = 407518, upload-time = "2026-06-18T16:13:14.233Z" }, - { url = "https://files.pythonhosted.org/packages/79/d3/36a46a8ed992b781acbc05928bd5bee3c810cb0c3563bf81a7b0c04a1a76/msgpack-1.2.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:787c9bebb5833e8f6fc8abca3c0597683d8d87f56a8842b6b89c75a5f3176e2d", size = 416405, upload-time = "2026-06-18T16:13:15.435Z" }, - { url = "https://files.pythonhosted.org/packages/f9/84/e8e9598b557c0ba6ddae901a73780a4c75ac667dddf59414b1e56a42fb34/msgpack-1.2.1-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:dc871b997a9370d855b7394465f2f350e847a5b806dd38dcc9c989e7d87da155", size = 376257, upload-time = "2026-06-18T16:13:17.022Z" }, - { url = "https://files.pythonhosted.org/packages/40/16/738fe6d875ad7e2a9429c165322a4ec088f4f273cdfae63d96a89c467961/msgpack-1.2.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:85f57e960d877f2977f6430896191b04a21f8901b3b4baf2e4604329f4db5402", size = 397469, upload-time = "2026-06-18T16:13:18.287Z" }, - { url = "https://files.pythonhosted.org/packages/ca/be/6d5952df75a7f24f35833af764c3a6860780364cb3a0030beb8099e1b2b4/msgpack-1.2.1-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:1233ee2dd0cefba127583de50ea654677277047d238303521db35def3d7b2e7c", size = 372802, upload-time = "2026-06-18T16:13:19.685Z" }, - { url = "https://files.pythonhosted.org/packages/e1/39/e2ef7dbf0473bcb8dc7c50bf782a892d67414877b63e47fc88eb189ef5e6/msgpack-1.2.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e3dc2feb0876209d9c38aa56cb1de169bd6c4348f1aa48271f241226590993e6", size = 411273, upload-time = "2026-06-18T16:13:21.028Z" }, - { url = "https://files.pythonhosted.org/packages/ef/c5/133f4512a56e983a93445c836c9d94d88f3bc2e0980ff4b9e577bd8416ce/msgpack-1.2.1-cp313-cp313-win32.whl", hash = "sha256:6d09badf350af2be9d189184e04e64cf54ad93569ab3d96fca58bd3e84aad707", size = 64471, upload-time = "2026-06-18T16:13:22.293Z" }, - { url = "https://files.pythonhosted.org/packages/e2/98/577e10b055096a7dd40732358cabaf7180a20c79ed1dcdbb618e4b9deac7/msgpack-1.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:33f14fba63278b714efe6ad07e50ea5f03d91537aa6a1c5f1ceca4cf44013ca9", size = 71274, upload-time = "2026-06-18T16:13:23.455Z" }, - { url = "https://files.pythonhosted.org/packages/ba/ee/0c0048e7cfbef23c6a94791b8959ab28155232e7956de8a305b5ff588f05/msgpack-1.2.1-cp313-cp313-win_arm64.whl", hash = "sha256:afc5febcd4c99effbc02b528e49d6fd0760b2b7d48c05239e345a5fa6e743d9a", size = 64795, upload-time = "2026-06-18T16:13:24.687Z" }, - { url = "https://files.pythonhosted.org/packages/77/58/cce442852c6b9e1639c7c8ac8fd9143121cb32dab0f308df4d1426a8eb9c/msgpack-1.2.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:05f340e47e7e47d2da8db9b53e1bb1d294369e9ef45a747441309f6650b8351d", size = 83610, upload-time = "2026-06-18T16:13:25.724Z" }, - { url = "https://files.pythonhosted.org/packages/60/5c/15b4c7a0182f75ffa90751958ba36a9c01cafee367d49a3edc10ed140b01/msgpack-1.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:810b916696c86ef0deb3b74588480224df4c1b071136c34183e4a2a4284d7ac7", size = 83138, upload-time = "2026-06-18T16:13:26.781Z" }, - { url = "https://files.pythonhosted.org/packages/b8/a6/99e58722feaffc5f2fbcc0c8c0d1451ab9f84097f7af87291b46af2390f4/msgpack-1.2.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ca0dacff965c47afdc3749a8469d7302a8f801d6a28758d55120d75e66ce6889", size = 406090, upload-time = "2026-06-18T16:13:28.072Z" }, - { url = "https://files.pythonhosted.org/packages/19/03/8c63e8cf52958534ef688625965ab04c269a6cadd8caef16758b380a821a/msgpack-1.2.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e2bf9280bceb5efca998435904b5d3e9fdbcc11d90dc9df30aec7973252b720", size = 412106, upload-time = "2026-06-18T16:13:29.427Z" }, - { url = "https://files.pythonhosted.org/packages/63/d2/155d9e71b40e41fd934bc0c48b9b2770f22263e1ac20aad8e29fdca7be3f/msgpack-1.2.1-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:aa6c4be5d1c02a42b066ca6ddb71adf36432868fdcdb6ee87e634e86e0674190", size = 374851, upload-time = "2026-06-18T16:13:30.631Z" }, - { url = "https://files.pythonhosted.org/packages/98/48/deaf2326262a8d5ea3295ce9649912ecd3f551ba7ec8e33c665d2ba583f3/msgpack-1.2.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec0e675d59150a6269ddc9139087c722292664a37d071a849c05c473350f1f2d", size = 396168, upload-time = "2026-06-18T16:13:31.977Z" }, - { url = "https://files.pythonhosted.org/packages/10/2a/b4410f906c2ec0008f1608d3ab5143afc3ad3f4e6da0fed3ea2231d0bef4/msgpack-1.2.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:dd3bfe82d53edfe4b7fc9a7ec9761e23a7a5b1dac22264505af428253c29ed24", size = 371959, upload-time = "2026-06-18T16:13:33.282Z" }, - { url = "https://files.pythonhosted.org/packages/59/86/1edc67270099a528fa2093ea60fe191233cd238e4bd30cfacf7db79fc959/msgpack-1.2.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5ad5467fc3f68b5468e06c5f788d712e9f8ffc8b0cd1bcb160c105c1ee92dae7", size = 408457, upload-time = "2026-06-18T16:13:34.567Z" }, - { url = "https://files.pythonhosted.org/packages/82/90/8b630fef07d8c5ab457b71ff2c217910c83d333c7a68472c186e87cc504a/msgpack-1.2.1-cp314-cp314-win32.whl", hash = "sha256:98b58bdb89c46190e4609bb36abe17c6d4105ad13f9c5f8f6f64d320f8ced3fb", size = 65942, upload-time = "2026-06-18T16:13:36.056Z" }, - { url = "https://files.pythonhosted.org/packages/16/f1/467b81e98b24dd3885d7b1857728797b4ffc76a7a7483af4fb321a07de3c/msgpack-1.2.1-cp314-cp314-win_amd64.whl", hash = "sha256:74847557e28ce71bd3c438a447ca90e4b507e997ddbdef8a12a7b283b86c156b", size = 72627, upload-time = "2026-06-18T16:13:37.079Z" }, - { url = "https://files.pythonhosted.org/packages/a7/1d/5d8c4c89985feb6acefb82a09e501c60392261856d2408d20bfe4f0360b1/msgpack-1.2.1-cp314-cp314-win_arm64.whl", hash = "sha256:b50b727bd652bdc37d950336c848ef20ec54a4cafc38dce19b1cd86ad625d0f7", size = 66908, upload-time = "2026-06-18T16:13:38.23Z" }, - { url = "https://files.pythonhosted.org/packages/1b/02/ad2afb678b4de94496cd432b581759b756a92c1192d8c767edd6b132efdc/msgpack-1.2.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:8d00f177ca88a77c1cf848d204a38f249751650b601cb6532acc68805d8a8273", size = 86000, upload-time = "2026-06-18T16:13:39.44Z" }, - { url = "https://files.pythonhosted.org/packages/54/74/0b797484013128837f3b1cbb6cea019277c4de4e377dc512b4d9a0f92940/msgpack-1.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5bb9c386f0a329c035ddbab4b72d1028bf9627add8dda41070288563d57ed1b1", size = 86544, upload-time = "2026-06-18T16:13:40.447Z" }, - { url = "https://files.pythonhosted.org/packages/a9/b4/b774d7eb95561739907fec675582f83203cf41c597a418c2589b4bfb8e9d/msgpack-1.2.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:20466cca18c49c7292a8984bc15d65857b171e7264bdcb5f96baf8be238791fc", size = 427661, upload-time = "2026-06-18T16:13:41.574Z" }, - { url = "https://files.pythonhosted.org/packages/b2/f9/3243191dc9937e00756c8bc1b0272fed8f23758e43df2a3b46f533e5090f/msgpack-1.2.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:196300e7e5d6e74d50f1607ab9c06c4a1484c383cd22defd727902591f7e8dde", size = 426375, upload-time = "2026-06-18T16:13:42.936Z" }, - { url = "https://files.pythonhosted.org/packages/23/c7/1693111db9944ba4ad4b67a1e788400d78a0b6af7a6523dc7e4e58f8274b/msgpack-1.2.1-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:575957e79cd51903a4e8495a242442949641e08f1efd5197b43bebd3ea7682b4", size = 380495, upload-time = "2026-06-18T16:13:44.306Z" }, - { url = "https://files.pythonhosted.org/packages/3e/2b/92f86956a0c13e8662f7e2ad630c4eb4db07497b967589bd5245e018b2c1/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8c2ed1e48cc0f460bf3c7780e7137ff21a4e18433451916f2442c1b21036cd7d", size = 410897, upload-time = "2026-06-18T16:13:45.629Z" }, - { url = "https://files.pythonhosted.org/packages/da/ea/1479f72d200313a76fc2f823a79d1e07ed052ab7b8a0280640aa7b95de42/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:5f6277e5f783c36786a145e0247fc189a03f35f84b251646e53592d2bc12b355", size = 378519, upload-time = "2026-06-18T16:13:46.998Z" }, - { url = "https://files.pythonhosted.org/packages/f5/4d/fa006060ffa1011d32bfae826fe766fe73e02982183601633b7121058ab3/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:f9389552ecf4784886345ead0647e4edc96bee37cbab05b75540f542f766c48c", size = 419815, upload-time = "2026-06-18T16:13:48.205Z" }, - { url = "https://files.pythonhosted.org/packages/2f/e1/aab6c946570496b78e67804721f3d5e2d62a93081b9b37df77764ef56347/msgpack-1.2.1-cp314-cp314t-win32.whl", hash = "sha256:c1c79a604a2969a868a78b6ebd27a887e00c624f14f66b3038e0590cb23332d1", size = 70914, upload-time = "2026-06-18T16:13:49.385Z" }, - { url = "https://files.pythonhosted.org/packages/13/0a/e608956488a2af014cfe6e3d665e090b8ee42aa14b07f8f95b8880d66b09/msgpack-1.2.1-cp314-cp314t-win_amd64.whl", hash = "sha256:f12038a35fabd52e56a3547bab42401af49a45caa6dd00b34c44de235bc93ee2", size = 77999, upload-time = "2026-06-18T16:13:50.467Z" }, - { url = "https://files.pythonhosted.org/packages/d2/8a/27e2e57055176e366a46b85d02d68e7a5bcfbdd8474c9706375d965f24d3/msgpack-1.2.1-cp314-cp314t-win_arm64.whl", hash = "sha256:0adcf06ffde0777c0e1a9b771a2b1c4226ba1bbf748c8efcc02fcdeca3299107", size = 71160, upload-time = "2026-06-18T16:13:51.498Z" }, + { url = "https://files.pythonhosted.org/packages/6d/aa/5b6b09f835791045282dc5d08431db599a5f4743a69fe2f6670045a2cd85/msgpack-1.2.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ec0030361cc861ac699b2ef1c695b741fa145c88f8667fa3d7e3f73deeb648a3", size = 90927, upload-time = "2026-09-29T02:31:28.286Z" }, + { url = "https://files.pythonhosted.org/packages/c9/91/7b288e9133bd1ba92ca0ca4e7f2a4cfc53cf467d99d8d2f57b9939908fac/msgpack-1.2.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:5c1efdd9181cb1b719ee46865f368a927f1c0c65d577798340b1194545b7515a", size = 89798, upload-time = "2026-09-29T02:31:30.028Z" }, + { url = "https://files.pythonhosted.org/packages/71/9b/5c3dbc450d14645dcec987970692d6ab24008cc33d2155474b1d818486f9/msgpack-1.2.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c309a7abae1d14ba29a8bd0ddbd704a5e469d8e9bd9c3dee0e4ff53d7ae01d56", size = 450687, upload-time = "2026-09-29T02:31:32.407Z" }, + { url = "https://files.pythonhosted.org/packages/2b/21/ea60a8fd0d9e0897fce823e9fd9bf6742567784b35c7eee8f4a18a56eb19/msgpack-1.2.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5bf390259cb25a6a1cd197c65810999b811f64cd38683251538bcc5a1e41f7d3", size = 459808, upload-time = "2026-09-29T02:31:34.282Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f7/42140e6afdac8e94bfedae4cfb67ee004b6ad5c4cadd024df42f759bf3b5/msgpack-1.2.3-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:39b6986c19e1f2dfa549d185dba6ccf1de2e4c0ba10d8cfc0048935b1c5f9109", size = 423845, upload-time = "2026-09-29T02:31:35.713Z" }, + { url = "https://files.pythonhosted.org/packages/19/7b/cd54f27b59dfbdc438a12361fbb6798b66d377a978f946bc9512598290e9/msgpack-1.2.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:fcc6800daac4922960f6eeb7a0dda3dd4105e0bf7bce0e83ebc465a78cb7bdba", size = 445608, upload-time = "2026-09-29T02:31:37.65Z" }, + { url = "https://files.pythonhosted.org/packages/57/38/52bc0dc44cc9f7c2339b632f93d02f8badc78cfb0bb070f2a50a51945e53/msgpack-1.2.3-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:968583e956d0427878050b371308c5f8647088732ef3e66a117dbe1192ec91e0", size = 421721, upload-time = "2026-09-29T02:31:39.151Z" }, + { url = "https://files.pythonhosted.org/packages/89/e6/451c9a42274fb2be82d8ba8b76a5219c613e20f8de1da521d10cb758a9ef/msgpack-1.2.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:1d6bcec3dbbdb89ca385d3a73e63ceae7b841fa0d7ca7c676f1a7bfe7fb2cdb8", size = 460430, upload-time = "2026-09-29T02:31:40.843Z" }, + { url = "https://files.pythonhosted.org/packages/57/bb/663e3100327b58caaa5fb66379e557a2717dac08bb586f22f885756bee47/msgpack-1.2.3-cp310-cp310-win32.whl", hash = "sha256:a6b63917d60d6df451f328bd6afba8565e33c4afe1f62ec4ad758b78731c827b", size = 67987, upload-time = "2026-09-29T02:31:42.157Z" }, + { url = "https://files.pythonhosted.org/packages/28/7a/a00d5d7abc5601099260e0d0af8fadc54fbfac2191315aa56eaee3641d9d/msgpack-1.2.3-cp310-cp310-win_amd64.whl", hash = "sha256:4c0780095871ecc49a58b2ff6b1b43b25214704da67646557ca287a3f49fb2dd", size = 75572, upload-time = "2026-09-29T02:31:43.544Z" }, + { url = "https://files.pythonhosted.org/packages/2a/95/b9c651ccb9d720b2e2c8d537954dff528ab869a03bf89598145716db823c/msgpack-1.2.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ec90a9ae3e1169fa1171147340f0e97d941aa19fcd3b34e8339a55933ed042af", size = 90404, upload-time = "2026-09-29T02:31:44.826Z" }, + { url = "https://files.pythonhosted.org/packages/50/cd/fc9e2e367e80f1493e2ec5f610dda558b344eeede296f88976db133e8f2c/msgpack-1.2.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9d7e9cbb0998bbfd363fd9a09c330520d5e9cb323c05b5a1a05865d23ccf2226", size = 89683, upload-time = "2026-09-29T02:31:46.413Z" }, + { url = "https://files.pythonhosted.org/packages/19/9e/1028485c6886c1c117f777cc9b053e541eff0fedb3292dfb1da95040edb5/msgpack-1.2.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6707d2fa2aa1bb5424ea0b05f44ffc989b15ab41a73ff5855bff4944fec7c8ac", size = 465347, upload-time = "2026-09-29T02:31:47.934Z" }, + { url = "https://files.pythonhosted.org/packages/aa/83/800570e6a22376eb8d599920f70aead4779a63611696f567477c4e85a70f/msgpack-1.2.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:382b219de3d436de3baba0f4b0c6d4336e8f5858d0eb047918b13b69a71c6c55", size = 477820, upload-time = "2026-09-29T02:31:49.479Z" }, + { url = "https://files.pythonhosted.org/packages/ab/ff/817e4a2052f848d3fb67726908d6e4e7c19f68ee7c19553a82ce7b0ed415/msgpack-1.2.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:186e6c602b8a9968b8e864c67d622a69279f7d1e55ae25f40e3bff7e815b2b62", size = 436656, upload-time = "2026-09-29T02:31:51.18Z" }, + { url = "https://files.pythonhosted.org/packages/3d/42/040cc55dde6a7d92057baac8d1fc9cfb9f4fd4162900e2ec16dc33917a7d/msgpack-1.2.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:9276ba88891338f2617044429dfd080ae008c9868a25f6f1a7d004a35dc9ac0a", size = 460939, upload-time = "2026-09-29T02:31:53.026Z" }, + { url = "https://files.pythonhosted.org/packages/09/93/4dc007bdef930eed247346773bc0189b710078961d3218d5ee7ba59f322c/msgpack-1.2.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:c942c21a93f36b3a69e828c8945bb72c94dc2ffe488a2086950c812f3edf046c", size = 433608, upload-time = "2026-09-29T02:31:54.981Z" }, + { url = "https://files.pythonhosted.org/packages/c0/97/a1b944046f283ec89445cb2a982c42233b5b07cc630f9be739f4f1d469a3/msgpack-1.2.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:18a6ed513023001b28dcd3ba54966f6bb90a38274ba8d2640464bcab3a1b81d4", size = 477373, upload-time = "2026-09-29T02:31:56.713Z" }, + { url = "https://files.pythonhosted.org/packages/59/79/ab411d0d172743732ab2503f4c32a22dd1a7d1436a6feecbb160e4b6376a/msgpack-1.2.3-cp311-cp311-win32.whl", hash = "sha256:d0238cd05dec9ffbe0de1071df685ba63e30a36ac155285b1a094e727c38cbe9", size = 67514, upload-time = "2026-09-29T02:31:58.267Z" }, + { url = "https://files.pythonhosted.org/packages/63/8d/6f0cb2b84e484e96278455c26870196d025bb0cec312b226a663f1fa9000/msgpack-1.2.3-cp311-cp311-win_amd64.whl", hash = "sha256:30e1522e4173230dca4d9ad896f038f73c0da6c1edd42f4dbad88ac583cf5d46", size = 75850, upload-time = "2026-09-29T02:31:59.449Z" }, + { url = "https://files.pythonhosted.org/packages/aa/25/f99e13a2c1d3f5a1dcaa5aab27f474e8c4358188bbc68ad79fecb0d1aefe/msgpack-1.2.3-cp311-cp311-win_arm64.whl", hash = "sha256:8ca67f77938ea6a3663aa9bd22b3e031f6da84d665be850abab910ee90728dfd", size = 72338, upload-time = "2026-09-29T02:32:00.885Z" }, + { url = "https://files.pythonhosted.org/packages/af/12/4d7c6d6203416d9fbf0f59ebaa805e70fb929b93a41b611bc821ec5964a0/msgpack-1.2.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:89c930aece4e972b208ba589c8410b4167b05e411a5ea2cb25fd96f8bc47ee43", size = 91577, upload-time = "2026-09-29T02:32:02.141Z" }, + { url = "https://files.pythonhosted.org/packages/eb/c7/8576ad39f4ca42ddad26f68eb8621d2d0a60501193d480f504bd9d7f36c4/msgpack-1.2.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:905a189853d6bdb204c7ae5f4ab77fb857448abfff574d3d93c62e2815b24b4f", size = 90027, upload-time = "2026-09-29T02:32:03.508Z" }, + { url = "https://files.pythonhosted.org/packages/0a/3a/aa9c580aea1314529a0f3562461479780b0d254b064f0880956bfbcc74a8/msgpack-1.2.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f3d7b3d0018746b5997dd6b14a1870b07cc4c327d9101145d94a1fc264a51a06", size = 460343, upload-time = "2026-09-29T02:32:04.906Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cf/9c2e4d6c179529d5bf4a64cff76fa581486569e9fbdd35bd98f51cb624bf/msgpack-1.2.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ede33b2892ceb976283e009ad12fa1834cfdf1f9c43ee9c97849fc588d00a618", size = 472998, upload-time = "2026-09-29T02:32:06.69Z" }, + { url = "https://files.pythonhosted.org/packages/7b/41/915c81fe6df2d3cbdb0dece4f1a5cd313e1cd2abd9f501d0f50c0582517e/msgpack-1.2.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:666ef5601ab0e6e345e47febc96aa81143cc932201543480cbb9499164f05ffb", size = 423216, upload-time = "2026-09-29T02:32:08.739Z" }, + { url = "https://files.pythonhosted.org/packages/a2/e7/7dda8b1039abfd9bba4c5068172c67135c9e33089f503512db9226f23c24/msgpack-1.2.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:87cf2ef05ff2f2493ba29fcdaef27e960ca64dacfd13460ae29e6f92e0ed05bb", size = 451218, upload-time = "2026-09-29T02:32:10.517Z" }, + { url = "https://files.pythonhosted.org/packages/16/5b/ce995c1ed4a0522b7f2d034bc2034fd63005f240b945961b70fb56fbaf3d/msgpack-1.2.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:b774ff994d844e541439ac5d2d49a14def4104830c3465e9394c153f86200ffb", size = 422453, upload-time = "2026-09-29T02:32:11.956Z" }, + { url = "https://files.pythonhosted.org/packages/d2/3f/ce191fb87e2650d0166b34c437e499ee4a7f9db9c1eb164f41725eb6160e/msgpack-1.2.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:eaf7e82249837e3aa97297b34a0bb9ff562027381631e057cea6e1367f10b438", size = 469003, upload-time = "2026-09-29T02:32:13.663Z" }, + { url = "https://files.pythonhosted.org/packages/42/35/539123407fe200fb16609c835675496fbeb6017ace9fc93909f0613223ae/msgpack-1.2.3-cp312-cp312-win32.whl", hash = "sha256:7c047250096f9fc19dba26e3d1639b5e7a84114003605c94def667149a70ced1", size = 68303, upload-time = "2026-09-29T02:32:15.02Z" }, + { url = "https://files.pythonhosted.org/packages/6f/4c/331b45f9b86fbda6b9e103244d189068e51f726d8c40021ed66e1f2c415e/msgpack-1.2.3-cp312-cp312-win_amd64.whl", hash = "sha256:3ec409b0d6aa8e9eec6eaf881b893caa215dbe68c5319ca96e8a271d81bb111d", size = 76744, upload-time = "2026-09-29T02:32:16.344Z" }, + { url = "https://files.pythonhosted.org/packages/13/9f/fb572dc42b9fac06c7ea848aaee6e140d84469743bd1402bc07089fc4566/msgpack-1.2.3-cp312-cp312-win_arm64.whl", hash = "sha256:59612b4ed48a04cf024584218e813562f3b30a3bafa5f55abe300b15da314751", size = 71580, upload-time = "2026-09-29T02:32:17.617Z" }, + { url = "https://files.pythonhosted.org/packages/1f/8b/3824d65e912e925d09ce30d9130fa9970d6d2855d7888b13639a6604967f/msgpack-1.2.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:21bfa4d2aa0b04c1806ef778a1199e9e53ea2441bcbf284420a32083896320b8", size = 91728, upload-time = "2026-09-29T02:32:18.949Z" }, + { url = "https://files.pythonhosted.org/packages/05/e6/df7f2c9ebb94760113debbcea2bd3afe5fdab88a4f7bec1b618755517460/msgpack-1.2.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:db84203b13aecc222f465061397fdd5b53b7ae73d2c95ffc1c8dc5be0153a709", size = 89955, upload-time = "2026-09-29T02:32:20.224Z" }, + { url = "https://files.pythonhosted.org/packages/08/6a/e5fc57136e8bacccb2b39627dea2cd546540a06181e22fe6db90e15b3ae4/msgpack-1.2.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e0d7950ca3c1bbae291d0552dd3bb2792fc680629c4c0d44e47e5bab969f3ca", size = 454930, upload-time = "2026-09-29T02:32:21.771Z" }, + { url = "https://files.pythonhosted.org/packages/b0/30/c394d37898db9212d1693456cdf363c7e1a097d0b63e10664007f3df3ec1/msgpack-1.2.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:07c9733089d1b176c3dd2f7fa268452f9d5d784d076473499d754a58e8d1fbbb", size = 466866, upload-time = "2026-09-29T02:32:23.742Z" }, + { url = "https://files.pythonhosted.org/packages/4a/c8/1e4ddf6f6b829b3ee6c530c79dfae89cb609d2b0eedb5e0ae716851c52d1/msgpack-1.2.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f24a43b3560e20f825b807fe1e874bd73d53abaf8bbdcf258a6eb152cddbc1f5", size = 418715, upload-time = "2026-09-29T02:32:25.262Z" }, + { url = "https://files.pythonhosted.org/packages/11/a5/f460ba6d7a12d4301002f3efbb8f841e8bdc9c5fc98d771689677a352885/msgpack-1.2.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6576f348ed6cc4f31db6fd915a8e94245f042f50eae08d48732425e70638ea37", size = 446489, upload-time = "2026-09-29T02:32:26.988Z" }, + { url = "https://files.pythonhosted.org/packages/49/23/adface88db909bed321c85dd673655152d4a514c67e1f0800eb51c777d07/msgpack-1.2.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:cd5a9f9f86a52c24713679aa2631956835f3842512964ff93f736ff76f1f530d", size = 416998, upload-time = "2026-09-29T02:32:28.606Z" }, + { url = "https://files.pythonhosted.org/packages/36/00/5bb3a239ccfc3763c4d0fa49b13b1b7010b00182c499ab3c1fecfe6294bc/msgpack-1.2.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f9ddd28d3e9bbc602a9dced1591882c7fb9ab776eef8837da2c326fde19e2853", size = 463288, upload-time = "2026-09-29T02:32:30.375Z" }, + { url = "https://files.pythonhosted.org/packages/29/8c/456df77f00d701df9d6980ffb80291bce6e4e2e112e25a4dfae216f0715a/msgpack-1.2.3-cp313-cp313-pyemscripten_2025_0_wasm32.whl", hash = "sha256:62cc1a4ef0e553bac32c8342e1f04834aca7de276b92744eb7307db77759b890", size = 53347, upload-time = "2026-09-29T02:32:31.867Z" }, + { url = "https://files.pythonhosted.org/packages/9d/22/ce780be666f89b77cdb855daa9ec62e87bb7f69e9f403e4a5d83a2b2208f/msgpack-1.2.3-cp313-cp313-win32.whl", hash = "sha256:d2f9c4f85e47a44d26d5baf3b041eef23436e224d44eed273f01bd8a12048d9f", size = 68258, upload-time = "2026-09-29T02:32:33.163Z" }, + { url = "https://files.pythonhosted.org/packages/51/06/c3def9bc4db283103c5901b302ee2a4305cb1e69729244f94d9bd8f8e8e7/msgpack-1.2.3-cp313-cp313-win_amd64.whl", hash = "sha256:bb89b5dc30469c84bbf8684826eb851d82412ca95690e111b9ac5e8fb343961a", size = 76569, upload-time = "2026-09-29T02:32:34.412Z" }, + { url = "https://files.pythonhosted.org/packages/12/9f/cef344073858b80adb92d6ea342e20b0eae7a8f6fe70281b69cf03707270/msgpack-1.2.3-cp313-cp313-win_arm64.whl", hash = "sha256:471e12a6a42498a31490c206e0069e343b6a7c35db540be73a879eb06f5be047", size = 71530, upload-time = "2026-09-29T02:32:35.892Z" }, + { url = "https://files.pythonhosted.org/packages/3f/8e/f777f74e38731c428857933c8011596f2d2f3160c821152f23b6ffba862f/msgpack-1.2.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:3a31905206722103a84c1f72633fe30692cff6732c9d262e09a27dbc468797c8", size = 92042, upload-time = "2026-09-29T02:32:37.464Z" }, + { url = "https://files.pythonhosted.org/packages/a0/71/551608543ee5d590f7e8d522267665d6d9946866ad2a2a70a770f7c70793/msgpack-1.2.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:3372475211a9ce1a23acefe512cb3e121d18c95dc74ed56cb1819ef40836ebf4", size = 90578, upload-time = "2026-09-29T02:32:38.883Z" }, + { url = "https://files.pythonhosted.org/packages/ea/11/6d78ce5a9a58bf9ba7b1b6a8f649173b030e6770c8019cf330b91825ee5d/msgpack-1.2.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9324c54995641c3d1f92a9d55093c8cde0ffa2fbc87a467a688ef60428393220", size = 454352, upload-time = "2026-09-29T02:32:40.34Z" }, + { url = "https://files.pythonhosted.org/packages/3d/08/feb9a196269ba7809f44f9117d9e4a601c41c313f6144fd0c337293a5488/msgpack-1.2.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d8ef3a66e4b52d2d7fdd90df2984670124b2ff7546d76bb25dcf68ef47f7df58", size = 462562, upload-time = "2026-09-29T02:32:42.176Z" }, + { url = "https://files.pythonhosted.org/packages/f5/77/3a674f366def24140b103d1ffd4fd27b3d912a13e47da67422afa16bebb3/msgpack-1.2.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:902f3490db0e07a7d40b48536a85c9b28fbf1397e7e1658a45a55f958e303620", size = 418134, upload-time = "2026-09-29T02:32:43.693Z" }, + { url = "https://files.pythonhosted.org/packages/48/82/944e71f280577490d99a3951cbce21aa4cbe04e7ab42cb373fd668af883c/msgpack-1.2.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:8e51eca14fbb65c4e0a5a9657346962bd3dca78c08e04e3d4dee70ef48687d30", size = 445937, upload-time = "2026-09-29T02:32:45.739Z" }, + { url = "https://files.pythonhosted.org/packages/b1/ec/feddd629c4a3edf1395313680450c525086cceab56dec0d4de9da9ccb618/msgpack-1.2.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:f42f146752eedb6765f07dcc04d72dab0a25779ec8d4a88c0085263ce114f22c", size = 416450, upload-time = "2026-09-29T02:32:47.558Z" }, + { url = "https://files.pythonhosted.org/packages/e4/59/263a10f8c4613ba0713f48cbda7695ac8dd6d6fab2fcbc9168f03f23a94d/msgpack-1.2.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:0ed5823c4efc20fe87d3530665f40ec18a002be003114814c21235cc8d256207", size = 459546, upload-time = "2026-09-29T02:32:49.145Z" }, + { url = "https://files.pythonhosted.org/packages/1e/21/addcfa1e583cfc8a22fbdc57526621b5decd7ad676ae12e9150b7be1be5d/msgpack-1.2.3-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:2487453ca1b6104442c6442f9a1a8fee1fe8f428a70d99d4cba799108b304150", size = 53462, upload-time = "2026-09-29T02:32:50.708Z" }, + { url = "https://files.pythonhosted.org/packages/8d/2c/3cb5c8524a1335ee27ca952c7ab78d375a16fea8e18ae3767ba0c880416c/msgpack-1.2.3-cp314-cp314-win32.whl", hash = "sha256:6df430419f2338cb71e4a34d6e64f83c88ccd321f91f40ba4513400b36d864ec", size = 70294, upload-time = "2026-09-29T02:32:52.037Z" }, + { url = "https://files.pythonhosted.org/packages/23/f9/9172ff3cdb85d160ad06df5e2708a5fce7682982a5eee8d31869b9f69d2e/msgpack-1.2.3-cp314-cp314-win_amd64.whl", hash = "sha256:84a6616d396ec1bc18a1e83e67c96a393ec35dfe5e17434a5be7b9aa0fe988ab", size = 77778, upload-time = "2026-09-29T02:32:53.429Z" }, + { url = "https://files.pythonhosted.org/packages/04/e8/b4c23178bcf605ae17cec48a75530dd69d49b0a5a6f5f4df5c47d59f746e/msgpack-1.2.3-cp314-cp314-win_arm64.whl", hash = "sha256:7a003b02c6ee2eea6dfe0bb08818631e3597e69f0131f2a8250488a1cc553290", size = 73794, upload-time = "2026-09-29T02:32:54.763Z" }, + { url = "https://files.pythonhosted.org/packages/66/b1/92704be352c4f428b7e0a0e0fb210cb1aa2b1c42c102b8dc22d34b82fac0/msgpack-1.2.3-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:ccea05b5542f6d283fef3f0a8e93a7f0be90af0ddeeef84c25c0216ba76dcae1", size = 93721, upload-time = "2026-09-29T02:32:56.342Z" }, + { url = "https://files.pythonhosted.org/packages/49/78/9c91f1e86cadcbc100b3780fd429c3715648704032a612e77a00646ebe79/msgpack-1.2.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b1631e12fe572e181cd77e831f69335d6cd5278eac22e3db3f33cf264ac2ac18", size = 94256, upload-time = "2026-09-29T02:32:58.056Z" }, + { url = "https://files.pythonhosted.org/packages/91/4d/270f9725921ae88a29d37a774a77ac24f0ef1411fc960a63f5a4665e81b4/msgpack-1.2.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e54394b7dbe2e12ab032d9d21feef7bb61a90a150a2623633ba3781ba69dcb1f", size = 471673, upload-time = "2026-09-29T02:32:59.886Z" }, + { url = "https://files.pythonhosted.org/packages/48/b8/eaa8d930f72dc1d1dd79511dc2ccf965922b059f2f0ed3b30aebac8c4b11/msgpack-1.2.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:63bb7448a1e9111319ae2430c09a5596140c160422830d6271bc75730ff2ff9a", size = 466257, upload-time = "2026-09-29T02:33:01.517Z" }, + { url = "https://files.pythonhosted.org/packages/5b/5a/97adc805037bc7e24c4e2f711bbcd3b28be8ec9aea3e778f18208cfbdb46/msgpack-1.2.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:382bc88fe90f29f5ac8a0b65c7046ff255356f2f2f3186c30e370215736fa1dc", size = 418484, upload-time = "2026-09-29T02:33:03.402Z" }, + { url = "https://files.pythonhosted.org/packages/0d/7e/1c53302606fe436ab48ba539ebafafe4a6a9efe12c4f04dc7eb36912d93e/msgpack-1.2.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c77e27790ad72989db783d5303825fba0b71550f00a490efba35cde7dc4b719f", size = 454064, upload-time = "2026-09-29T02:33:04.977Z" }, + { url = "https://files.pythonhosted.org/packages/00/2d/9ee0170f638907b396c15c6cd26b3e54f869159efc6206683acfd8f696e1/msgpack-1.2.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:700bc0fc9e968a292b9137ee70e7a012f7e115bf0107ce45e3a88202788dfc1e", size = 417901, upload-time = "2026-09-29T02:33:06.489Z" }, + { url = "https://files.pythonhosted.org/packages/cc/d2/905c84490a75cd15a27065407cd085d201f7d392e1e0411f49f03fd31ade/msgpack-1.2.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5bd5f91ea75c45cafcc5433ba8fae59b708b736ec178d2441c40c499e9e079db", size = 459896, upload-time = "2026-09-29T02:33:08.361Z" }, + { url = "https://files.pythonhosted.org/packages/37/cd/4ce5809b9ab3b114d7cca64863e436820fa1614b49d55ccb93d49824ac2d/msgpack-1.2.3-cp314-cp314t-win32.whl", hash = "sha256:7995a7c6a62a1d6e7df211b4a16de513bd99fd053525050a319f80f44fb8015e", size = 75983, upload-time = "2026-09-29T02:33:10.023Z" }, + { url = "https://files.pythonhosted.org/packages/8a/31/853bb580744c24be0dbd8b090c3e6987dce466a1fc840fe50c0ac2ef9044/msgpack-1.2.3-cp314-cp314t-win_amd64.whl", hash = "sha256:bfe7d5b62cbe7aa664f0b3e2c49077f10fcdd06183d3014f8271ff3c5edbfbf9", size = 83757, upload-time = "2026-09-29T02:33:11.441Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/9f1b2ee484414eef9e21ee2b2b23b482bb71433ab9bac1da03cbda15ebf5/msgpack-1.2.3-cp314-cp314t-win_arm64.whl", hash = "sha256:1f585407f740a9eac04a3bb82c61d68a0ea78f90e29e670bfb086b9ce3a518dd", size = 78128, upload-time = "2026-09-29T02:33:13.063Z" }, + { url = "https://files.pythonhosted.org/packages/47/b8/50db4235407c3802f622b4ccdf65c6fe1e48d3c3eab6981fa6a9a5e53f11/msgpack-1.2.3-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:13221a6c81ebb8e43ea63a7251c35d54e4175cea37ebf3a62e911bdf42562a3c", size = 92111, upload-time = "2026-09-29T02:33:14.476Z" }, + { url = "https://files.pythonhosted.org/packages/15/56/50cf2a45c6163edafd737e2fd555103a26ce6748e1e241fb56ed445ea835/msgpack-1.2.3-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:0955b9000725573d1457c1676944b370dd9643c8d18f25bda5ac72913f850949", size = 90583, upload-time = "2026-09-29T02:33:15.924Z" }, + { url = "https://files.pythonhosted.org/packages/2a/fd/8cc02f767c3bc94d2649c954d28dea935ce9398eb9c93ce2444bb9474cc1/msgpack-1.2.3-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0c91762c48cd686dc9cf2b142c0bc544083952de32f5853d6624c956e54b85e5", size = 454751, upload-time = "2026-09-29T02:33:17.475Z" }, + { url = "https://files.pythonhosted.org/packages/80/c9/ddb896767808e3e022453d8dfae26fd52ed404b0aa6fb7f752d39c040208/msgpack-1.2.3-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1f4ae8bd4ad9ba085fde95e95d055a896d19210238a4199a771a3cf36dceed49", size = 463597, upload-time = "2026-09-29T02:33:19.309Z" }, + { url = "https://files.pythonhosted.org/packages/4d/a5/e7c261abf75783c07dcac89951cb31dd0c123bf02fbdeda0c67303e698d8/msgpack-1.2.3-cp315-cp315-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7013534a7163aa4f213c4d9864f1a8a7555daac6fcd48f699a198e29b436bfab", size = 422661, upload-time = "2026-09-29T02:33:21.093Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8e/466d5133f9e1c2e232e15e304f715b62f6f0e28332d18e37d975fe174315/msgpack-1.2.3-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:6a834097144aabe948b8ca9020a833e8026f7d0abbd0ec54bc7e50f45a8ce012", size = 445188, upload-time = "2026-09-29T02:33:22.877Z" }, + { url = "https://files.pythonhosted.org/packages/d4/b4/33e7ad987ee2f4b3d449a6cbf28f574ed222987ca7f65ad277072646ac5e/msgpack-1.2.3-cp315-cp315-musllinux_1_2_riscv64.whl", hash = "sha256:d31864ba3933a589b6a00249f89c0eb422197f49128fc10da550e57e9cb0f377", size = 420451, upload-time = "2026-09-29T02:33:24.485Z" }, + { url = "https://files.pythonhosted.org/packages/34/2c/9d8be0d6c16e7e6131cd7da20257dd3da65473e3e6df0c00572fb10a195c/msgpack-1.2.3-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:e15f70588f4db8cd10df0930145b186de70feb9db51710cd378b1399009655bd", size = 460624, upload-time = "2026-09-29T02:33:26.063Z" }, + { url = "https://files.pythonhosted.org/packages/6a/e7/3a04783582c6f44f398cbfcf5f07a111192126ec4e63edf7f5640143bf64/msgpack-1.2.3-cp315-cp315-pyemscripten_2026_5_wasm32.whl", hash = "sha256:b949cc25e4a09252cbcc54e66e507de914d0e94a3a7039bd54c299bf7037c098", size = 53474, upload-time = "2026-09-29T02:33:27.83Z" }, + { url = "https://files.pythonhosted.org/packages/68/fb/db07359851644e258609d84f8e4fe0030ef448c108e20afe73f2a3bf539c/msgpack-1.2.3-cp315-cp315-win32.whl", hash = "sha256:8ec7a1d49ca6c2569d722ab5ec86e90089b0713900aa31905b47b4c4d9e78ce0", size = 70344, upload-time = "2026-09-29T02:33:29.382Z" }, + { url = "https://files.pythonhosted.org/packages/5b/e4/cf5584d2f2a2e4465d5896a855a3e75a34a20ab172360b3d42ad862dd1ce/msgpack-1.2.3-cp315-cp315-win_amd64.whl", hash = "sha256:79dfa38faf92f804aa61beec140d70b18418e1dde1778dbb77a87a4cce85aa8a", size = 77800, upload-time = "2026-09-29T02:33:30.941Z" }, + { url = "https://files.pythonhosted.org/packages/63/f9/518ad4e8a580027b507eafdd26de7aae661a714e43d7c111c212482e4a1b/msgpack-1.2.3-cp315-cp315-win_arm64.whl", hash = "sha256:ed899d73a22f286a72bd9528d63f2ab3030dbad8bf1527fc249319a50d61fb9d", size = 73871, upload-time = "2026-09-29T02:33:32.406Z" }, + { url = "https://files.pythonhosted.org/packages/a4/79/254d4c9ad642b2a3ba84e646787892b34cc815eb36c9976f67a1c4f38515/msgpack-1.2.3-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:f56fba61b2516be7917cb00151f0d060b5b21184e3499bb57f0f7d9259bea124", size = 93370, upload-time = "2026-09-29T02:33:33.87Z" }, + { url = "https://files.pythonhosted.org/packages/3d/6f/5a2ba167646a25e84eaa8894e12935351e4331b80c28a9237ce6fe8d375f/msgpack-1.2.3-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:69ad12cedb674c73527bed869cddb42b742cac79a207a614202a4abaa24ea173", size = 93959, upload-time = "2026-09-29T02:33:35.503Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a1/2b44612e55f7cf5d5e4b580294959b4429bbbcb1991177888e3e18668137/msgpack-1.2.3-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:db9fb67a3a2e75247bae569d34ebb5ff61c0448a4f0d6dbf991dae68af39b007", size = 467921, upload-time = "2026-09-29T02:33:37.023Z" }, + { url = "https://files.pythonhosted.org/packages/0b/6e/3309798ed1c11d7fcfdc7b946642685b0ff1588477925bc0d26bee7dcaae/msgpack-1.2.3-cp315-cp315t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2574ef81c1c8c38b10e330f3f9406fd09198a776b002030fafcf8e7647e9e06e", size = 467310, upload-time = "2026-09-29T02:33:38.799Z" }, + { url = "https://files.pythonhosted.org/packages/6f/79/9c799f489fa4146de4e00cfe9fee17afe33d8012f88ddffffea94f7c4700/msgpack-1.2.3-cp315-cp315t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fafc3b8898b432b841d30a61082c599fa7f4d06885f9dc58ad72259e12059fa6", size = 420178, upload-time = "2026-09-29T02:33:40.781Z" }, + { url = "https://files.pythonhosted.org/packages/94/c6/5850dc9cafcd2ea315692e65db0e222d20923dd55f44adf35061003de27e/msgpack-1.2.3-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:a393e428f6ffb0dcb73308c1fff5593041c16ff42da66e5bac8a83a6107a54b0", size = 450248, upload-time = "2026-09-29T02:33:42.366Z" }, + { url = "https://files.pythonhosted.org/packages/a9/d2/b4c806e3497fe21f0b353568266aec14ff735d092aea672de7b2955db03f/msgpack-1.2.3-cp315-cp315t-musllinux_1_2_riscv64.whl", hash = "sha256:d1c1e8989a855b7f1f2a64ec4a80b23a631822903952770813857b2e4f460471", size = 418431, upload-time = "2026-09-29T02:33:44.178Z" }, + { url = "https://files.pythonhosted.org/packages/b0/f5/f4ecc3ddac4d551bf2f3cdb283ec546dcc826fe7c500074be61aa273e08a/msgpack-1.2.3-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:e0bd394e999949c814f7912284243298de1b5a17b6a3dcb6cc8a79b156ffc4fa", size = 457543, upload-time = "2026-09-29T02:33:45.978Z" }, + { url = "https://files.pythonhosted.org/packages/a4/69/1c821d8386fae5cecc5fcaacf3de3947ff0a23f16bb481b5532b5868372a/msgpack-1.2.3-cp315-cp315t-win32.whl", hash = "sha256:3d4c807ed050fe3ddbea5ba7e9f63d7136871ce42861be1f50ff739f0e91047a", size = 75820, upload-time = "2026-09-29T02:33:47.596Z" }, + { url = "https://files.pythonhosted.org/packages/68/9e/41e2f7343a3764a9c1fb10c79f9a6a05db9df93dedd76401d1b511f5a685/msgpack-1.2.3-cp315-cp315t-win_amd64.whl", hash = "sha256:5f304123b90e8b2e49867981b7f6061612c39f50cca51ee88de007c084cf68d3", size = 83345, upload-time = "2026-09-29T02:33:49.325Z" }, + { url = "https://files.pythonhosted.org/packages/80/cd/0c3aa439bc7a7bf24684fef3a0ad776cba170e18ed94445e723bce42fce7/msgpack-1.2.3-cp315-cp315t-win_arm64.whl", hash = "sha256:f41ca154b7737b11893cdce3c78c61d703398a1cd54d4297bdad908392338a8e", size = 77572, upload-time = "2026-09-29T02:33:50.729Z" }, ] [[package]] @@ -8743,7 +8814,7 @@ crypto = [ [[package]] name = "pylint" -version = "4.0.5" +version = "4.0.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "astroid" }, @@ -8755,9 +8826,9 @@ dependencies = [ { name = "tomli", marker = "python_full_version < '3.11'" }, { name = "tomlkit" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e4/b6/74d9a8a68b8067efce8d07707fe6a236324ee1e7808d2eb3646ec8517c7d/pylint-4.0.5.tar.gz", hash = "sha256:8cd6a618df75deb013bd7eb98327a95f02a6fb839205a6bbf5456ef96afb317c", size = 1572474, upload-time = "2026-02-20T09:07:33.621Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/5e/d5ec2c861010cf8577a3e5baaf2405fdf1004f1657d9245c4efb006c1dfd/pylint-4.0.9.tar.gz", hash = "sha256:6f1305780229ae720e89b8aff2b3ff857cbb268fbab98c8b586960267667ad34", size = 1604187, upload-time = "2026-09-23T15:26:43.144Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d5/6f/9ac2548e290764781f9e7e2aaf0685b086379dabfb29ca38536985471eaf/pylint-4.0.5-py3-none-any.whl", hash = "sha256:00f51c9b14a3b3ae08cff6b2cdd43f28165c78b165b628692e428fb1f8dc2cf2", size = 536694, upload-time = "2026-02-20T09:07:31.028Z" }, + { url = "https://files.pythonhosted.org/packages/66/46/12e94ae1f9db89152223bdb17920b12c79e94ac51e19dce1da9fb63627b0/pylint-4.0.9-py3-none-any.whl", hash = "sha256:c4ba5bbaf1e70ecf05f08cd7d6091ea04824b79a2b2d9f5b10dc25fc34cb1124", size = 542068, upload-time = "2026-09-23T15:26:41.626Z" }, ] [[package]] diff --git a/vscode-extension/package-lock.json b/vscode-extension/package-lock.json index 81e9425ef4e..5f6cf7edf4e 100644 --- a/vscode-extension/package-lock.json +++ b/vscode-extension/package-lock.json @@ -2920,9 +2920,9 @@ } }, "node_modules/qs": { - "version": "6.16.0", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.16.0.tgz", - "integrity": "sha512-h6fhOIaRrID2CbEY2fqs+7t+UXZo+MLAnU5gRIq85uFtdiUPCdsApMlHhXogKVM4HM2DVbIjGNTTYH2OcmP1vA==", + "version": "6.17.0", + "resolved": "https://registry.npmjs.org/qs/-/qs-6.17.0.tgz", + "integrity": "sha512-5LK5BZ2XDHMvcUJaf0CQ5HL3aVgKbpklvByC7qLSxYkfE8COEEbJg5xCHWCTqzqi0Gs4Jn5Mv5v3R38W6edscA==", "dev": true, "license": "BSD-3-Clause", "dependencies": {