diff --git a/.circleci/config.yml b/.circleci/config.yml index 0d442e32b0a..1edf22578ff 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1291,7 +1291,7 @@ jobs: name: Install MCP SDK1 peer command: | uv venv --python 3.12 .venv-mcp-peer - uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' + uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - when: condition: diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index d2ab3478397..a33695d0dbe 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -33,7 +33,7 @@ while IFS= read -r file || [ -n "$file" ]; do *) has_backend=true ;; esac case "$file" in - model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json) + model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json) has_cost_map=true ;; tests/unit/proxy/*) : ;; *) outside_cost_map_set=true ;; diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index eefda48b682..c593acbe8fb 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -139,6 +139,7 @@ jobs: test-path: |- tests/unit/batch_completion tests/unit/decisions + tests/unit/extraction tests/unit/litellm_core_utils tests/unit/systemone python-version: "3.12" @@ -433,6 +434,13 @@ jobs: - shard: rust-bridge-harness test-path: >- + tests/test_litellm_rust/test_dispatch.py + tests/test_litellm_rust/test_public_entries.py + tests/test_litellm_rust/chat_completions/test_dispatch.py + tests/test_litellm_rust/embeddings/test_dispatch.py + tests/test_litellm_rust/messages/test_dispatch.py + tests/test_litellm_rust/ocr/test_dispatch.py + tests/test_litellm_rust/responses/test_dispatch.py tests/unit/compression tests/unit/harness tests/unit/rust_bridge diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index 932da84e5d7..c1aeaabc5d2 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -4,8 +4,7 @@ Every pull request whose diff against its merge base touches one of the two cost checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always -get the file checks and additionally may only touch those two files or the legacy schema -path and may only add or update models. +get the file checks and additionally may only touch those two files and may only add or update models. """ from __future__ import annotations @@ -21,7 +20,6 @@ from typing import Final COST_MAP_PATH: Final = "model_prices_and_context_window.json" BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" -LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH) BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_" SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"}) @@ -110,7 +108,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str *( f"bot PRs may only change the cost map files, not {path}" for path in changed_files - if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH + if path not in GUARDED_PATHS ), *(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map), *(f"bot PRs may not remove fields: {ref}" for ref in removed_fields), 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 bd9440d4754..df1f263a947 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, ManagedObjectRow, + row_metadata, ) import litellm @@ -90,11 +91,6 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li return TeamRepository(prisma_client).table -def _row_metadata(row: object) -> dict[str, object]: - metadata: Final = getattr(row, "metadata", None) - return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column - - class CheckBatchCost: def __init__( self, @@ -1063,7 +1059,9 @@ class CheckBatchCost: if (has_key_row and key_row is None) or (team_id and team_row is None): return True settings: Final = get_dynamic_logging_metadata( - UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id), + UserAPIKeyAuth( + metadata=dict(row_metadata(key_row)), team_metadata=dict(row_metadata(team_row)), team_id=team_id + ), proxy_config, ) callback_vars: Final = (settings.callback_vars if settings is not None else None) or {} diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 1fa461d7ab7..07048cf2c2f 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -7,6 +7,12 @@ attribution, priced against the deployment that served it. The poll's status che with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the managed files hook, so both stay free. The spend event carries the id the client holds, the row's unified_object_id, so its spend log is found under that id with the _background_cost suffix. +The spend event reaches the creating key's or team's own logging callbacks the way that creator's +requests do. Its message content is redacted when the creator's key or team logging turns message +logging off, else as the proxy is set, the way the batch poll decides it. The row is claimed before +the creator is read, so reads that lose the race skip the lookups. A creator whose lookup fails, or +whose logging settings the logger rejects, is still billed, through the proxy's own callbacks only +and with its message content redacted. A read that finds no row, because the create could not store one, stores it under the reading key and bills it the same way, so a response the create left untracked is still billed once. A streamed background create is free on its own; its stream stores the row under the creating key and bills @@ -29,6 +35,7 @@ from typing import TYPE_CHECKING, Final, cast from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, ManagedObjectRow, + ResponseCreator, ) from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -54,11 +61,13 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.proxy.proxy_server import ProxyConfig from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.prisma_protocols import TableActions from litellm.router import Router _STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) +_CALLBACK_LIST_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) class _RoutedModelInfo(BaseModel): @@ -73,6 +82,8 @@ class _RoutedDeployment(BaseModel): _ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...]) _RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses" +_CALLBACK_LIST_FIELDS: Final = frozenset({"success_callback", "failure_callback", "litellm_disabled_callbacks"}) +_REJECTED_CREATOR_LOGGING: Final[Mapping[str, object]] = MappingProxyType({"turn_off_message_logging": True}) @dataclass(frozen=True, slots=True) @@ -89,6 +100,7 @@ class BilledBackgroundResponse: custom_llm_provider: str | None litellm_params: Mapping[str, object] metadata: Mapping[str, object] + creator_logging: Mapping[str, object] @dataclass(frozen=True, slots=True) @@ -107,11 +119,51 @@ def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse: return unpriced -async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +def _string_callbacks(callbacks: object) -> tuple[str, ...]: + if not isinstance(callbacks, list): + return () + return tuple( + callback for callback in _CALLBACK_LIST_ADAPTER.validate_python(callbacks) if isinstance(callback, str) + ) - logging_obj: Final = LiteLLMLogging( + +def _running_proxy_config() -> "ProxyConfig": + from litellm.proxy.proxy_server import proxy_config + + return proxy_config + + +def creator_logging_fields(creator: ResponseCreator, proxy_config: "ProxyConfig | None") -> Mapping[str, object]: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields + + if not creator.settings_read: + return _REJECTED_CREATOR_LOGGING + try: + return key_and_team_logging_fields( + user_api_key_dict=UserAPIKeyAuth( + metadata=dict(creator.key_metadata), + team_id=creator.team_id, + team_metadata=dict(creator.team_metadata), + ), + proxy_config=proxy_config if proxy_config is not None else _running_proxy_config(), + ) + except Exception as e: + verbose_proxy_logger.error( + f"CheckResponsesCost: could not read the creator's logging settings, so only the proxy's " + f"callbacks get the billed event: {e}" + ) + return _REJECTED_CREATOR_LOGGING + + +def _billed_logging_obj(billed: BilledBackgroundResponse, creator_logging: Mapping[str, object]) -> "LiteLLMLogging": + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.utils import split_dynamic_success_callbacks + + sync_callbacks, async_callbacks = split_dynamic_success_callbacks( + _string_callbacks(creator_logging.get("success_callback")) + ) + return LiteLLMLogging( model=billed.model, messages=[{"role": "user", "content": ""}], stream=False, @@ -119,7 +171,27 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi start_time=datetime.now(), litellm_call_id=str(uuid.uuid4()), function_id=str(uuid.uuid4()), + dynamic_success_callbacks=list(sync_callbacks), + dynamic_async_success_callbacks=list(async_callbacks) or None, + kwargs={key: value for key, value in creator_logging.items() if key not in _CALLBACK_LIST_FIELDS}, ) + + +def _creator_or_proxy_logging_obj(billed: BilledBackgroundResponse) -> "LiteLLMLogging": + try: + return _billed_logging_obj(billed, billed.creator_logging) + except Exception as e: + verbose_proxy_logger.error( + f"CheckResponsesCost: the creator's logging settings were rejected, so only the proxy's " + f"callbacks get the billed event: {e}" + ) + return _billed_logging_obj(billed, _REJECTED_CREATOR_LOGGING) + + +async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj: Final = _creator_or_proxy_logging_obj(billed) logging_obj.update_environment_variables( litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)}, optional_params={}, @@ -187,6 +259,7 @@ class CheckResponsesCost: prisma_client: "PrismaClient", llm_router: "Router", log_billed_response: BilledResponseLogger = log_billed_background_response, + proxy_config: "ProxyConfig | None" = None, ): from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -196,6 +269,7 @@ class CheckResponsesCost: self.llm_router: Router = llm_router self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost") self._log_billed_response: Final = log_billed_response + self._proxy_config: Final = proxy_config def _resolve_deployment(self, response_id: str) -> Deployment | None: model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) @@ -307,11 +381,12 @@ class CheckResponsesCost: data={"status": job.status or PENDING_RESPONSE_STATUSES[0]}, ) - async def _billed_response( + def _billed_response( self, job: ManagedObjectRow, target: _PollTarget, response: ResponsesAPIResponse, + creator: ResponseCreator, ) -> BilledBackgroundResponse: deployment: Final = target.deployment return BilledBackgroundResponse( @@ -323,13 +398,14 @@ class CheckResponsesCost: litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}), metadata=MappingProxyType( { - **(await self._creator_attribution.response_metadata(job, job.unified_object_id)), + **creator.metadata, INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, **target.model_metadata, **(_deployment_metadata(deployment) if deployment is not None else {}), **self._requested_model_group(target, job.team_id), } ), + creator_logging=creator_logging_fields(creator, self._proxy_config), ) def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]: @@ -354,7 +430,8 @@ class CheckResponsesCost: verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping") return None try: - charge: Final = await self._log_billed_response(await self._billed_response(job, target, response)) + creator: Final = await self._creator_attribution.response_creator(job, job.unified_object_id) + charge: Final = await self._log_billed_response(self._billed_response(job, target, response, creator)) except BaseException: await self._release_billing_claim(job) raise diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index 2e2c07fd544..a4a40a7a3d7 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -10,12 +10,16 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime -from typing import TYPE_CHECKING, Final, Protocol, cast +from typing import TYPE_CHECKING, Final, Protocol, TypeGuard, cast from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.constants import CLI_SESSION_KEY_PREFIX +from litellm.repositories.table_repositories import ( + DeletedVerificationTokenRepository, + DeprecatedVerificationTokenRepository, +) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository @@ -70,6 +74,7 @@ class ManagedObjectRow(Protocol): _TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object]) _KEY_OWNER_SCOPE_PREFIX: Final = "key:" +_MAX_ROTATION_HOPS: Final = 16 def _string_tags(request_tags: object) -> list[str] | None: @@ -96,6 +101,20 @@ def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteL return TeamRepository(prisma_client).table +def _rotated_token_table( + prisma_client: PrismaClient, +) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]: + return DeprecatedVerificationTokenRepository(prisma_client).table + + +def _deleted_token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedVerificationToken]: + return DeletedVerificationTokenRepository(prisma_client).table + + +def _deleted_team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedTeamTable]: + return TeamRepository(prisma_client).deleted_table + + @dataclass(frozen=True, slots=True) class _CreatorRows: key: prisma_models.LiteLLM_VerificationToken | None @@ -103,6 +122,15 @@ class _CreatorRows: user: prisma_models.LiteLLM_UserTable | None +@dataclass(frozen=True, slots=True) +class ResponseCreator: + metadata: Mapping[str, object] + key_metadata: Mapping[str, object] + team_id: str | None + team_metadata: Mapping[str, object] + settings_read: bool = True + + class ManagedObjectCreatorAttribution: """Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll job's spend row is attributed like the creating request: key hash, team, organization, @@ -117,7 +145,12 @@ class ManagedObjectCreatorAttribution: A background response also carries its end user, the creating key's project and its alias, and the key, team, and user budgets and per-model budgets, the way the read that billed it before did, and never names a key-scoped creator as its user. A key that sets no budget of - its own takes its budget tier's, the same rule auth applies. + its own takes its budget tier's, the same rule auth applies. Its creator's key and team + metadata come along too, so the billed event reaches the key's and team's logging callbacks + like that read did. A rotated key's metadata is read from the live key its rotations lead + to, and a deleted key's or team's from its deleted row. A lookup that fails leaves the + creator's settings unread and the attribution to the rows that still load, so the + response still bills and still counts toward their budgets. """ def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None: @@ -128,31 +161,70 @@ class ManagedObjectCreatorAttribution: rows: Final = await self._rows(job, object_id) return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)} - async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]: - rows: Final = await self._rows(job, object_id) - end_user_id: Final = getattr(job, "end_user_id", None) - return { - "user_api_key_user_id": _creator_user_id(job.created_by), - **_shared_metadata(job, rows), - **({"user_api_key_end_user_id": end_user_id} if end_user_id else {}), - "user_api_key_project_id": getattr(rows.key, "project_id", None), - "user_api_key_project_alias": getattr( - getattr(rows.key, "litellm_project_table", None), "project_alias", None + async def response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator: + try: + return await self._read_response_creator(job, object_id) + except Exception as e: + verbose_proxy_logger.error( + f"{self._poller_name}: could not read the creator of {object_id}, so it bills without their " + f"settings: {e}" + ) + return ResponseCreator( + metadata=_response_creator_metadata(job, await self._rows(job, object_id)), + key_metadata={}, + team_id=getattr(job, "team_id", None), + team_metadata={}, + settings_read=False, + ) + + async def _read_response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator: + api_key: Final = getattr(job, "api_key", None) + team_id: Final = getattr(job, "team_id", None) + rows: Final = _CreatorRows( + key=await self._live_key_row(api_key, job.created_by), + team=await self._live_team_row(team_id), + user=await self._user_row(object_id, _creator_user_id(job.created_by)), + ) + return ResponseCreator( + metadata=_response_creator_metadata(job, rows), + key_metadata=row_metadata(rows.key) if rows.key is not None else await self._retired_key_metadata(api_key), + team_id=team_id, + team_metadata=( + row_metadata(rows.team) if rows.team is not None else await self._deleted_team_metadata(team_id) ), - **_key_budget_fields(rows.key), - **_budget_fields( - rows.team, - "user_api_key_team_spend", - "user_api_key_team_max_budget", - "user_api_key_team_model_max_budget", - ), - **_budget_fields( - rows.user, - "user_api_key_user_spend", - "user_api_key_user_max_budget", - "user_api_key_user_model_max_budget", - ), - } + ) + + async def _retired_key_metadata(self, api_key: str | None) -> Mapping[str, object]: + if not _names_a_key_row(api_key): + return {} + current_token: Final = await self._token_rotated_to(api_key, _MAX_ROTATION_HOPS) + current_key: Final = ( + await _token_table(self._prisma_client).find_unique(where={"token": current_token}) + if current_token != api_key + else None + ) + if current_key is not None: + return row_metadata(current_key) + return row_metadata( + await _deleted_token_table(self._prisma_client).find_first( + where={"token": current_token}, order={"deleted_at": "desc"} + ) + ) + + async def _token_rotated_to(self, token: str, hops_left: int) -> str: + rotation: Final = await _rotated_token_table(self._prisma_client).find_unique(where={"token": token}) + if rotation is None or hops_left == 0: + return token + return await self._token_rotated_to(rotation.active_token_id, hops_left - 1) + + async def _deleted_team_metadata(self, team_id: str | None) -> Mapping[str, object]: + if not team_id: + return {} + return row_metadata( + await _deleted_team_table(self._prisma_client).find_first( + where={"team_id": team_id}, order={"deleted_at": "desc"} + ) + ) async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows: return _CreatorRows( @@ -175,25 +247,57 @@ class ManagedObjectCreatorAttribution: async def _key_row( self, object_id: str, api_key: str | None, created_by: str | None ) -> prisma_models.LiteLLM_VerificationToken | None: - if not api_key or _is_cli_session_key(api_key, created_by): - return None try: - return await _token_table(self._prisma_client).find_unique( - where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True} - ) + return await self._live_key_row(api_key, created_by) except Exception as e: verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}") return None - async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: - if not team_id: + async def _live_key_row( + self, api_key: str | None, created_by: str | None + ) -> prisma_models.LiteLLM_VerificationToken | None: + if not api_key or _is_cli_session_key(api_key, created_by): return None + return await _token_table(self._prisma_client).find_unique( + where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True} + ) + + async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: try: - return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + return await self._live_team_row(team_id) except Exception as e: verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}") return None + async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: + if not team_id: + return None + return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + + +def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: + end_user_id: Final = getattr(job, "end_user_id", None) + return { + "user_api_key_user_id": _creator_user_id(job.created_by), + **_shared_metadata(job, rows), + **({"user_api_key_end_user_id": end_user_id} if end_user_id else {}), + "user_api_key_project_id": getattr(rows.key, "project_id", None), + "user_api_key_project_alias": getattr(getattr(rows.key, "litellm_project_table", None), "project_alias", None), + **_key_budget_fields(rows.key), + **_budget_fields( + rows.team, + "user_api_key_team_spend", + "user_api_key_team_max_budget", + "user_api_key_team_model_max_budget", + ), + **_budget_fields( + rows.user, + "user_api_key_user_spend", + "user_api_key_user_max_budget", + "user_api_key_user_model_max_budget", + ), + } + def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: api_key: Final = getattr(job, "api_key", None) @@ -221,6 +325,19 @@ def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, } +def row_metadata(row: object) -> Mapping[str, object]: + metadata: Final = getattr(row, "metadata", None) + if not isinstance(metadata, dict): + return {} + return cast(Mapping[str, object], metadata) # cast-ok: prisma Json object column + + +def _names_a_key_row(api_key: object) -> TypeGuard[str]: + from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash + + return isinstance(api_key, str) and is_valid_sha256_hash(api_key) + + def _is_cli_session_key(api_key: str, created_by: str | None) -> bool: return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}" diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 63f6979ca3d..17209a29e15 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.78" +version = "0.1.79" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.78" +version = "0.1.79" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index cced49f788a..cc87a2ed16b 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -1,7 +1,7 @@ """Path allowlist for the gateway component. The gateway exposes the LLM data-plane surface: chat/completions, embeddings, -audio, batches, files, fine-tuning, rerank, decisions, ocr, rag, video, search, image, +audio, batches, files, fine-tuning, rerank, decisions, extraction, ocr, rag, video, search, image, responses, vector stores, passthrough providers, realtime websockets, MCP tool-call endpoints, and operational endpoints (/health, /metrics, and the /debug/memory/summary read of the serving worker's RSS). @@ -64,6 +64,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( "/decisions", "/v1/systemone", "/systemone", + "/v1/extraction", + "/extraction", "/v1/ocr", "/ocr", "/v1/rag/", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index a1dee5ac671..6de44be0c12 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.110" +version = "0.4.111" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -30,7 +30,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.110" +version = "0.4.111" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 5a8f9bcf94f..e190b12e665 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -4,7 +4,7 @@ mod selection; pub(crate) use native::NativeCacheHandle; pub(crate) use python::{CacheCall, PythonCache}; -pub(crate) use selection::{Cached, Selection, admit_native, configure, configured_native}; +pub(crate) use selection::{Cached, Selection, configure, configured_native}; use litellm_cache::Error; use pyo3::{ 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 82771644351..eb09b8131c3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs @@ -279,7 +279,7 @@ pub(in crate::cache) fn configured( )> { let handle = native_handle(configured)?.ok_or_else(|| { pyo3::exceptions::PyRuntimeError::new_err( - "the configured cache changed to a Python cache after native admission", + "native inference needs a v2 cache, the catalog keeps Python caches on Python", ) })?; let cache = handle.extract::>()?; diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index b2dc67fe965..cd19533fc73 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -76,21 +76,6 @@ fn selected_cache<'py>( Ok(Some(configured)) } -pub(crate) fn admit_native( - py: Python<'_>, - kwargs: &Bound<'_, PyDict>, - call_type: &str, -) -> PyResult<()> { - if let Some(configured) = selected_cache(py, kwargs, call_type)? - && native::v2::native_handle(&configured)?.is_none() - { - return Err(crate::errors::RustBridgeDeclined::new_err( - "the configured cache requires Python inference", - )); - } - Ok(()) -} - pub(crate) fn configured_native( py: Python<'_>, kwargs: &Bound<'_, PyDict>, diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index a5be7f51828..3b590a522fe 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -177,7 +177,6 @@ where } else { R::SYNC_CALL_TYPE }; - crate::cache::admit_native(py, &call.kwargs, call_type)?; let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?; super::run_public_call( py, diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index 9caeae6c3da..9c045d65c13 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -8,59 +8,19 @@ use litellm_auth::AuthServices; use litellm_cache_response::{CachePolicy, ScopedCache}; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_host::{call::HostedMachine, protocol::Protocol}; -use litellm_host_python::present; use litellm_inference_responses::{ResponsesRoute, route::Responses}; use litellm_secrets::source::SecretSource; use pyo3::prelude::*; -use serde_json::Value; pub(crate) use websocket::ResponsesWebSocketConnection; use super::{ NativeCall, inference::{InferenceHost, InferenceRoute, run_inference}, }; -use crate::errors::RustBridgeDeclined; const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses"; fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { - let resolved = call.resolved()?; - if let Some(reason) = py - .import(ROUTE_HOST_MODULE)? - .getattr("decline_reason")? - .call1((&resolved,))? - .extract::>()? - { - return Err(RustBridgeDeclined::new_err(reason)); - } - let argument = |name: &str| present(&call.kwargs, &call.base, name); - let model = argument("model")? - .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))? - .extract::()?; - let provider = argument("custom_llm_provider")? - .map(|value| value.extract::()) - .transpose()?; - if provider - .as_deref() - .is_some_and(|provider| provider != "openai") - || model - .strip_prefix("openai/") - .unwrap_or(&model) - .contains('/') - { - return Err(RustBridgeDeclined::new_err( - "native HTTP responses provider", - )); - } - if argument("stream")? - .map(|value| litellm_host_python::from_py::(&value)) - .transpose()? - .is_some_and(|value| value == Value::Bool(true)) - { - return Err(RustBridgeDeclined::new_err( - "native Python responses streaming", - )); - } let host = InferenceHost::new(call.view()?, ROUTE_HOST_MODULE); run_inference::(py, call, asynchronous, ResponsesPythonHost::new(host)) } diff --git a/litellm/__init__.py b/litellm/__init__.py index f6604a3b4e5..1784deeb797 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1505,6 +1505,7 @@ from .rag.main import * from .sandbox.main import * from .decisions.main import * from .systemone.main import * +from .extraction.main import * from .tool_loop import ToolLoopMaxRoundsExceeded, arun_tool_loop, run_tool_loop from .search.main import * from .realtime_api.main import ( @@ -1711,6 +1712,9 @@ if TYPE_CHECKING: from .llms.cloudflare.decisions.transformation import ( CloudflareDecisionsConfig as CloudflareDecisionsConfig, ) + from .llms.cloudflare.extraction.transformation import ( + CloudflareExtractionConfig as CloudflareExtractionConfig, + ) from .llms.strands_decider.decisions.transformation import ( StrandsDeciderDecisionsConfig as StrandsDeciderDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 8abc24d1417..e0cb7c680bc 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -166,6 +166,7 @@ LLM_CONFIG_NAMES: Final = ( "OpenRouterDecisionsConfig", "FireworksAIDecisionsConfig", "CloudflareDecisionsConfig", + "CloudflareExtractionConfig", "StrandsDeciderDecisionsConfig", "DatabricksDecisionsConfig", "HostedVLLMDecisionsConfig", @@ -735,6 +736,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"), "FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"), "CloudflareDecisionsConfig": (".llms.cloudflare.decisions.transformation", "CloudflareDecisionsConfig"), + "CloudflareExtractionConfig": (".llms.cloudflare.extraction.transformation", "CloudflareExtractionConfig"), "StrandsDeciderDecisionsConfig": ( ".llms.strands_decider.decisions.transformation", "StrandsDeciderDecisionsConfig", diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index cb6426f64ff..1cc3e106aa0 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -95,7 +95,10 @@ class InMemoryCache(BaseCache): """ Check if a specific key is expired """ - return key in self.ttl_dict and self._clock() > self.ttl_dict[key] + try: + return self._clock() > self.ttl_dict[key] + except KeyError: + return False def _cache_keys(self) -> tuple[str, ...]: return tuple(self.cache_dict) # pyright: ignore[reportUnknownVariableType] # cache_dict is untyped legacy state @@ -227,7 +230,10 @@ class InMemoryCache(BaseCache): if key in self.cache_dict: if self.evict_element_if_expired(key): return None - original_cached_response: Final = self.cache_dict[key] + try: + original_cached_response: Final = self.cache_dict[key] + except KeyError: + return None try: cached_response = json.loads(original_cached_response) except Exception: diff --git a/litellm/chat_completions/dispatch.py b/litellm/chat_completions/dispatch.py index 99900177820..06ba8e9ad76 100644 --- a/litellm/chat_completions/dispatch.py +++ b/litellm/chat_completions/dispatch.py @@ -1,15 +1,19 @@ import inspect from collections.abc import Awaitable, Callable, Coroutine, Mapping +from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable from litellm import main -from litellm.rust_bridge.catalog import Route, RouteContext +from litellm.rust_bridge import detectors +from litellm.rust_bridge.catalog import Feature, Route, RouteContext +from litellm.rust_bridge.chat_completions import detectors as route_detectors from litellm.rust_bridge.chat_completions.entrypoints import ( NATIVE_ACOMPLETION, NATIVE_COMPLETION, ) from litellm.rust_bridge.dispatch import PublicDispatch from litellm.rust_bridge.public_call import ( + AsyncFlag, NativeCall, bind, native_call, @@ -44,6 +48,7 @@ def _python_acompletion() -> PythonAcompletion: _PYTHON_COMPLETION: Final = _python_completion() _COMPLETION: Final = signature(_PYTHON_COMPLETION) +_FLAG: Final = AsyncFlag("acompletion", sync=_COMPLETION) _PYTHON_ACOMPLETION: Final = _python_acompletion() _ACOMPLETION: Final = signature(_PYTHON_ACOMPLETION) @@ -61,11 +66,32 @@ def _public_request( return native_call(legacy, args, kwargs) +_FEATURES: Final[detectors.Detectors] = MappingProxyType( + { + **detectors.COMMON, + Feature.MOCK_RESPONSE: route_detectors.chat_mock_response, + Feature.AGENTIC_LOOP_HOOK: route_detectors.chat_agentic_loop_hook, + Feature.RESPONSE_CACHE: detectors.response_cache, + Feature.STREAMING: detectors.streaming, + Feature.DROP_PARAMS: detectors.drop_params, + Feature.MODIFY_PARAMS: detectors.modify_params, + Feature.EXTRA_BODY: detectors.extra_body, + Feature.RETRIES_FALLBACKS: route_detectors.chat_retries_fallbacks, + Feature.PROMPT_MANAGEMENT: detectors.prompt_management, + Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, + Feature.MODEL_PARAM_GATING: route_detectors.chat_model_param_gating, + Feature.MODEL_LIST: route_detectors.model_list, + Feature.MODEL_ALIAS_MAP: route_detectors.model_alias_map, + } +) + + def _context(request: NativeCall) -> RouteContext: return RouteContext( Route.CHAT_COMPLETIONS, provider=optional_str(request.resolved.get("custom_llm_provider")), model=str(request.resolved["model"]), + features=detectors.detected(_FEATURES, request), ) @@ -73,7 +99,6 @@ _DISPATCH: Final = PublicDispatch( route=Route.CHAT_COMPLETIONS, request=lambda args, kwargs: _public_request(_COMPLETION, args, kwargs), context=_context, - bypass=lambda request: request.kwargs.get("acompletion") is True, ) _ADISPATCH: Final = PublicDispatch( @@ -88,9 +113,11 @@ def completion( **kwargs: object, # kwargs-ok: preserve the public chat completions call shape ) -> ChatResult | Coroutine[object, object, ChatResult]: python: Final = _PYTHON_COMPLETION + if _FLAG.requested(kwargs): + return acompletion(**_FLAG.keywords(args, kwargs)) return _DISPATCH.run( args, - kwargs, + _FLAG.dropped(kwargs), python=python, binding=NATIVE_COMPLETION, native=native_call_hook, @@ -101,7 +128,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs- python: Final = _PYTHON_ACOMPLETION return await _ADISPATCH.arun( args, - kwargs, + _FLAG.dropped(kwargs), python=python, binding=NATIVE_ACOMPLETION, native=native_call_hook, diff --git a/litellm/constants.py b/litellm/constants.py index cb50d74f0cf..6a503463a8c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -25,6 +25,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code" MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC" GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth" GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type" +GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host" +GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com" GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60 GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60 GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow" @@ -154,6 +156,9 @@ DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLIN DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1) DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250) HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0 +HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( + "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS", 300, 1, 86400 +) # Maximum wall-clock seconds a streaming response is allowed to run. # Streams exceeding this duration are terminated with a Timeout error. @@ -256,6 +261,9 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0) +MCP_MAX_SSE_EVENT_SIZE_BYTES: Final = get_env_int_in_range( + "LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", 64 * 1024 * 1024, 1, sys.maxsize +) MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0) MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES: Final = get_env_int("LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES", 3) MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS: Final = get_env_float( @@ -1649,6 +1657,7 @@ BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags" VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata" ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" +BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt" 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" @@ -1812,6 +1821,7 @@ LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, get_env_int("LITELLM_EXECUTED ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli" LITELLM_CLI_SESSION_TOKEN_PREFIX: Final = "litellm-session-token" +CP_RETURN_TO_COOKIE_PREFIX: Final = "return_to.v1." CLI_SSO_SESSION_CACHE_KEY_PREFIX: Final = "cli_sso_session" CLI_SSO_SESSION_TTL_SECONDS: Final = 600 CLI_SESSION_KEY_PREFIX: Final = "cli-session" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 67178c8959c..f952883bc87 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -112,6 +112,7 @@ from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_toke from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.agents import LiteLLMSendMessageResponse from litellm.types.decisions import DecisionsResponse, DecisionsUsage, OpenAIDecisionResponse, OpenAIDecisionUsage +from litellm.types.extraction import ExtractionResponse, ExtractionUsage from litellm.types.llms.base import CachedTokensDetails, LiteLLMBaseModel from litellm.types.llms.openai import ( HttpxBinaryResponseContent, @@ -1171,6 +1172,8 @@ def get_usage_object( if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)): return None if completion_response.usage is None else _decisions_usage(completion_response.usage) + if isinstance(completion_response, ExtractionResponse): + return _extraction_usage(completion_response.usage) if isinstance(completion_response, HttpxBinaryResponseContent): return completion_response.usage if usage_obj is None: @@ -1227,6 +1230,14 @@ def _decisions_usage(usage: DecisionsUsage | OpenAIDecisionUsage) -> Usage: ) +def _extraction_usage(usage: ExtractionUsage) -> Usage: + return Usage( + prompt_tokens=usage.input_tokens, + completion_tokens=usage.output_tokens, + total_tokens=usage.input_tokens + usage.output_tokens, + ) + + def _is_known_usage_objects(usage_obj): """Returns True if the usage obj is a known Usage type""" return ( @@ -1234,6 +1245,7 @@ def _is_known_usage_objects(usage_obj): or isinstance(usage_obj, ResponseAPIUsage) or isinstance(usage_obj, DecisionsUsage) or isinstance(usage_obj, OpenAIDecisionUsage) + or isinstance(usage_obj, ExtractionUsage) or TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj) ) @@ -1263,6 +1275,8 @@ def infer_call_type(call_type: CallTypesLiteral | None, completion_response: obj return "send_message" elif isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)): return "adecisions" + elif isinstance(completion_response, ExtractionResponse): + return "aextraction" return call_type @@ -1674,6 +1688,8 @@ def completion_cost( ) if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)): _usage = _decisions_usage(usage_obj).model_dump() + elif isinstance(usage_obj, ExtractionUsage): + _usage = _extraction_usage(usage_obj).model_dump() elif usage_obj is None: _usage = {} elif isinstance(usage_obj, BaseModel): @@ -2199,7 +2215,8 @@ def response_cost_calculator( | Response | SearchResponse | DecisionsResponse - | OpenAIDecisionResponse, + | OpenAIDecisionResponse + | ExtractionResponse, model: str, custom_llm_provider: str | None, call_type: Literal[ @@ -2225,6 +2242,8 @@ def response_cost_calculator( "adecisions", "systemone", "asystemone", + "extraction", + "aextraction", ], optional_params: dict, cache_hit: bool | None = None, diff --git a/litellm/embeddings/dispatch.py b/litellm/embeddings/dispatch.py index 3fd8893464d..772ac68737d 100644 --- a/litellm/embeddings/dispatch.py +++ b/litellm/embeddings/dispatch.py @@ -9,7 +9,15 @@ from litellm.rust_bridge.embeddings.entrypoints import ( NATIVE_AEMBEDDING, NATIVE_EMBEDDING, ) -from litellm.rust_bridge.public_call import NativeCall, bind, native_call, native_call_hook, optional_str, signature +from litellm.rust_bridge.public_call import ( + AsyncFlag, + NativeCall, + bind, + native_call, + native_call_hook, + optional_str, + signature, +) from litellm.types.utils import EmbeddingResponse __all__ = ("aembedding", "embedding") @@ -24,6 +32,7 @@ _PYTHON_AEMBEDDING: Final = cast( # cast-ok: [LIT006] preserve the legacy publi PythonAembedding, main.aembedding ) _EMBEDDING_SIGNATURE: Final = signature(_PYTHON_EMBEDDING) +_FLAG: Final = AsyncFlag("aembedding", sync=_EMBEDDING_SIGNATURE) def _public_request( @@ -50,7 +59,6 @@ _DISPATCH: Final = PublicDispatch( route=Route.EMBEDDINGS, request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs), context=_context, - bypass=lambda request: request.kwargs.get("aembedding") is True, ) _ADISPATCH: Final = PublicDispatch( @@ -64,9 +72,11 @@ def embedding( *args: object, **kwargs: object, # kwargs-ok: preserve the public embedding call shape ) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]: + if _FLAG.requested(kwargs): + return aembedding(**_FLAG.keywords(args, kwargs)) return _DISPATCH.run( args, - kwargs, + _FLAG.dropped(kwargs), python=_PYTHON_EMBEDDING, binding=NATIVE_EMBEDDING, native=native_call_hook, @@ -76,7 +86,7 @@ def embedding( async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # kwargs-ok: preserve the public call shape return await _ADISPATCH.arun( args, - kwargs, + _FLAG.dropped(kwargs), python=_PYTHON_AEMBEDDING, binding=NATIVE_AEMBEDDING, native=native_call_hook, diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ed6d89397a0..79887e4ab99 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -95,6 +95,7 @@ from pydantic import AnyUrl, TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( MCP_CLIENT_TIMEOUT, + MCP_MAX_SSE_EVENT_SIZE_BYTES, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_MAX_PAGES, MCP_TOOL_LISTING_TIMEOUT, @@ -292,6 +293,7 @@ class _MCPResponseStream(httpx2.AsyncByteStream): class _MCPHTTPClient(httpx2.AsyncClient): cleanup_scope: anyio.CancelScope | None = None cleanup_errors: tuple[Exception, ...] = () + post_failure: asyncio.Future[Exception] | None = None def _record_cleanup_error(self, error: Exception) -> None: if self.cleanup_scope is not None and self.cleanup_scope.shield: @@ -319,7 +321,12 @@ class _MCPHTTPClient(httpx2.AsyncClient): response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects) if request.method == "POST" and response.is_error and response.status_code != 404: await response.aclose() - response.raise_for_status() + try: + response.raise_for_status() + except httpx2.HTTPStatusError as error: + if self.post_failure is not None and not self.post_failure.done(): + self.post_failure.set_result(error) + raise if stream: response.stream = _MCPResponseStream( cast(httpx2.AsyncByteStream, response.stream), self._record_cleanup_error @@ -641,6 +648,7 @@ class MCPClient: transport_ctx: Final = streamable_http_client( url=self.server_url, http_client=http_client, + max_sse_event_size=MCP_MAX_SSE_EVENT_SIZE_BYTES, ) return transport_ctx, http_client @@ -757,6 +765,8 @@ class MCPClient: read_stream: Final = transport[0] write_stream: Final = transport[1] stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future() + if isinstance(http_client, _MCPHTTPClient): + http_client.post_failure = stream_error # rebind-ok: operation-scoped stream error _notify_stream_error(on_stream_error, stream_error) session_kwargs: Final = ( @@ -823,6 +833,7 @@ class MCPClient: if isinstance(http_client, _MCPHTTPClient): http_client.cleanup_errors = () http_client.cleanup_scope = None + http_client.post_failure = None # rebind-ok: clear operation-scoped stream error await anyio.lowlevel.checkpoint_if_cancelled() if cleanup_scope.cancel_called: raise ( diff --git a/litellm/extraction/__init__.py b/litellm/extraction/__init__.py new file mode 100644 index 00000000000..bdacdae6ea4 --- /dev/null +++ b/litellm/extraction/__init__.py @@ -0,0 +1,3 @@ +from litellm.extraction.main import aextraction, extraction + +__all__ = ["aextraction", "extraction"] diff --git a/litellm/extraction/main.py b/litellm/extraction/main.py new file mode 100644 index 00000000000..94198e911ef --- /dev/null +++ b/litellm/extraction/main.py @@ -0,0 +1,351 @@ +import asyncio +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import partial +from types import MappingProxyType +from typing import Final, TypeAlias + +import httpx +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm.decisions.call import log_caller_request +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.extraction import ( + EXTRACTION_STATE_REWRITTEN_METADATA_KEY, + ExtractionAnswer, + ExtractionBase64Image, + ExtractionField, + ExtractionOptions, + ExtractionRequestBody, + ExtractionResponse, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders, is_litellm_owned_kwarg +from litellm.utils import ( + ProviderConfigManager, + client, # pyright: ignore[reportUnknownVariableType] # client is an untyped decorator + get_model_cost_key, +) + +ExtractionFields: TypeAlias = Mapping[str, ExtractionField | Mapping[str, object]] +ExtractionImages: TypeAlias = Sequence[str | ExtractionBase64Image | Mapping[str, object]] +ExtractionOptionsParam: TypeAlias = ExtractionOptions | Mapping[str, object] + +_REQUEST_ADAPTER: Final[TypeAdapter[ExtractionRequestBody]] = TypeAdapter(ExtractionRequestBody) +_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_HANDLER: Final = BaseLLMHTTPHandler() +_LOCAL_KWARGS: Final = frozenset( + ( + *GenericLiteLLMParams.model_fields, + "extra_headers", + "base_url", + "default_headers", + "deployment_id", + "additional_drop_params", + ) +) + + +@dataclass(frozen=True, slots=True, repr=False) +class _ExtractionCall: + model: str + requested_model: str + custom_llm_provider: str + provider_config: BaseExtractionConfig + body: Mapping[str, object] + api_base: str + api_key: str | None + logging_obj: LiteLLMLoggingObj | None + headers: Mapping[str, str] + timeout: float | httpx.Timeout | None + litellm_params: Mapping[str, object] + client: HTTPHandler | AsyncHTTPHandler | None + + +def supported_extraction_providers() -> tuple[str, ...]: + return tuple( + provider.value + for provider in LlmProviders + if ProviderConfigManager.get_provider_extraction_config(model="", provider=provider) is not None + ) + + +def _provider_config(model: str, custom_llm_provider: str) -> BaseExtractionConfig: + provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None) + provider_config: Final = ( + None if provider is None else ProviderConfigManager.get_provider_extraction_config(model, provider) + ) + if provider_config is None: + raise litellm.BadRequestError( + message=( + f"Unknown Extraction provider '{custom_llm_provider}'. " + f"Supported providers: {', '.join(supported_extraction_providers())}" + ), + model=model, + llm_provider=custom_llm_provider, + ) + return provider_config + + +def _request_fields( + *, + state: str | None, + fields: ExtractionFields | None, + images: ExtractionImages | None, + reference_date: str | None, + options: ExtractionOptionsParam | None, +) -> Mapping[str, object]: + sent: Final[Mapping[str, object | None]] = MappingProxyType( + {"state": state, "fields": fields, "images": images, "reference_date": reference_date, "options": options} + ) + return MappingProxyType({name: value for name, value in sent.items() if value is not None}) + + +def _is_unsupported(name: str, value: object, policy: DropPolicy) -> bool: + return ( + value is not None + and value is not False + and name not in _LOCAL_KWARGS + and not is_litellm_owned_kwarg(name) + and not policy.drops(name) + ) + + +def _unsupported_params(kwargs: Mapping[str, object]) -> tuple[str, ...]: + policy: Final = drop_policy(kwargs) + return tuple(sorted(name for name, value in kwargs.items() if _is_unsupported(name, value, policy))) + + +def _prepare_call( + *, + model: str, + request_fields: Mapping[str, object], + api_key: str | None, + api_base: str | None, + timeout: float | httpx.Timeout | None, + custom_llm_provider: str | None, + extra_headers: Mapping[str, str] | None, + kwargs: Mapping[str, object], +) -> _ExtractionCall: + log_caller_request(kwargs, request_fields) + upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + ) + provider_config: Final = _provider_config(model, provider) + unsupported: Final = _unsupported_params(kwargs) + if unsupported: + raise litellm.UnsupportedParamsError( + message=( + f"{provider} does not support parameters: {list(unsupported)}, for model={model}. " + "To drop these, set `litellm.drop_params=True` or for proxy:\n\n" + "`litellm_settings:\n drop_params: true`\n" + ), + model=model, + llm_provider=provider, + ) + try: + request: Final = _REQUEST_ADAPTER.validate_python(request_fields) + except ValidationError as error: + raise litellm.BadRequestError( + message=f"Invalid Extraction request: {error}", model=model, llm_provider=provider + ) from error + canonical_model: Final = provider_config.canonical_model(upstream_model) + resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base) + if resolved_api_base is None: + raise litellm.BadRequestError( + message=provider_config.missing_api_base_message(provider), model=model, llm_provider=provider + ) + resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key) + if resolved_api_key is None: + raise litellm.BadRequestError( + message=f"Missing API key for Extraction provider '{provider}'", model=model, llm_provider=provider + ) + logging_obj: Final = kwargs.get("litellm_logging_obj") + http_client: Final = kwargs.get("client") + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.update_from_kwargs( # pyright: ignore[reportUnknownMemberType] # Logging takes untyped dicts + kwargs=dict(kwargs), + model=canonical_model, + litellm_params={ + "litellm_call_id": kwargs.get("litellm_call_id"), + "api_base": provider_config.get_complete_url(resolved_api_base, canonical_model), + }, + custom_llm_provider=provider, + ) + return _ExtractionCall( + model=canonical_model, + requested_model=model, + custom_llm_provider=provider, + provider_config=provider_config, + body=provider_config.transform_extraction_request(canonical_model, request), + api_base=resolved_api_base, + api_key=resolved_api_key, + logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None, + headers=MappingProxyType(dict(extra_headers or {})), + timeout=timeout, + litellm_params=kwargs, + client=http_client if isinstance(http_client, HTTPHandler | AsyncHTTPHandler) else None, + ) + + +def _state_rewritten_by_guardrail(litellm_params: Mapping[str, object]) -> bool: + metadata: Final = litellm_params.get(get_metadata_variable_name_from_kwargs(litellm_params)) + return ( + isinstance(metadata, Mapping) + and _METADATA_ADAPTER.validate_python(metadata).get(EXTRACTION_STATE_REWRITTEN_METADATA_KEY) is True + ) + + +def _without_spans(answers: Mapping[str, ExtractionAnswer]) -> Mapping[str, ExtractionAnswer]: + return {name: answer.model_copy(update={"span": None, "spans": None}) for name, answer in answers.items()} + + +def _formatted_response(response: ExtractionResponse, call: _ExtractionCall) -> ExtractionResponse: + spans_index_hidden_text: Final = response.state_text is None and _state_rewritten_by_guardrail(call.litellm_params) + formatted: Final = response.model_copy( + update={ + "model": call.requested_model, + **({"answers": _without_spans(response.answers)} if spans_index_hidden_text else {}), + } + ) + formatted.set_hidden_params( + { + "model": f"{call.custom_llm_provider}/{call.model}", + "custom_llm_provider": call.custom_llm_provider, + "provider_response_model": ( + get_model_cost_key(f"{call.custom_llm_provider}/{response.model}") + or f"{call.custom_llm_provider}/{call.model}" + ), + } + ) + return formatted + + +def _upstream_exception(error: Exception, call: _ExtractionCall) -> Exception: + if isinstance(error, BaseLLMException) and error.status_code_is_synthesized: + return litellm.APIConnectionError( + message=f"{call.custom_llm_provider.capitalize()}Exception - {error.message}", + llm_provider=call.custom_llm_provider, + model=f"{call.custom_llm_provider}/{call.model}", + ) + return litellm.exception_type( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped + model=f"{call.custom_llm_provider}/{call.model}", + custom_llm_provider=call.custom_llm_provider, + original_exception=error, + ) + + +async def _asend(call: _ExtractionCall) -> ExtractionResponse: + try: + response: Final = await _HANDLER.aextraction( + model=call.model, + custom_llm_provider=call.custom_llm_provider, + logging_obj=call.logging_obj, + provider_config=call.provider_config, + body=call.body, + api_base=call.api_base, + api_key=call.api_key, + headers=call.headers, + timeout=call.timeout, + litellm_params=call.litellm_params, + client=call.client if isinstance(call.client, AsyncHTTPHandler) else None, + ) + except Exception as error: + raise _upstream_exception(error, call) from error + return _formatted_response(response, call) + + +def _send(call: _ExtractionCall) -> ExtractionResponse: + try: + response: Final = _HANDLER.extraction( + model=call.model, + custom_llm_provider=call.custom_llm_provider, + logging_obj=call.logging_obj, + provider_config=call.provider_config, + body=call.body, + api_base=call.api_base, + api_key=call.api_key, + headers=call.headers, + timeout=call.timeout, + litellm_params=call.litellm_params, + client=call.client if isinstance(call.client, HTTPHandler) else None, + ) + except Exception as error: + raise _upstream_exception(error, call) from error + return _formatted_response(response, call) + + +@client +async def aextraction( + model: str, + fields: ExtractionFields, + state: str | None = None, + images: ExtractionImages | None = None, + reference_date: str | None = None, + options: ExtractionOptionsParam | None = None, + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, + extra_headers: Mapping[str, str] | None = None, + **kwargs: object, # kwargs-ok: litellm_params forwarded by @client +) -> ExtractionResponse: + call: Final = await asyncio.to_thread( + partial( + _prepare_call, + model=model, + request_fields=_request_fields( + state=state, fields=fields, images=images, reference_date=reference_date, options=options + ), + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) + ) + return await _asend(call) + + +@client +def extraction( + model: str, + fields: ExtractionFields, + state: str | None = None, + images: ExtractionImages | None = None, + reference_date: str | None = None, + options: ExtractionOptionsParam | None = None, + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, + extra_headers: Mapping[str, str] | None = None, + **kwargs: object, # kwargs-ok: litellm_params forwarded by @client +) -> ExtractionResponse: + call: Final = _prepare_call( + model=model, + request_fields=_request_fields( + state=state, fields=fields, images=images, reference_date=reference_date, options=options + ), + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) + return _send(call) + + +__all__ = ["aextraction", "extraction"] diff --git a/litellm/harness/endpoint.py b/litellm/harness/endpoint.py index 0abf7c11280..71e2742fa2e 100644 --- a/litellm/harness/endpoint.py +++ b/litellm/harness/endpoint.py @@ -37,6 +37,7 @@ from litellm.harness.context import GatewayTarget from litellm.harness.errors import HarnessError, HarnessInstallFailed from litellm.harness.types import Harness, Usage from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -690,18 +691,21 @@ class ModelEndpoint: async def _invoke_sdk( self, route: str, - kwargs: dict[str, Any], # mutable-ok: injects stream_options into the SDK kwargs + kwargs: dict[str, Any], # mutable-ok: SDK call kwargs built by _sdk_kwargs ) -> object: if route == ROUTE_MESSAGES: return await litellm.anthropic.messages.acreate(**kwargs) if route == ROUTE_CHAT: - if kwargs.get("stream"): - stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field - kwargs["stream_options"] = { # mutable-ok: JSON field sent to litellm.acompletion - "include_usage": True, - **stream_options, + if not kwargs.get("stream"): + return await litellm.acompletion(**kwargs) + stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field + return await litellm.acompletion( + **{ + **kwargs, + "stream_options": {"include_usage": True, **stream_options}, + STREAM_USAGE_INJECTED_FIELD: "include_usage" not in stream_options, } - return await litellm.acompletion(**kwargs) + ) return await litellm.aresponses(**kwargs) async def _call_sdk(self, route: str, body: Mapping[str, object]) -> Response: diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 657fd3abfbc..85557090ae9 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -77,6 +77,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "otpm", "use_xai_oauth", "github_copilot_auth_type", + "github_copilot_host", "github_copilot_user_session", "fireworks_forward_user_id", "responses_api_version", diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 03af9076c79..b19bcd941fa 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -4,6 +4,7 @@ Helper functions for health check calls. import base64 from collections.abc import Awaitable, Callable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal from litellm.llms.base_llm.ocr.transformation import DocumentType @@ -25,6 +26,9 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = ( ) ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16 +EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( + {"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})} +) def native_health_check_mode(model: str, custom_llm_provider: str | None) -> Literal["anthropic_messages"] | None: @@ -215,6 +219,7 @@ class HealthCheckHelpers: "anthropic_messages", "ocr", "decisions", + "extraction", ], Callable, ]: @@ -237,7 +242,11 @@ class HealthCheckHelpers: from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_for_health_check, ) - from litellm.litellm_core_utils.health_check_utils import DECISIONS_CALL_PARAMS, _filter_model_params + from litellm.litellm_core_utils.health_check_utils import ( + DECISIONS_CALL_PARAMS, + EXTRACTION_CALL_PARAMS, + _filter_model_params, + ) from litellm.realtime_api.main import realtime_health_check return { @@ -322,4 +331,13 @@ class HealthCheckHelpers: } ) ), + "extraction": lambda: litellm.aextraction( + **EXTRACTION_CALL_PARAMS.validate_python( + { + "state": prompt or "health check", + "fields": EXTRACTION_HEALTH_CHECK_FIELDS, + **_filter_model_params(model_params=model_params), + } + ) + ), } diff --git a/litellm/litellm_core_utils/health_check_utils.py b/litellm/litellm_core_utils/health_check_utils.py index 9b344d08d4b..74b6f6dea5d 100644 --- a/litellm/litellm_core_utils/health_check_utils.py +++ b/litellm/litellm_core_utils/health_check_utils.py @@ -8,8 +8,10 @@ from typing import Final from pydantic import TypeAdapter from litellm.types.decisions import DecisionsCallParams +from litellm.types.extraction import ExtractionCallParams DECISIONS_CALL_PARAMS: Final[TypeAdapter[DecisionsCallParams]] = TypeAdapter(DecisionsCallParams) +EXTRACTION_CALL_PARAMS: Final[TypeAdapter[ExtractionCallParams]] = TypeAdapter(ExtractionCallParams) OPTIONAL_STR: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b071cb8f65e..2a841385756 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -12,7 +12,7 @@ import subprocess import sys import time import traceback -from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime as dt_object from functools import lru_cache @@ -132,11 +132,13 @@ from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.agents import LiteLLMSendMessageResponse from litellm.types.containers.main import ContainerObject from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse +from litellm.types.extraction import ExtractionResponse from litellm.types.integrations.s3_v2 import S3PartitionGranularity from litellm.types.interactions import ( InteractionsAPIResponse, InteractionsAPIStreamingResponse, ) +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.openai import ( AllMessageValues, Batch, @@ -573,6 +575,27 @@ class UncachedPrice: cost_breakdown: CostBreakdown | None +@dataclass(frozen=True, slots=True) +class BilledRoundsSpend: + usage: Usage + cost: float + cost_breakdown: CostBreakdown | None + + +@dataclass(frozen=True, slots=True) +class BilledRoundsDispatch: + """A parked success log for a stream bridged through chat completions that billed several rounds, and the spend + its failure log bills when a post-call hook rejects the response, as it does for a rejected one-round bridged + stream.""" + + dispatch: Coroutine[object, object, None] + spend: BilledRoundsSpend + + +def _fixed_margin(cost_breakdown: CostBreakdown | None) -> float: + return 0.0 if cost_breakdown is None else cost_breakdown.get("margin_fixed_amount", 0.0) + + def _timestamp_seconds(moment: object) -> float | None: if isinstance(moment, datetime.datetime): return moment.timestamp() @@ -594,6 +617,7 @@ class Logging(LiteLLMLoggingBaseClass): prometheusLogger, \ slack_app custom_pricing: bool = False + replaced_global_callback_keys: frozenset[str] = frozenset() stream_options = None litellm_request_debug: bool = False streamed_anthropic_message_id: str | None = None @@ -648,11 +672,11 @@ class Logging(LiteLLMLoggingBaseClass): litellm_call_id: str, function_id: str, litellm_trace_id: str | None = None, - dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, - dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None, - dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, + dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, applied_guardrails: list[str] | None = None, kwargs: dict | None = None, log_raw_request_response: bool = False, @@ -697,6 +721,7 @@ class Logging(LiteLLMLoggingBaseClass): self._pre_call_session_id: str = session_id_var.get() _sid: Final = kwargs.get("litellm_session_id") if kwargs else None self.litellm_session_id: str = str(_sid) if _sid else "" + self.stream_usage_injected: Final = kwargs is not None and kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True # supports_correlation_logging is False for calls originating from the # sync client entry point (wrapper() in utils.py): a plain OS thread # has no per-call context isolation the way an asyncio Task does, and @@ -725,15 +750,17 @@ class Logging(LiteLLMLoggingBaseClass): self.raw_request_only = raw_request_only # Initialize dynamic callbacks - self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks + self.dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = dynamic_input_callbacks self.dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_success_callbacks ) self.dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_async_success_callbacks ) - self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks - self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = ( + self.dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( + dynamic_failure_callbacks + ) + self.dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_async_failure_callbacks ) @@ -796,6 +823,7 @@ class Logging(LiteLLMLoggingBaseClass): self.defer_async_logging: bool = False self.enqueue_deferred_logging: Callable[[], None] | None = None self._async_success_scheduled: bool = False + 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 @@ -850,10 +878,32 @@ class Logging(LiteLLMLoggingBaseClass): self.dynamic_async_failure_callbacks, callback ) + def replace_global_callback(self, callback: CustomLogger) -> None: + self.add_dynamic_callback(callback) + self.replaced_global_callback_keys = self.replaced_global_callback_keys | { + litellm.logging_callback_manager.get_custom_logger_key(callback) + } + + def add_dynamic_function_callbacks( + self, + *, + async_success: Callable[..., object], + failure: Callable[..., object], + async_failure: Callable[..., object], + ) -> None: + self.dynamic_async_success_callbacks = self._with_dynamic_callback( + self.dynamic_async_success_callbacks, async_success + ) + self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, failure) + self.dynamic_async_failure_callbacks = self._with_dynamic_callback( + self.dynamic_async_failure_callbacks, async_failure + ) + @staticmethod def _with_dynamic_callback( - callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger - ) -> list[str | Callable | CustomLogger]: + callbacks: Sequence[str | Callable[..., object] | CustomLogger] | None, + callback: Callable[..., object] | CustomLogger, + ) -> list[str | Callable[..., object] | CustomLogger]: existing: Final = tuple(callbacks or ()) return [*existing, *(() if callback in existing else (callback,))] @@ -890,9 +940,9 @@ class Logging(LiteLLMLoggingBaseClass): def _process_dynamic_callback_list( self, - callback_list: list[str | Callable | CustomLogger] | None, + callback_list: list[str | Callable[..., object] | CustomLogger] | None, dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"], - ) -> list[str | Callable | CustomLogger] | None: + ) -> list[str | Callable[..., object] | CustomLogger] | None: """ Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks @@ -904,7 +954,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback_list is None: return None - processed_list: Final[list[str | Callable | CustomLogger]] = [] + processed_list: Final[list[str | Callable[..., object] | CustomLogger]] = [] for callback in callback_list: if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks: for callback_instance in self._resolve_dynamic_callback_string(callback): @@ -1565,7 +1615,10 @@ class Logging(LiteLLMLoggingBaseClass): "api_call_start_time" ] # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) + callbacks: Final = ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) for callback in callbacks: try: if callback == "supabase" and supabaseClient is not None: @@ -1738,7 +1791,10 @@ class Logging(LiteLLMLoggingBaseClass): ) # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) + callbacks: Final = ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) for callback in callbacks: try: if callback == "sentry" and add_breadcrumb: @@ -2209,20 +2265,23 @@ class Logging(LiteLLMLoggingBaseClass): usage=usage, ) - def custom_pricing_for(self, result: object) -> bool: - litellm_params: Final = getattr(self, "litellm_params", None) + def spilled_over_ptu_for(self, result: object) -> bool: result_hidden_params: Final = getattr(result, HIDDEN_PARAMS_ATTR, None) or MappingProxyType({}) additional_headers: Final = ( result_hidden_params.get("additional_headers") if isinstance(result_hidden_params, dict) else getattr(result_hidden_params, "additional_headers", None) ) - spilled_over: Final = is_spilled_over_ptu_request( - model_info=_deployment_model_info(litellm_params), + return is_spilled_over_ptu_request( + model_info=_deployment_model_info(getattr(self, "litellm_params", None)), response_headers=self.model_call_details.get("response_headers"), additional_headers=additional_headers, ) - return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params) + + def custom_pricing_for(self, result: object) -> bool: + if self.spilled_over_ptu_for(result): + return False + return use_custom_pricing_for_model(litellm_params=getattr(self, "litellm_params", None)) def _prompt_for_cost_calculation(self) -> str: """ @@ -2384,6 +2443,31 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def record_failure_after_earlier_rounds( + self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float + ) -> None: + """Bill a failed round on the failure log after the rounds of its request before it were billed on the success + log. + + The fixed margin is a per-request fee, so it is charged once: the failed round's margin is dropped when the + earlier rounds already paid one.""" + from litellm.responses.mcp.round_billing import breakdown_less_repeated_fixed_margin + + failed_breakdown: Final = None if self.cost_breakdown is earlier.cost_breakdown else self.cost_breakdown + repeated_fixed_margin: Final = ( + 0.0 + if failed_breakdown is None + else min(_fixed_margin(earlier.cost_breakdown), _fixed_margin(failed_breakdown), earlier.cost, failed_cost) + ) + self.model_call_details["combined_usage_object"] = failed_usage + self.model_call_details["response_cost"] = failed_cost - repeated_fixed_margin + self.cost_breakdown = breakdown_less_repeated_fixed_margin(failed_breakdown, repeated_fixed_margin) + + def record_billed_rounds_for_failure(self, spend: BilledRoundsSpend) -> None: + """Bill every round of a fully streamed request on the failure log when a post-call hook rejects it.""" + self.record_partial_usage_for_failure(spend.usage, spend.cost) + self.cost_breakdown = spend.cost_breakdown + def attempt_scoped_copy(self) -> "Logging": """A copy with its own model_call_details, so logging one attempt of a request through it leaves the state the request's other attempts share untouched.""" @@ -2472,7 +2556,7 @@ class Logging(LiteLLMLoggingBaseClass): await self.async_failure_handler(exception, traceback_exception) finally: if self._should_run_sync_failure_callbacks_for_async_calls(): - executor.submit(self.failure_handler, exception, traceback_exception) + executor.submit(self.failure_handler, exception, traceback_exception, for_async_request=True) def should_run_logging( self, @@ -2487,6 +2571,16 @@ class Logging(LiteLLMLoggingBaseClass): except Exception: return True + @property + def async_success_log_enqueued(self) -> bool: + """True once the success log task was actually created, not merely claimed, so a later failure raised + around it (a disconnect cancelling a still-pending guardrail task) must not bill the prompt a second + time. A claimed log parked in the deferred slot never runs once a 499 interrupts, so it stays False.""" + return self._async_success_enqueued + + def mark_async_success_log_enqueued(self) -> None: + self._async_success_enqueued = True + def claim_async_success_log(self) -> bool: """One async success log per request. The innermost @client wrapper always exits first, since the outer one is awaiting it, so it claims the log here and the outer wrapper gets @@ -2915,6 +3009,7 @@ class Logging(LiteLLMLoggingBaseClass): or isinstance(logging_result, SearchResponse) # Search API or isinstance(logging_result, DecisionsResponse) or isinstance(logging_result, OpenAIDecisionResponse) + or isinstance(logging_result, ExtractionResponse) or ( isinstance(logging_result, InteractionsAPIResponse) and logging_result.usage is not None @@ -4086,15 +4181,21 @@ class Logging(LiteLLMLoggingBaseClass): traceback_exception: str, start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, + *, + for_async_request: bool = False, ) -> None: """Restores trace_id/session_id contextvars once this attempt's own failure - logging (including any nested calls its callbacks trigger) is fully done.""" + logging (including any nested calls its callbacks trigger) is fully done. + + ``for_async_request`` marks the sync run of a failure whose async run also logs it, so a + CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook.""" try: return self._failure_handler_body( exception=exception, traceback_exception=traceback_exception, start_time=start_time, end_time=end_time, + for_async_request=for_async_request, ) finally: self._restore_correlation_context() @@ -4105,6 +4206,7 @@ class Logging(LiteLLMLoggingBaseClass): traceback_exception: str, start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, + for_async_request: bool = False, ) -> None: verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback) if not self.should_run_logging(event_type="sync_failure"): # prevent double logging @@ -4197,6 +4299,7 @@ class Logging(LiteLLMLoggingBaseClass): isinstance(callback, CustomLogger) and is_sync_request and self.call_type != CallTypes.pass_through.value + and not (for_async_request and self._async_handler_logs_to(callback)) ): # custom logger class callback.log_failure_event( start_time=start_time, @@ -4499,14 +4602,26 @@ class Logging(LiteLLMLoggingBaseClass): _filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks) return len(_filtered_failure_callbacks) > 0 + def _without_replaced_global_callbacks( + self, global_callbacks: Iterable[str | Callable[..., object] | CustomLogger] + ) -> tuple[str | Callable[..., object] | CustomLogger, ...]: + if not self.replaced_global_callback_keys: + return tuple(global_callbacks) + return tuple( + litellm.logging_callback_manager.without_custom_logger_keys( + global_callbacks, self.replaced_global_callback_keys + ) + ) + def get_combined_callback_list( self, dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None, global_callbacks: list[str | Callable[..., object] | CustomLogger], ) -> list[str | Callable[..., object] | CustomLogger]: + kept_global_callbacks: Final = self._without_replaced_global_callbacks(global_callbacks) if dynamic_success_callbacks is None: - return list(global_callbacks) - return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks)) + return list(kept_global_callbacks) + return list(dict.fromkeys((*dynamic_success_callbacks, *kept_global_callbacks))) def _remove_internal_litellm_callbacks(self, callbacks: list) -> list: """ @@ -4559,6 +4674,12 @@ class Logging(LiteLLMLoggingBaseClass): cb_name: Final = self._get_callback_name(cb) return any(prefix in cb_name for prefix in INTERNAL_PREFIXES) + def _async_handler_logs_to(self, logger: CustomLogger) -> bool: + """True when the async failure run of this call hands ``logger`` its own ``async_log_failure_event``.""" + if type(logger).async_log_failure_event is CustomLogger.async_log_failure_event: # pyright: ignore[reportUnknownMemberType] # CustomLogger's hook signatures are untyped + return False + return logger in (self.dynamic_async_failure_callbacks or []) or logger in litellm._async_failure_callback # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # callback lists are untyped and have no public accessor + def _remove_internal_custom_logger_callbacks(self, callbacks: list) -> list: """ Removes internal custom logger callbacks from the list. @@ -7227,13 +7348,16 @@ def get_standard_logging_object_payload( saved_cache_cost: Final = (uncached_price.cost if uncached_price is not None else None) or 0.0 ## Get model cost information ## - base_model = get_base_model_from_metadata(model_call_details=kwargs) - # The router overrides completion_response.model to the model-group alias before - # this payload is built, so cost-map lookup via that alias always misses. - # Fall back to the actual deployment model set by the router in metadata. - if base_model is None: - base_model = metadata.get("deployment") - custom_pricing: Final = logging_obj.custom_pricing_for(init_response_obj) + configured_base_model: Final = get_base_model_from_metadata(model_call_details=kwargs) + spilled_over: Final = logging_obj.spilled_over_ptu_for(init_response_obj) + # The router overrides completion_response.model to the model-group alias before this payload is + # built, so cost-map lookup via that alias misses and the deployment model stands in for base_model. + # A PTU request Azure spilled onto pay-as-you-go capacity is billed at the served model's rates, + # which the deployment's zeroed PTU entry would misreport, so that fallback is skipped there. + base_model: Final = ( + metadata.get("deployment") if configured_base_model is None and not spilled_over else configured_base_model + ) + custom_pricing: Final = False if spilled_over else logging_obj.custom_pricing_for(init_response_obj) raw_response_cost: Final = kwargs.get("response_cost") llm_response_cost: Final[float] = raw_response_cost or 0.0 guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information")) diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index ff82aacf393..ace57a64aae 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -4,6 +4,7 @@ from functools import reduce from typing import Final import httpx +from pydantic import TypeAdapter, ValidationError from litellm.constants import LITELLM_DETAILED_TIMING from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers @@ -101,6 +102,32 @@ def response_timing_metrics( return {"_response_ms": total_response_time_ms, "litellm_overhead_time_ms": overhead_ms} +_HIDDEN_PARAMS_DICT: Final = TypeAdapter(dict[str, object]) + + +def _hidden_params_mapping(hidden_params: object) -> Mapping[str, object]: + try: + return _HIDDEN_PARAMS_DICT.validate_python(hidden_params) + except ValidationError: + return {} + + +def _existing_provider_response_model(result: object) -> str | None: + """The pricing model a response already carries, like the decisions path sets, wins over the response model.""" + hidden_params: Final[object] = getattr(result, "_hidden_params", None) + carried: Final[object] = _hidden_params_mapping(hidden_params).get("provider_response_model") or getattr( + hidden_params, "provider_response_model", None + ) + return carried if isinstance(carried, str) and carried else None + + +def _provider_response_model(result: object) -> str | None: + """The model the provider answered with, kept because the proxy later rewrites ``result.model`` to the alias the + client asked for and cost-map lookups for the standard logging payload would otherwise miss.""" + model: Final = getattr(result, "model", None) if isinstance(result, ModelResponse) else None + return model if isinstance(model, str) and model else None + + class ResponseMetadata: """ Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses @@ -147,6 +174,8 @@ class ResponseMetadata: preserve_litellm_internal_headers=True, ), "litellm_model_name": model, + "provider_response_model": _existing_provider_response_model(self.result) + or _provider_response_model(self.result), } self._update_hidden_params(new_params) diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9f521ffc6cc..307353d081b 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -1,4 +1,6 @@ -from collections.abc import Callable +import threading +from collections.abc import Callable, Iterable, Iterator +from types import MethodType from typing import TYPE_CHECKING, Final import litellm @@ -15,6 +17,7 @@ else: _custom_logger_compatible_callbacks_literal = str _generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {} +_callback_lists_lock: Final = threading.RLock() class LoggingCallbackManager: @@ -43,16 +46,18 @@ class LoggingCallbackManager: Add a input callback to litellm.input_callback. Auto-routes async callbacks to litellm._async_input_callback. """ - if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) + with _callback_lists_lock: + if not isinstance(callback, str) and self._is_async_callable(callback): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) def add_litellm_service_callback(self, callback: CustomLogger | str | Callable[..., object]): """ Add a service callback to litellm.service_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) def add_litellm_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -60,10 +65,11 @@ class LoggingCallbackManager: Ensures no duplicates are added. """ - self._safe_add_callback_to_list( - callback=callback, - parent_list=litellm.callbacks, - ) + with _callback_lists_lock: + self._safe_add_callback_to_list( + callback=callback, + parent_list=litellm.callbacks, + ) def add_litellm_success_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -71,37 +77,41 @@ class LoggingCallbackManager: Auto-routes async callbacks to litellm._async_success_callback. Special-cases 'dynamodb' and 'openmeter' as async callbacks. """ - if ( - isinstance(callback, str) - and callback in ("dynamodb", "openmeter") - or not isinstance(callback, str) - and self._is_async_callable(callback) - ): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) + with _callback_lists_lock: + if ( + isinstance(callback, str) + and callback in ("dynamodb", "openmeter") + or not isinstance(callback, str) + and self._is_async_callable(callback) + ): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) def add_litellm_failure_callback(self, callback: CustomLogger | str | Callable[..., object]): """ Add a failure callback to `litellm.failure_callback`. Auto-routes async callbacks to litellm._async_failure_callback. """ - if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) + with _callback_lists_lock: + if not isinstance(callback, str) and self._is_async_callable(callback): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) def add_litellm_async_success_callback(self, callback: CustomLogger | Callable[..., object] | str): """ Add a success callback to litellm._async_success_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) def add_litellm_async_failure_callback(self, callback: CustomLogger | Callable[..., object] | str): """ Add a failure callback to litellm._async_failure_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) def remove_callback_from_list_by_object(self, callback_list, obj, require_self=True): """ @@ -110,13 +120,14 @@ class LoggingCallbackManager: if not isinstance(callback_list, list): # Not list -> do nothing return - if require_self: - remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] - else: - remove_list = [c for c in callback_list if c == obj] + with _callback_lists_lock: + if require_self: + remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] + else: + remove_list = [c for c in callback_list if c == obj] - for c in remove_list: - callback_list.remove(c) + for c in remove_list: + callback_list.remove(c) def remove_callbacks_by_type(self, callback_list, callback_type): """ @@ -134,10 +145,11 @@ class LoggingCallbackManager: if not isinstance(callback_list, list): return - remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)] + with _callback_lists_lock: + remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)] - for c in remove_list: - callback_list.remove(c) + for c in remove_list: + callback_list.remove(c) def _add_string_callback_to_list( self, callback: str, parent_list: list[CustomLogger | Callable[..., object] | str] @@ -336,6 +348,55 @@ class LoggingCallbackManager: return "-".join(key_parts) + get_custom_logger_key = _get_custom_logger_key + callback_lists_lock: Final = _callback_lists_lock + + def drop_callbacks(self, callbacks: Iterable[object], method_owners: Iterable[object] = ()) -> frozenset[str]: + dropped_ids: Final = frozenset(map(id, callbacks)) + owner_ids: Final = frozenset(map(id, method_owners)) + + def keep(callback: object) -> bool: + if id(callback) in dropped_ids: + return False + return not (isinstance(callback, MethodType) and id(callback.__self__) in owner_ids) + + with _callback_lists_lock: + vacated_keys: Final = frozenset( + self._get_custom_logger_key(callback) + for callback in litellm.callbacks + if isinstance(callback, CustomLogger) and id(callback) in dropped_ids + ) + litellm.callbacks = [callback for callback in litellm.callbacks if keep(callback)] + litellm.input_callback = [callback for callback in litellm.input_callback if keep(callback)] + litellm.success_callback = [callback for callback in litellm.success_callback if keep(callback)] + litellm.failure_callback = [callback for callback in litellm.failure_callback if keep(callback)] + litellm.service_callback = [callback for callback in litellm.service_callback if keep(callback)] + litellm._async_success_callback = [ + callback for callback in litellm._async_success_callback if keep(callback) + ] + litellm._async_failure_callback = [ + callback for callback in litellm._async_failure_callback if keep(callback) + ] + return vacated_keys + + def with_custom_logger_keys( + self, candidates: Iterable[CustomLogger], keys: frozenset[str] + ) -> tuple[CustomLogger, ...]: + if not keys: + return () + class_names: Final = frozenset(key.split("-", 1)[0] for key in keys) + same_class: Final = (candidate for candidate in candidates if type(candidate).__name__ in class_names) + return tuple(candidate for candidate in same_class if self._get_custom_logger_key(candidate) in keys) + + def without_custom_logger_keys( + self, callbacks: Iterable[str | Callable[..., object] | CustomLogger], keys: frozenset[str] + ) -> Iterator[str | Callable[..., object] | CustomLogger]: + return ( + callback + for callback in callbacks + if not (isinstance(callback, CustomLogger) and self._get_custom_logger_key(callback) in keys) + ) + def _reset_all_callbacks(self): """ Reset all callbacks to an empty list @@ -363,19 +424,21 @@ class LoggingCallbackManager: _get_all_callbacks = get_all_callbacks - def remove_callback_from_all_lists(self, obj, require_self=False) -> None: + def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None: """ Remove a callback object from every callback list it may have been promoted into, so a re-initialized callback leaves no stale instance behind. """ - for callback_list in ( - litellm.callbacks, - litellm.success_callback, - litellm.failure_callback, - litellm._async_success_callback, - litellm._async_failure_callback, - ): - self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) + with _callback_lists_lock: + for callback_list in ( + litellm.callbacks, + litellm.input_callback, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ): + self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) def get_active_additional_logging_utils_from_custom_logger( self, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ec426d6da5..5e43157a55d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -10,7 +10,7 @@ import traceback from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, NoReturn, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx @@ -60,6 +60,9 @@ from .exception_mapping_utils import exception_type, with_deployment_api_base from .llm_response_utils.get_api_base import get_api_base from .rules import Rules +if TYPE_CHECKING: + from .litellm_logging import BilledRoundsSpend + # Constants for special delta attribute names AUDIO_ATTRIBUTE: Final = "audio" IMAGE_ATTRIBUTE: Final = "images" @@ -259,6 +262,8 @@ def _provider_hidden_params( class CustomStreamWrapper: + _earlier_rounds_spend: "BilledRoundsSpend | None" = None + @property def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation return self._hidden_params @@ -342,6 +347,10 @@ class CustomStreamWrapper: self.messages = getattr(logging_obj, "messages", None) self.sent_stream_usage = False self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False + self.caller_requested_stream_usage: Final = ( + self.check_send_stream_usage(getattr(logging_obj, "stream_options", None)) + and getattr(logging_obj, "stream_usage_injected", False) is not True + ) self.tool_call = False # keep track of the returned chunks - used for calculating the input/output tokens for stream options self.chunks: list[ModelResponseStream] = [] @@ -2070,7 +2079,7 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: complete_streaming_response: Final = self._build_complete_streaming_response() - response = self.model_response_creator() + response = self._new_usage_chunk() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2222,27 +2231,10 @@ class CustomStreamWrapper: ) if _has_usage: - # Store a copy ONLY when usage stripping below will mutate - # the chunk. For non-usage chunks (vast majority), store - # directly to avoid expensive model_copy() per chunk. - self.chunks.append(processed_chunk.model_copy()) - - # Strip usage from the outgoing chunk so it's not sent twice - # (once in the chunk, once in _hidden_params). - obj_dict = processed_chunk.model_dump() - if "usage" in obj_dict: - del obj_dict["usage"] - processed_chunk = self.model_response_creator( - chunk=obj_dict, - hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage - Mapping[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR) - ), - ) - is_empty = is_model_response_stream_empty( - model_response=cast(ModelResponseStream, processed_chunk) - ) - if is_empty: + usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) + if usage_free_chunk is None: continue + processed_chunk = usage_free_chunk else: # No usage data — safe to store directly without copying self.chunks.append(processed_chunk) @@ -2273,13 +2265,18 @@ class CustomStreamWrapper: if processed_chunk is None: continue - choice = processed_chunk.choices[0] - if isinstance(choice, StreamingChoices): - self._accumulate_streamed_delta(choice.delta) + if processed_chunk.choices: + choice = processed_chunk.choices[0] + if isinstance(choice, StreamingChoices): + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) - # RETURN RESULT - self.chunks.append(processed_chunk) - return processed_chunk + if not self.send_stream_usage or getattr(processed_chunk, "usage", None) is None: + self.chunks.append(processed_chunk) + return processed_chunk + usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) + if usage_free_chunk is None: + continue + return usage_free_chunk except (StopAsyncIteration, StopIteration): return await self._finalize_completed_stream(cache_hit=cache_hit) except httpx.TimeoutException as e: # if httpx read timeout error occues @@ -2306,7 +2303,7 @@ class CustomStreamWrapper: async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": if self.sent_last_chunk is True: complete_streaming_response: Final = await self._async_build_complete_streaming_response() - response: Final = self.model_response_creator() + response: Final = self._new_usage_chunk() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2407,6 +2404,9 @@ class CustomStreamWrapper: return await asyncio.wait((self._failure_logging,), timeout=timeout) + def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: + self._earlier_rounds_spend = earlier + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2433,10 +2433,12 @@ class CustomStreamWrapper: if self.model: partial_response.model = self.model backfill_missing_cache_usage_fields(usage) - self.logging_obj.model_call_details["combined_usage_object"] = usage - self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 - ) + partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 + if self._earlier_rounds_spend is None: + self.logging_obj.model_call_details["combined_usage_object"] = usage + self.logging_obj.model_call_details["response_cost"] = partial_cost + return + self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, usage, partial_cost) except Exception as recover_error: verbose_logger.debug( "could not recover partial usage for interrupted stream: %s", @@ -2463,6 +2465,22 @@ class CustomStreamWrapper: return False return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) + def _new_usage_chunk(self) -> ModelResponseStream: + usage_chunk: Final = self.model_response_creator() + if self.caller_requested_stream_usage: + usage_chunk.choices = [] + return usage_chunk + + def _record_and_strip_usage_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: + self.chunks.append(chunk.model_copy()) + usage_free_chunk: Final = self.model_response_creator( + chunk={key: value for key, value in chunk.model_dump().items() if key != "usage"}, + hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage + Mapping[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR) + ), + ) + return None if is_model_response_stream_empty(model_response=usage_free_chunk) else usage_free_chunk + def _accumulate_streamed_delta(self, delta: object) -> None: # Shared by the sync, async, and non-aiohttp iteration sites so answer # text and the disqualifying-content latch stay in step across all three. diff --git a/litellm/llms/anthropic/pass_through/messages/interceptors/advisor.py b/litellm/llms/anthropic/pass_through/messages/interceptors/advisor.py index 03b27201398..85d76de5402 100644 --- a/litellm/llms/anthropic/pass_through/messages/interceptors/advisor.py +++ b/litellm/llms/anthropic/pass_through/messages/interceptors/advisor.py @@ -15,7 +15,7 @@ How it works: """ import uuid -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import litellm @@ -48,7 +48,7 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): def can_handle( self, - tools: list[dict] | None, + tools: Sequence[Mapping[str, object]] | None, custom_llm_provider: str | None, ) -> bool: if not tools: diff --git a/litellm/llms/anthropic/pass_through/messages/interceptors/base.py b/litellm/llms/anthropic/pass_through/messages/interceptors/base.py index c32b21564a7..3692c03b6a5 100644 --- a/litellm/llms/anthropic/pass_through/messages/interceptors/base.py +++ b/litellm/llms/anthropic/pass_through/messages/interceptors/base.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping, Sequence from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -21,7 +21,7 @@ class MessagesInterceptor(ABC): @abstractmethod def can_handle( self, - tools: list[dict] | None, + tools: Sequence[Mapping[str, object]] | None, custom_llm_provider: str | None, ) -> bool: """Return True if this interceptor should handle the request.""" diff --git a/litellm/llms/base_llm/extraction/__init__.py b/litellm/llms/base_llm/extraction/__init__.py new file mode 100644 index 00000000000..fe48996e8b4 --- /dev/null +++ b/litellm/llms/base_llm/extraction/__init__.py @@ -0,0 +1,3 @@ +from .transformation import BaseExtractionConfig + +__all__ = ["BaseExtractionConfig"] diff --git a/litellm/llms/base_llm/extraction/transformation.py b/litellm/llms/base_llm/extraction/transformation.py new file mode 100644 index 00000000000..8461256b883 --- /dev/null +++ b/litellm/llms/base_llm/extraction/transformation.py @@ -0,0 +1,86 @@ +from abc import ABC, abstractmethod +from collections.abc import Mapping +from typing import Final + +import httpx +from pydantic import TypeAdapter, ValidationError + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.extraction import ExtractionRequestBody, ExtractionResponse + +_PAYLOAD_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object) +_RESPONSE_ADAPTER: Final[TypeAdapter[ExtractionResponse]] = TypeAdapter(ExtractionResponse) +_RESERVED_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "content-type"}) + + +class BaseExtractionConfig(ABC): + api_key_env: tuple[str, ...] = () + api_base_env: tuple[str, ...] = () + + def get_default_api_base(self) -> str | None: + return None + + def missing_api_base_message(self, custom_llm_provider: str) -> str: + return f"api_base is required for Extraction provider '{custom_llm_provider}'" + + def resolve_api_base(self, api_base: str | None) -> str | None: + return api_base or self._first_secret(self.api_base_env) or self.get_default_api_base() + + def resolve_api_key(self, api_key: str | None) -> str | None: + return api_key or self._first_secret(self.api_key_env) + + @staticmethod + def _first_secret(names: tuple[str, ...]) -> str | None: + return next((value for value in (get_secret_str(name) for name in names) if value), None) + + def canonical_model(self, model: str) -> str: + return model + + def request_model(self, model: str) -> str: + return model + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> Mapping[str, str]: + return { + **{name: value for name, value in headers.items() if name.lower() not in _RESERVED_HEADERS}, + **({"Authorization": f"Bearer {api_key}"} if api_key is not None else {}), + "Content-Type": "application/json", + } + + @abstractmethod + def get_complete_url(self, api_base: str, model: str) -> str: ... + + def transform_extraction_request(self, model: str, request: ExtractionRequestBody) -> Mapping[str, object]: + return {**request.model_dump(mode="json", exclude_unset=True), "model": self.request_model(model)} + + def unwrap_response(self, payload: object) -> object: + return payload + + def transform_extraction_response( + self, + model: str, + custom_llm_provider: str, + raw_response: httpx.Response, + ) -> ExtractionResponse: + payload: Final[object] = _PAYLOAD_ADAPTER.validate_json(raw_response.content) + try: + return _RESPONSE_ADAPTER.validate_python(self.unwrap_response(payload)) + except ValidationError as error: + raise BaseLLMException( + status_code=500, + message=f"Extraction provider '{custom_llm_provider}' returned an unexpected response: {error}", + ) from error + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: Mapping[str, str] | httpx.Headers, + ) -> BaseLLMException: + return BaseLLMException(status_code=status_code, message=error_message, headers=httpx.Headers(headers)) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index f026e4f9a99..87d4fc0b32e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -60,7 +60,7 @@ def _response_header(headers: Mapping[str, str], name: str) -> str | None: return headers.get(name, None) -def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": +def bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": try: guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config) except ValidationError as e: @@ -440,7 +440,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): guardrail_headers: Final = ( () if raw_guardrail_config is None - else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) + else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) ) owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 9d895616336..edd30fff11b 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -14,6 +14,9 @@ from litellm.constants import AWS_ERROR_TYPE_HEADER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + bedrock_invoke_guardrail_headers, +) from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage from litellm.types.utils import ModelResponseStream @@ -42,6 +45,9 @@ if TYPE_CHECKING: BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"}) _BEDROCK_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) _BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"}) +_BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset( + {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} +) def _bedrock_endpoint_action(endpoint: str) -> str: @@ -123,6 +129,31 @@ def _with_invoke_request_metadata( return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body +def _converse_guardrail_config(action: str, guardrail_config: object) -> object: + fields: Final = as_json_object(guardrail_config) + if action == "converse-stream" or fields is None: + return guardrail_config + return {key: value for key, value in fields.items() if key != "streamProcessingMode"} + + +def _with_deployment_guardrail( + action: str, + headers: Mapping[str, str], + body: Mapping[str, object], + guardrail_config: object, +) -> tuple[Mapping[str, str], Mapping[str, object]]: + if guardrail_config is None: + return headers, body + if action in _BEDROCK_CONVERSE_ACTIONS: + return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)} + if action in _BEDROCK_INVOKE_ACTIONS: + guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items()) + return merge_bedrock_invoke_headers( + dict(headers), (), guardrail_headers, _BEDROCK_INVOKE_GUARDRAIL_HEADERS + ), body + return headers, body + + def _is_cache_point(entry: object) -> bool: fields: Final = as_json_object(entry) return fields is not None and "cachePoint" in fields @@ -453,11 +484,14 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD body=_converse_body_to_send(body) if action in _BEDROCK_CONVERSE_ACTIONS else body, litellm_params=litellm_params, ) + guarded_headers, guarded_body = _with_deployment_guardrail( + action, signed_headers, signed_body, litellm_params.get("guardrailConfig") + ) return self._sign_request( service_name="bedrock", - headers=dict(signed_headers), + headers=dict(guarded_headers), optional_params=optional_params, - request_data=dict(signed_body), + request_data=dict(guarded_body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py index 71a2dc78578..95a66387b28 100644 --- a/litellm/llms/bedrock/request_metadata.py +++ b/litellm/llms/bedrock/request_metadata.py @@ -256,12 +256,11 @@ def merge_bedrock_invoke_headers( """ Merge the ``X-Amzn-*`` headers the Invoke paths derive from params. - ``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is - the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's - headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry - proxy-authenticated identity into an AWS billing record that the caller must not be able to - write. Names are compared case-insensitively so a caller cannot leave a second spelling in - the dict and let the transport pick the winner. + ``caller_owned`` defers to a header the caller already set. ``proxy_owned_names`` are dropped + from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, for names + the caller must not be able to write, such as proxy-authenticated identity headed for an AWS + billing record or a deployment's guardrail. Names are compared case-insensitively so a caller + cannot leave a second spelling in the dict and let the transport pick the winner. """ if not caller_owned and not proxy_owned and not proxy_owned_names: return headers diff --git a/litellm/llms/cloudflare/common_utils.py b/litellm/llms/cloudflare/common_utils.py new file mode 100644 index 00000000000..9f1293cb9c4 --- /dev/null +++ b/litellm/llms/cloudflare/common_utils.py @@ -0,0 +1,53 @@ +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.secret_managers.main import ( + get_secret_str, + normalize_nonempty_secret_str, +) + +CLOUDFLARE_MISSING_API_BASE_MESSAGE: Final = ( + "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly" +) +_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def workers_ai_default_api_base() -> str | None: + account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID")) + if account_id is None: + return None + return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run" + + +def workers_ai_canonical_model(model: str) -> str: + if model.startswith("@cf/"): + return model + return f"@cf/cloudflare/{model}" + + +def workers_ai_request_model(model: str) -> str: + return model.rsplit("/", maxsplit=1)[-1] + + +def workers_ai_run_url(api_base: str, model: str) -> str: + normalized_api_base: Final = api_base.rstrip("/") + if normalized_api_base.endswith("/ai/v1"): + return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}" + if normalized_api_base.endswith("/ai/run"): + return f"{normalized_api_base}/{model}" + return f"{normalized_api_base}/ai/run/{model}" + + +def unwrap_workers_ai_result(payload: object) -> object: + try: + response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload) + except ValidationError: + return payload + if "answers" in response_mapping: + return response_mapping + try: + return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result")) + except ValidationError: + return response_mapping diff --git a/litellm/llms/cloudflare/decisions/transformation.py b/litellm/llms/cloudflare/decisions/transformation.py index 967935edfb4..7c776367886 100644 --- a/litellm/llms/cloudflare/decisions/transformation.py +++ b/litellm/llms/cloudflare/decisions/transformation.py @@ -1,16 +1,13 @@ -from collections.abc import Mapping -from typing import Final - -from pydantic import TypeAdapter, ValidationError - from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig -from litellm.secret_managers.main import ( - get_secret_str, - normalize_nonempty_secret_str, +from litellm.llms.cloudflare.common_utils import ( + CLOUDFLARE_MISSING_API_BASE_MESSAGE, + unwrap_workers_ai_result, + workers_ai_canonical_model, + workers_ai_default_api_base, + workers_ai_request_model, + workers_ai_run_url, ) -_RESPONSE_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) - class CloudflareDecisionsConfig(BaseDecisionsConfig): api_key_env = ("CLOUDFLARE_API_KEY",) @@ -23,38 +20,19 @@ class CloudflareDecisionsConfig(BaseDecisionsConfig): inlines_remote_images = True def get_default_api_base(self) -> str | None: - account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID")) - if account_id is None: - return None - return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run" + return workers_ai_default_api_base() def missing_api_base_message(self, custom_llm_provider: str) -> str: - return "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly" + return CLOUDFLARE_MISSING_API_BASE_MESSAGE def canonical_model(self, model: str) -> str: - if model.startswith("@cf/"): - return model - return f"@cf/cloudflare/{model}" + return workers_ai_canonical_model(model) def request_model(self, model: str) -> str: - return model.rsplit("/", maxsplit=1)[-1] + return workers_ai_request_model(model) def get_complete_url(self, api_base: str, model: str) -> str: - normalized_api_base: Final = api_base.rstrip("/") - if normalized_api_base.endswith("/ai/v1"): - return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}" - if normalized_api_base.endswith("/ai/run"): - return f"{normalized_api_base}/{model}" - return f"{normalized_api_base}/ai/run/{model}" + return workers_ai_run_url(api_base, model) def unwrap_response(self, payload: object) -> object: - try: - response_mapping: Final = _RESPONSE_MAPPING_ADAPTER.validate_python(payload) - except ValidationError: - return payload - if "answers" in response_mapping: - return response_mapping - try: - return _RESPONSE_MAPPING_ADAPTER.validate_python(response_mapping.get("result")) - except ValidationError: - return response_mapping + return unwrap_workers_ai_result(payload) diff --git a/litellm/llms/cloudflare/extraction/guardrail_translation/__init__.py b/litellm/llms/cloudflare/extraction/guardrail_translation/__init__.py new file mode 100644 index 00000000000..3fda941854a --- /dev/null +++ b/litellm/llms/cloudflare/extraction/guardrail_translation/__init__.py @@ -0,0 +1,11 @@ +from typing import Final + +from litellm.llms.cloudflare.extraction.guardrail_translation.handler import ExtractionHandler +from litellm.types.utils import CallTypes + +guardrail_translation_mappings: Final = { + CallTypes.extraction: ExtractionHandler, + CallTypes.aextraction: ExtractionHandler, +} + +__all__ = ["ExtractionHandler", "guardrail_translation_mappings"] diff --git a/litellm/llms/cloudflare/extraction/guardrail_translation/handler.py b/litellm/llms/cloudflare/extraction/guardrail_translation/handler.py new file mode 100644 index 00000000000..607d6fac48b --- /dev/null +++ b/litellm/llms/cloudflare/extraction/guardrail_translation/handler.py @@ -0,0 +1,276 @@ +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain, groupby +from types import MappingProxyType +from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias, TypeVar + +from pydantic import Field, JsonValue, TypeAdapter + +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.extraction import ( + EXTRACTION_STATE_REWRITTEN_METADATA_KEY, + ExtractionAnswer, + ExtractionBase64Image, + ExtractionResponse, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + +_FIELDS_ADAPTER: Final[TypeAdapter[Mapping[str, Mapping[str, object]]]] = TypeAdapter( + Mapping[str, Mapping[str, object]] +) +_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_IMAGES_ADAPTER: Final[TypeAdapter[tuple[ExtractionBase64Image | JsonValue, ...]]] = TypeAdapter( + tuple[Annotated[ExtractionBase64Image | JsonValue, Field(union_mode="left_to_right")], ...] +) + + +_AnswerKey: TypeAlias = Literal["value", "guess", "mean", "median", "p10", "p90", "quantiles"] +_ESTIMATE_KEYS: Final[tuple[_AnswerKey, ...]] = ("value", "guess", "mean", "median", "p10", "p90", "quantiles") +_STATISTIC_KEYS: Final[frozenset[_AnswerKey]] = frozenset(("mean", "median", "p10", "p90", "quantiles")) +_NULLED_ANSWER: Final[Mapping[str, None]] = MappingProxyType( + dict.fromkeys(("value", "guess", "mean", "median", "p10", "p90", "quantiles", "steps", "span", "spans"), None) +) + + +@dataclass(frozen=True, slots=True) +class _AnswerText: + name: str + key: _AnswerKey + item: int | None + text: str + + +_SPEC_TEXT_KEYS: Final = ("instructions", "unit") + + +@dataclass(frozen=True, slots=True) +class _SpecText: + name: str + key: str + text: str + + +def _spec_texts_of(name: str, spec: Mapping[str, object]) -> tuple[_SpecText, ...]: + return tuple( + _SpecText(name=name, key=key, text=text) for key in _SPEC_TEXT_KEYS if isinstance(text := spec.get(key), str) + ) + + +def _spec_texts(fields: Mapping[str, Mapping[str, object]]) -> tuple[_SpecText, ...]: + return tuple(chain.from_iterable(_spec_texts_of(name, spec) for name, spec in fields.items())) + + +_NamedText = TypeVar("_NamedText", _AnswerText, _SpecText) + + +def _scans_by_name(scanned: Iterable[tuple[_NamedText, str]]) -> Mapping[str, tuple[tuple[_NamedText, str], ...]]: + return {name: tuple(group) for name, group in groupby(scanned, key=lambda pair: pair[0].name)} + + +def _guardrailed_spec(spec: Mapping[str, object], scanned: Sequence[tuple[_SpecText, str]]) -> Mapping[str, object]: + return {**spec, **{text.key: guardrailed for text, guardrailed in scanned}} + + +def _metadata_key(data: Mapping[str, object]) -> Literal["metadata", "litellm_metadata"]: + return get_metadata_variable_name_from_kwargs(data) + + +def _metadata(data: Mapping[str, object]) -> Mapping[str, object]: + return _METADATA_ADAPTER.validate_python(data.get(_metadata_key(data)) or {}) + + +def _is_text_answer(answer: ExtractionAnswer, spec: Mapping[str, object] | None) -> bool: + return answer.type == "text" or (answer.type == "list" and (spec is None or spec.get("items") == "text")) + + +def _image_urls(images: object) -> tuple[str, ...]: + if not isinstance(images, list): + return () + return tuple( + f"data:{image.content_type};base64,{image.base64}" if isinstance(image, ExtractionBase64Image) else image + for image in _IMAGES_ADAPTER.validate_python(images) + if isinstance(image, (str, ExtractionBase64Image)) + ) + + +def _scalar_text(value: JsonValue, whole: bool) -> str | None: + if isinstance(value, bool): + return None + if isinstance(value, float) and whole: + return format(value, ".0f") + return str(value) if isinstance(value, (str, int, float)) else None + + +def _item_texts(items: Sequence[JsonValue], whole: bool) -> tuple[tuple[int, str], ...]: + return tuple((index, text) for index, item in enumerate(items) if (text := _scalar_text(item, whole)) is not None) + + +def _value_texts(name: str, key: _AnswerKey, value: JsonValue) -> tuple[_AnswerText, ...]: + whole: Final = key in _STATISTIC_KEYS + if isinstance(value, list): + return tuple( + _AnswerText(name=name, key=key, item=index, text=text) for index, text in _item_texts(value, whole) + ) + text: Final = _scalar_text(value, whole) + return () if text is None else (_AnswerText(name=name, key=key, item=None, text=text),) + + +def _step_texts(name: str, step: JsonValue) -> tuple[_AnswerText, ...]: + if not isinstance(step, dict): + return () + return tuple(chain.from_iterable(_value_texts(name, key, step.get(key)) for key in _ESTIMATE_KEYS)) + + +def _answer_texts(name: str, answer: ExtractionAnswer) -> tuple[_AnswerText, ...]: + return ( + *_value_texts(name, "value", answer.value), + *_value_texts(name, "guess", answer.guess), + *_value_texts(name, "mean", answer.mean), + *_value_texts(name, "median", answer.median), + *_value_texts(name, "p10", answer.p10), + *_value_texts(name, "p90", answer.p90), + *_value_texts(name, "quantiles", list(answer.quantiles or ())), + *chain.from_iterable(_step_texts(name, step) for step in answer.steps or ()), + ) + + +def _guardrailed_value(value: JsonValue, scanned: Mapping[int | None, str]) -> JsonValue: + if isinstance(value, list): + return [scanned.get(index, item) for index, item in enumerate(value)] + return scanned.get(None, value) + + +def _scans_of(key: _AnswerKey, scanned: Sequence[tuple[_AnswerText, str]]) -> Mapping[int | None, str]: + return {text.item: guardrailed for text, guardrailed in scanned if text.key == key} + + +def _guardrailed_answer( + answer: ExtractionAnswer, + scanned: Sequence[tuple[_AnswerText, str]], + drop_spans: bool, + is_text: bool, +) -> ExtractionAnswer: + if not is_text and any(text.text != guardrailed for text, guardrailed in scanned): + return answer.model_copy(update=_NULLED_ANSWER) + value_scans: Final = _scans_of("value", scanned if is_text else ()) + guess_scans: Final = _scans_of("guess", scanned if is_text else ()) + update: Final = { + **({"value": _guardrailed_value(answer.value, value_scans)} if value_scans else {}), + **({"guess": _guardrailed_value(answer.guess, guess_scans)} if guess_scans else {}), + **({"span": None, "spans": None} if drop_spans else {}), + } + return answer.model_copy(update=update) if update else answer + + +async def _scanned_texts( + texts: Sequence[str], + model: object, + guardrail_to_apply: "CustomGuardrail", + request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail takes the request as a plain dict + input_type: Literal["request", "response"], + litellm_logging_obj: "LiteLLMLoggingObj | None", + images: Sequence[str] = (), +) -> tuple[str, ...]: + unique_texts: Final = tuple(dict.fromkeys(texts)) + inputs: Final = GenericGuardrailAPIInputs(texts=list(unique_texts), model=model if isinstance(model, str) else None) + if images: + inputs["images"] = list(images) + guardrailed: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.apply_guardrail types request_data as a bare dict + inputs=inputs, + request_data=request_data, + input_type=input_type, + logging_obj=litellm_logging_obj, + ) + returned: Final = guardrailed.get("texts", []) + guardrailed_by_text: Final = MappingProxyType( + {text: returned[index] if index < len(returned) else text for index, text in enumerate(unique_texts)} + ) + return tuple(guardrailed_by_text[text] for text in texts) + + +class ExtractionHandler(BaseTranslation): + async def process_input_messages( + self, + data: dict[str, object], # mutable-ok: BaseTranslation hands over the proxy's request dict + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> Mapping[str, object]: + state: Final = data.get("state") + states: Final = (state,) if isinstance(state, str) and state else () + fields: Final = _FIELDS_ADAPTER.validate_python(data.get("fields") or {}) + spec_texts: Final = _spec_texts(fields) + images: Final = _image_urls(data.get("images")) + if not states and not spec_texts and not images: + return data + scanned: Final = await _scanned_texts( + (*states, *(text.text for text in spec_texts)), + data.get("model"), + guardrail_to_apply, + data, + "request", + litellm_logging_obj, + images, + ) + spec_scans: Final = _scans_by_name(zip(spec_texts, scanned[len(states) :])) + state_rewritten: Final = bool(states) and scanned[0] != states[0] + return { + **data, + **({"state": scanned[0]} if states else {}), + **( + {_metadata_key(data): {**_metadata(data), EXTRACTION_STATE_REWRITTEN_METADATA_KEY: True}} + if state_rewritten + else {} + ), + "fields": {name: _guardrailed_spec(spec, spec_scans.get(name, ())) for name, spec in fields.items()}, + } + + async def process_output_response( + self, + response: object, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, + request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation hands over the request dict + ) -> object: + if not isinstance(response, ExtractionResponse): + return response + texts: Final = tuple( + chain.from_iterable(_answer_texts(name, answer) for name, answer in response.answers.items()) + ) + state_texts: Final = (response.state_text,) if response.state_text else () + request: Final = request_data or {} + scanned: Final = ( + await _scanned_texts( + (*(text.text for text in texts), *state_texts), + response.model, + guardrail_to_apply, + request, + "response", + litellm_logging_obj, + ) + if texts or state_texts + else () + ) + answer_scans: Final = _scans_by_name(zip(texts, scanned)) + state_text: Final = next(iter(scanned[len(texts) :]), response.state_text) + specs: Final = _FIELDS_ADAPTER.validate_python(request.get("fields") or {}) + return response.model_copy( + update={ + "answers": { + name: _guardrailed_answer( + answer, + answer_scans.get(name, ()), + drop_spans=state_text != response.state_text, + is_text=_is_text_answer(answer, specs.get(name)), + ) + for name, answer in response.answers.items() + }, + **({"state_text": state_text} if state_texts else {}), + } + ) diff --git a/litellm/llms/cloudflare/extraction/transformation.py b/litellm/llms/cloudflare/extraction/transformation.py new file mode 100644 index 00000000000..7aa3bf6316c --- /dev/null +++ b/litellm/llms/cloudflare/extraction/transformation.py @@ -0,0 +1,32 @@ +from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig +from litellm.llms.cloudflare.common_utils import ( + CLOUDFLARE_MISSING_API_BASE_MESSAGE, + unwrap_workers_ai_result, + workers_ai_canonical_model, + workers_ai_default_api_base, + workers_ai_request_model, + workers_ai_run_url, +) + + +class CloudflareExtractionConfig(BaseExtractionConfig): + api_key_env = ("CLOUDFLARE_API_KEY",) + api_base_env = ("CLOUDFLARE_API_BASE",) + + def get_default_api_base(self) -> str | None: + return workers_ai_default_api_base() + + def missing_api_base_message(self, custom_llm_provider: str) -> str: + return CLOUDFLARE_MISSING_API_BASE_MESSAGE + + def canonical_model(self, model: str) -> str: + return workers_ai_canonical_model(model) + + def request_model(self, model: str) -> str: + return workers_ai_request_model(model) + + def get_complete_url(self, api_base: str, model: str) -> str: + return workers_ai_run_url(api_base, model) + + def unwrap_response(self, payload: object) -> object: + return unwrap_workers_ai_result(payload) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b1e55375f16..ac2ac1fe31b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -76,6 +76,7 @@ from litellm.llms.base_llm.containers.transformation import BaseContainerConfig from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig +from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, BaseFileUploadStream, @@ -128,6 +129,7 @@ from litellm.types.containers.main import ( DeleteContainerResult, ) from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse +from litellm.types.extraction import ExtractionResponse from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, @@ -1700,6 +1702,115 @@ class BaseLLMHTTPHandler: model=model, custom_llm_provider=custom_llm_provider, raw_response=response, request=request ) + def _prepare_extraction_request( + self, + model: str, + logging_obj: LiteLLMLoggingObj | None, + provider_config: BaseExtractionConfig, + body: Mapping[str, object], + api_base: str, + api_key: str | None, + headers: Mapping[str, str], + litellm_params: Mapping[str, object], + ) -> tuple[str, Mapping[str, str]]: + outbound_headers: Final = provider_config.validate_environment( + headers=headers, + model=model, + api_key=api_key, + litellm_params=litellm_params, + ) + url: Final = provider_config.get_complete_url(api_base=api_base, model=model) + if logging_obj is not None: + logging_obj.pre_call( # pyright: ignore[reportUnknownMemberType] # Logging.pre_call is untyped + input=dict(body), + api_key=api_key, + model=model, + additional_args={"api_base": url, "complete_input_dict": dict(body), "headers": outbound_headers}, + ) + return url, outbound_headers + + def extraction( + self, + model: str, + custom_llm_provider: str, + logging_obj: LiteLLMLoggingObj | None, + provider_config: BaseExtractionConfig, + body: Mapping[str, object], + api_base: str, + api_key: str | None, + headers: Mapping[str, str], + timeout: float | httpx.Timeout | None, + litellm_params: Mapping[str, object], + client: HTTPHandler | None = None, + ) -> ExtractionResponse: + url, outbound_headers = self._prepare_extraction_request( + model=model, + logging_obj=logging_obj, + provider_config=provider_config, + body=body, + api_base=api_base, + api_key=api_key, + headers=headers, + litellm_params=litellm_params, + ) + sync_httpx_client: Final = client if client is not None else get_httpx_client() + try: + response: Final = sync_httpx_client.post( # pyright: ignore[reportUnknownMemberType] # untyped + url, + json=dict(body), + headers=dict(outbound_headers), + timeout=timeout, + logging_obj=logging_obj, + ) + except httpx.HTTPError as e: + raise self._handle_error(e=e, provider_config=provider_config) + return provider_config.transform_extraction_response( + model=model, custom_llm_provider=custom_llm_provider, raw_response=response + ) + + async def aextraction( + self, + model: str, + custom_llm_provider: str, + logging_obj: LiteLLMLoggingObj | None, + provider_config: BaseExtractionConfig, + body: Mapping[str, object], + api_base: str, + api_key: str | None, + headers: Mapping[str, str], + timeout: float | httpx.Timeout | None, + litellm_params: Mapping[str, object], + client: AsyncHTTPHandler | None = None, + ) -> ExtractionResponse: + url, outbound_headers = self._prepare_extraction_request( + model=model, + logging_obj=logging_obj, + provider_config=provider_config, + body=body, + api_base=api_base, + api_key=api_key, + headers=headers, + litellm_params=litellm_params, + ) + async_httpx_client: Final = ( + client + if client is not None + else get_async_httpx_client(llm_provider=litellm.LlmProviders(custom_llm_provider)) + ) + try: + response: Final = await async_httpx_client.post( # pyright: ignore[reportUnknownMemberType] # untyped + url, + json=dict(body), + headers=dict(outbound_headers), + timeout=timeout, + logging_obj=logging_obj, + ) + except httpx.HTTPError as e: + raise self._handle_error(e=e, provider_config=provider_config) + return provider_config.transform_extraction_response( + model=model, custom_llm_provider=custom_llm_provider, raw_response=response + ) + def _prepare_audio_transcription_request( self, model: str, @@ -6462,6 +6573,7 @@ class BaseLLMHTTPHandler: provider_config: Union[ BaseConfig, BaseDecisionsConfig, + BaseExtractionConfig, BaseRerankConfig, BaseResponsesAPIConfig, BaseImageEditConfig, diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 143652ee77c..936919af83d 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -82,9 +82,9 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): return "tts" in model def get_supported_openai_params(self, model: str) -> list[str]: + sampling_params: Final = ("temperature", "top_p") if self.supports_sampling_params(model) else () supported_params: Final = [ - "temperature", - "top_p", + *sampling_params, "max_tokens", "max_completion_tokens", "stream", diff --git a/litellm/llms/github_copilot/github_host.py b/litellm/llms/github_copilot/github_host.py new file mode 100644 index 00000000000..4fc921d490d --- /dev/null +++ b/litellm/llms/github_copilot/github_host.py @@ -0,0 +1,78 @@ +import os +import re +from dataclasses import dataclass +from typing import Final + +from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST + +from .authenticator import ( + DEFAULT_GITHUB_ACCESS_TOKEN_URL, + DEFAULT_GITHUB_API_KEY_URL, + DEFAULT_GITHUB_CLIENT_ID, + DEFAULT_GITHUB_DEVICE_CODE_URL, +) + +_GHE_SUBDOMAIN_PATTERN: Final = re.compile(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?") + + +@dataclass(frozen=True, slots=True) +class GithubHostEndpoints: + host: str + device_code_url: str + access_token_url: str + api_key_url: str + user_url: str + client_id: str + copilot_api_hosts: frozenset[str] + copilot_api_host_suffix: str | None + default_copilot_api_base: str + + +def normalized_github_host(raw: object) -> str | None: + if raw is None or raw == "": + return GITHUB_COPILOT_DEFAULT_HOST + if not isinstance(raw, str): + return None + + host: Final = raw.lower() + if host == GITHUB_COPILOT_DEFAULT_HOST: + return host + if not host.endswith(".ghe.com"): + return None + + subdomain: Final = host.removesuffix(".ghe.com") + return host if _GHE_SUBDOMAIN_PATTERN.fullmatch(subdomain) is not None else None + + +def github_host_endpoints(host: str) -> GithubHostEndpoints: + normalized: Final = normalized_github_host(host) + if normalized is None: + raise ValueError("host must be github.com or a single-subdomain GHE.com hostname") + client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) + if normalized == GITHUB_COPILOT_DEFAULT_HOST: + return GithubHostEndpoints( + host=normalized, + device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL), + access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL), + api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL), + user_url="https://api.github.com/user", + client_id=client_id, + copilot_api_hosts=frozenset({"githubcopilot.com"}), + copilot_api_host_suffix=".githubcopilot.com", + default_copilot_api_base="https://api.githubcopilot.com", + ) + + subdomain: Final = normalized.removesuffix(".ghe.com") + api_host: Final = f"api.{subdomain}.ghe.com" + copilot_api_host: Final = f"copilot-api.{subdomain}.ghe.com" + return GithubHostEndpoints( + host=normalized, + device_code_url=f"https://{normalized}/login/device/code", + access_token_url=f"https://{normalized}/login/oauth/access_token", + api_key_url=f"https://{api_host}/copilot_internal/v2/token", + user_url=f"https://{api_host}/user", + client_id=client_id, + copilot_api_hosts=frozenset({copilot_api_host}), + copilot_api_host_suffix=None, + default_copilot_api_base=f"https://{copilot_api_host}", + ) diff --git a/litellm/llms/github_copilot/per_user_auth.py b/litellm/llms/github_copilot/per_user_auth.py index 3dc8c12fe06..75a10f2e671 100644 --- a/litellm/llms/github_copilot/per_user_auth.py +++ b/litellm/llms/github_copilot/per_user_auth.py @@ -5,8 +5,7 @@ touched on the per-user path. """ import asyncio -import hashlib -import os +import hmac import threading import time from collections.abc import Callable, Mapping @@ -27,6 +26,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( GITHUB_COPILOT_AUTH_TYPE_KEY, + GITHUB_COPILOT_HOST_KEY, GITHUB_COPILOT_PER_USER_AUTH_TYPE, GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS, ) @@ -38,14 +38,8 @@ from litellm.exceptions import ( ServiceUnavailableError, ) -from .authenticator import ( - DEFAULT_GITHUB_ACCESS_TOKEN_URL, - DEFAULT_GITHUB_API_KEY_URL, - DEFAULT_GITHUB_CLIENT_ID, - DEFAULT_GITHUB_DEVICE_CODE_URL, - github_api_headers, -) -from .common_utils import DEFAULT_GITHUB_COPILOT_API_BASE +from .authenticator import github_api_headers +from .github_host import GithubHostEndpoints, github_host_endpoints, normalized_github_host GITHUB_COPILOT_USER_SESSION_KWARG_KEY: Final = "github_copilot_user_session" _LLM_PROVIDER: Final = "github_copilot" @@ -98,7 +92,7 @@ _SESSION_CACHE: Final = cast( # cast-ok: pins the untyped module cache to _Sess ) _EXCHANGE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT)) _IN_FLIGHT: Final[ # mutable-ok: single-flight dedup registry mutated under its own locking - dict[str, "asyncio.Future[GithubCopilotUserSession]"] + dict[str, tuple[str, "asyncio.Future[GithubCopilotUserSession]"]] ] = {} @@ -141,7 +135,7 @@ _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload) _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload) _GITHUB_USER: Final = TypeAdapter(_GithubUserPayload) _COPILOT_ENDPOINTS: Final = TypeAdapter(_CopilotEndpointsPayload) -_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, float]) +_CACHED_SESSION_ENTRY: Final = TypeAdapter(tuple[str, str, str, float]) @dataclass(frozen=True, slots=True, repr=False) @@ -178,53 +172,72 @@ class GithubCopilotDeviceFlowPoll: access_token: str | None = field(default=None, repr=False) -def validated_copilot_api_base(endpoints_api: object) -> str: - """Only a genuine *.githubcopilot.com HTTPS endpoint may receive a Copilot - token; anything else falls back to the default base.""" +def validated_copilot_api_base(endpoints_api: object, endpoints: GithubHostEndpoints) -> str: if not isinstance(endpoints_api, str): - return DEFAULT_GITHUB_COPILOT_API_BASE - parsed: Final = urlparse(endpoints_api) + return endpoints.default_copilot_api_base + try: + parsed: Final = urlparse(endpoints_api) + port: Final = parsed.port + except ValueError: + return endpoints.default_copilot_api_base hostname: Final = (parsed.hostname or "").lower() if ( parsed.scheme.lower() != "https" - or not (hostname == "githubcopilot.com" or hostname.endswith(".githubcopilot.com")) + or not ( + hostname in endpoints.copilot_api_hosts + or (endpoints.copilot_api_host_suffix is not None and hostname.endswith(endpoints.copilot_api_host_suffix)) + ) or parsed.username is not None or parsed.password is not None + or port not in (None, 443) ): - return DEFAULT_GITHUB_COPILOT_API_BASE - try: - port: Final = parsed.port - except ValueError: - return DEFAULT_GITHUB_COPILOT_API_BASE - if port not in (None, 443): - return DEFAULT_GITHUB_COPILOT_API_BASE + return endpoints.default_copilot_api_base return endpoints_api.rstrip("/") -def _session_cache_key(user_id: str, github_token: str) -> str: - return hashlib.sha256(f"{user_id}\x00{github_token}".encode()).hexdigest() +def github_copilot_host_endpoints(credential_name: str) -> GithubHostEndpoints: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + + credential: Final = CredentialAccessor.find_credential(credential_name) + host_raw: Final[object] = ( + credential.credential_values.get(GITHUB_COPILOT_HOST_KEY) if credential is not None else None + ) + host: Final = normalized_github_host(host_raw) + if host is None: + raise BadRequestError( + message=f"Credential '{credential_name}' has an invalid {GITHUB_COPILOT_HOST_KEY}", + llm_provider=_LLM_PROVIDER, + model="", + ) + return github_host_endpoints(host) -def evict_copilot_user_session(user_id: str, github_token: str) -> None: - _SESSION_CACHE.delete_cache(_session_cache_key(user_id, github_token)) +def _session_cache_key(user_id: str, credential_name: str, github_host: str) -> str: + return f"{user_id}\x00{credential_name}\x00{github_host}" -def _cached_session(cache_key: str) -> GithubCopilotUserSession | None: +def evict_copilot_user_session(user_id: str, credential_name: str, github_host: str) -> None: + _SESSION_CACHE.delete_cache(_session_cache_key(user_id, credential_name, github_host)) + + +def _cached_session(cache_key: str, github_token: str) -> GithubCopilotUserSession | None: cached: Final[object] = _SESSION_CACHE.get_cache(cache_key) try: - token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached) + cached_github_token, token, api_base, expires_at = _CACHED_SESSION_ENTRY.validate_python(cached) except ValidationError: return None + if not hmac.compare_digest(cached_github_token, github_token): + return None if expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS <= time.time(): return None return GithubCopilotUserSession(token=token, api_base=api_base) -def _cache_session(cache_key: str, session: GithubCopilotUserSession, expires_at: float) -> None: +def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUserSession, expires_at: float) -> None: ttl: Final = expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS - time.time() if ttl <= 0: return - _SESSION_CACHE.set_cache(cache_key, (session.token, session.api_base, expires_at), ttl=ttl) + _SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl) def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationError: @@ -238,7 +251,13 @@ def _reject_connection(credential_name: str) -> CallerCredentialAuthenticationEr ) -def _session_from_response(response: httpx.Response, credential_name: str, cache_key: str) -> GithubCopilotUserSession: +def _session_from_response( + response: httpx.Response, + github_token: str, + credential_name: str, + cache_key: str, + endpoints: GithubHostEndpoints, +) -> GithubCopilotUserSession: if response.status_code in (401, 403, 404): _SESSION_CACHE.delete_cache(cache_key) raise _reject_connection(credential_name) @@ -263,15 +282,17 @@ def _session_from_response(response: httpx.Response, credential_name: str, cache model="", ) from e - endpoints: Final[object] = payload.get("endpoints") + payload_endpoints: Final[object] = payload.get("endpoints") endpoints_api: Final[object] = ( - _COPILOT_ENDPOINTS.validate_python(endpoints).get("api") if isinstance(endpoints, Mapping) else None + _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api") + if isinstance(payload_endpoints, Mapping) + else None ) - api_base: Final = validated_copilot_api_base(endpoints_api) + api_base: Final = validated_copilot_api_base(endpoints_api, endpoints) session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base) expires_at: Final = _expires_at_seconds(payload.get("expires_at")) if expires_at > 0: - _cache_session(cache_key, session, expires_at) + _cache_session(cache_key, github_token, session, expires_at) return session @@ -303,7 +324,12 @@ def _expires_at_seconds(raw: object) -> float: return 0.0 -def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession: +def _do_exchange( + github_token: str, + credential_name: str, + cache_key: str, + endpoints: GithubHostEndpoints, +) -> GithubCopilotUserSession: import litellm try: @@ -311,7 +337,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git _SyncGetClient, litellm.module_level_client ) response: Final = client.get( - DEFAULT_GITHUB_API_KEY_URL, + endpoints.api_key_url, headers=github_api_headers(github_token), ) except Exception as e: @@ -320,7 +346,7 @@ def _do_exchange(github_token: str, credential_name: str, cache_key: str) -> Git llm_provider=_LLM_PROVIDER, model="", ) from e - return _session_from_response(response, credential_name, cache_key) + return _session_from_response(response, github_token, credential_name, cache_key, endpoints) def _new_async_github_client() -> _AsyncGitHubClient: @@ -332,11 +358,16 @@ def _new_async_github_client() -> _AsyncGitHubClient: )(_LLM_PROVIDER) -async def _do_aexchange(github_token: str, credential_name: str, cache_key: str) -> GithubCopilotUserSession: +async def _do_aexchange( + github_token: str, + credential_name: str, + cache_key: str, + endpoints: GithubHostEndpoints, +) -> GithubCopilotUserSession: client: Final = _new_async_github_client() try: response: Final = await client.get( - DEFAULT_GITHUB_API_KEY_URL, + endpoints.api_key_url, headers=github_api_headers(github_token), ) except Exception as e: @@ -345,43 +376,50 @@ async def _do_aexchange(github_token: str, credential_name: str, cache_key: str) llm_provider=_LLM_PROVIDER, model="", ) from e - return _session_from_response(response, credential_name, cache_key) + return _session_from_response(response, github_token, credential_name, cache_key, endpoints) def exchange_github_token( user_id: str, github_token: str, credential_name: str, + endpoints: GithubHostEndpoints, ) -> GithubCopilotUserSession: - cache_key: Final = _session_cache_key(user_id, github_token) - cached: Final = _cached_session(cache_key) + cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host) + cached: Final = _cached_session(cache_key, github_token) if cached is not None: return cached - lock: Final = _EXCHANGE_LOCKS[int(cache_key, 16) % _SYNC_LOCK_COUNT] + lock: Final = _EXCHANGE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT] with lock: - rechecked: Final = _cached_session(cache_key) + rechecked: Final = _cached_session(cache_key, github_token) if rechecked is not None: return rechecked - return _do_exchange(github_token, credential_name, cache_key) + return _do_exchange(github_token, credential_name, cache_key, endpoints) async def aexchange_github_token( user_id: str, github_token: str, credential_name: str, + endpoints: GithubHostEndpoints, ) -> GithubCopilotUserSession: - cache_key: Final = _session_cache_key(user_id, github_token) - cached: Final = _cached_session(cache_key) + cache_key: Final = _session_cache_key(user_id, credential_name, endpoints.host) + cached: Final = _cached_session(cache_key, github_token) if cached is not None: return cached loop: Final = asyncio.get_running_loop() - in_flight: Final = _IN_FLIGHT.get(cache_key) - if in_flight is not None and not in_flight.done() and in_flight.get_loop() is loop: - return await asyncio.shield(in_flight) + in_flight_entry: Final = _IN_FLIGHT.get(cache_key) + if in_flight_entry is not None: + in_flight_token, in_flight = in_flight_entry + if not hmac.compare_digest(in_flight_token, github_token): + return await _do_aexchange(github_token, credential_name, cache_key, endpoints) + if not in_flight.done() and in_flight.get_loop() is loop: + return await asyncio.shield(in_flight) future: Final = loop.create_future() - _IN_FLIGHT[cache_key] = future + entry: Final = (github_token, future) + _IN_FLIGHT[cache_key] = entry try: - session: Final = await _do_aexchange(github_token, credential_name, cache_key) + session: Final = await _do_aexchange(github_token, credential_name, cache_key, endpoints) except BaseException as e: if not future.done(): future.set_exception(e) @@ -392,19 +430,17 @@ async def aexchange_github_token( future.set_result(session) return session finally: - if _IN_FLIGHT.get(cache_key) is future: + if _IN_FLIGHT.get(cache_key) is entry: _IN_FLIGHT.pop(cache_key, None) -async def astart_device_flow() -> GithubCopilotDeviceFlowStart: +async def astart_device_flow(endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowStart: client: Final = _new_async_github_client() - device_code_url: Final = os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL) - client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) try: response: Final = await client.post( - device_code_url, + endpoints.device_code_url, headers=github_api_headers(), - json={"client_id": client_id, "scope": _DEVICE_FLOW_SCOPE}, + json={"client_id": endpoints.client_id, "scope": _DEVICE_FLOW_SCOPE}, ) except httpx.HTTPStatusError as e: raise ServiceUnavailableError( @@ -435,16 +471,14 @@ async def astart_device_flow() -> GithubCopilotDeviceFlowStart: ) -async def apoll_device_flow(device_code: str) -> GithubCopilotDeviceFlowPoll: +async def apoll_device_flow(device_code: str, endpoints: GithubHostEndpoints) -> GithubCopilotDeviceFlowPoll: client: Final = _new_async_github_client() - access_token_url: Final = os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL) - client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) try: response: Final = await client.post( - access_token_url, + endpoints.access_token_url, headers=github_api_headers(), json={ - "client_id": client_id, + "client_id": endpoints.client_id, "device_code": device_code, "grant_type": "urn:ietf:params:oauth:grant-type:device_code", }, @@ -500,11 +534,11 @@ def _poll_payload_from_response(response: httpx.Response) -> GithubCopilotDevice ) -async def afetch_github_login(github_token: str) -> str: +async def afetch_github_login(github_token: str, endpoints: GithubHostEndpoints) -> str: client: Final = _new_async_github_client() try: response: Final = await client.get( - "https://api.github.com/user", + endpoints.user_url, headers=github_api_headers(github_token), ) except Exception as e: @@ -530,11 +564,17 @@ async def afetch_github_login(github_token: str) -> str: return payload["login"] -async def acheck_copilot_seat(user_id: str, github_token: str, credential_name: str) -> GithubCopilotUserSession: +async def acheck_copilot_seat( + user_id: str, + github_token: str, + credential_name: str, + endpoints: GithubHostEndpoints, +) -> GithubCopilotUserSession: return await aexchange_github_token( user_id=user_id, github_token=github_token, credential_name=credential_name, + endpoints=endpoints, ) @@ -567,7 +607,11 @@ def _connect_error(credential_name: str) -> CallerCredentialAuthenticationError: ) -def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tuple[str, str]: +def _caller_connection( + kwargs: Mapping[str, object], + credential_name: str, + endpoints: GithubHostEndpoints, +) -> tuple[str, str]: secret_fields_raw: Final[object] = kwargs.get("secret_fields") secret_fields: Final = ( cast(Mapping[object, object], secret_fields_raw) # cast-ok: value is Mapping-checked or dict-shaped at runtime @@ -585,8 +629,21 @@ def _caller_connection(kwargs: Mapping[str, object], credential_name: str) -> tu if isinstance(credentials_raw, Mapping) else None ) - token: Final[object] = credentials.get(credential_name) if credentials is not None else None - if not isinstance(user_id, str) or not user_id or not isinstance(token, str) or not token: + connection_raw: Final[object] = credentials.get(credential_name) if credentials is not None else None + connection: Final = ( + cast(Mapping[object, object], connection_raw) # cast-ok: runtime Mapping check validates the stored connection + if isinstance(connection_raw, Mapping) + else None + ) + token: Final[object] = connection.get("access_token") if connection is not None else None + github_host: Final = normalized_github_host(connection.get("github_host") if connection is not None else None) + if ( + not isinstance(user_id, str) + or not user_id + or not isinstance(token, str) + or not token + or github_host != endpoints.host + ): raise _connect_error(credential_name) return user_id, token @@ -705,11 +762,13 @@ def attach_github_copilot_user_session( if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession): return credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else "" - user_id, github_token = _caller_connection(kwargs, credential_name) + endpoints: Final = github_copilot_host_endpoints(credential_name) + user_id, github_token = _caller_connection(kwargs, credential_name, endpoints) session: Final = exchange_github_token( user_id=user_id, github_token=github_token, credential_name=credential_name, + endpoints=endpoints, ) kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel _strip_caller_authorization(kwargs) @@ -725,11 +784,13 @@ async def aattach_github_copilot_user_session( if isinstance(kwargs.get(GITHUB_COPILOT_USER_SESSION_KWARG_KEY), GithubCopilotUserSession): return credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else "" - user_id, github_token = _caller_connection(kwargs, credential_name) + endpoints: Final = github_copilot_host_endpoints(credential_name) + user_id, github_token = _caller_connection(kwargs, credential_name, endpoints) session: Final = await aexchange_github_token( user_id=user_id, github_token=github_token, credential_name=credential_name, + endpoints=endpoints, ) kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY] = session # rebind-ok: kwargs is the request mutation channel _strip_caller_authorization(kwargs) diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index 651a2215052..748203b22d9 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -16,7 +16,9 @@ The handler is intentionally a thin orchestration layer: 3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS) vs self-hosted (any host:port, optional TLS) modes. 4. Stream the audio through Riva's ``streaming_response_generator`` and - aggregate ``is_final`` results into a single transcript. + aggregate ``is_final`` results into a single transcript, or, when + ``riva_offline`` is set, send the whole buffer through the unary + ``offline_recognize`` RPC for NIMs that only ship offline engines. 5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed on ``_hidden_params`` so cost calculation works. @@ -26,13 +28,19 @@ without the optional STT extras installed. import asyncio import inspect -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol +import httpx + +import litellm +from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_name, process_audio_file, ) +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.llms.nvidia_riva.audio_transcription.audio_utils import ( resample_to_riva_pcm, ) @@ -45,6 +53,7 @@ from litellm.llms.nvidia_riva.common_utils import ( NvidiaRivaException, grpc_error_to_litellm_exception, ) +from litellm.secret_managers.main import str_to_bool from litellm.types.utils import FileTypes, TranscriptionResponse from litellm.utils import convert_to_model_response_object @@ -67,10 +76,20 @@ class _RivaAuth(Protocol): """Opaque ``riva.client.Auth`` handle.""" +class _RecognizeFuture(Protocol): + def result(self, timeout: float | None = None) -> object: ... + + def done(self) -> bool: ... + + def cancel(self) -> bool: ... + + class _AsrService(Protocol): @property def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ... + def offline_recognize(self, audio_bytes: bytes, config: object, future: bool) -> _RecognizeFuture: ... + class _EndpointingConfig(Protocol): """Opaque ``EndpointingConfig`` protobuf message.""" @@ -128,6 +147,38 @@ def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding: return riva_asr_module.AudioEncoding +def _parse_riva_offline(value: object, drop_params: bool) -> bool: + match value: + case None: + return False + case bool(): + return value + case str() if (parsed := str_to_bool(value)) is not None: + return parsed + case _ if drop_params: + return False + case _: + raise NvidiaRivaException( + status_code=400, + message=f"riva_offline must be true or false, got {value!r}. Set drop_params to ignore it.", + ) + + +def _deadline_seconds(timeout: float | str | httpx.Timeout, drop_params: bool) -> float: + try: + limit: Final = CompletionTimeout.as_limit(timeout) + except ValueError: + if drop_params: + return COMPLETION_HTTP_FALLBACK_SECONDS + raise NvidiaRivaException( + status_code=400, + message=f"timeout must be a number of seconds, got {timeout!r}. Set drop_params to ignore it.", + ) from None + if not isinstance(limit, httpx.Timeout): + return limit + return COMPLETION_HTTP_FALLBACK_SECONDS if limit.read is None else limit.read + + class NvidiaRivaAudioTranscription: """Sync + async entry point for Riva ASR.""" @@ -138,7 +189,7 @@ class NvidiaRivaAudioTranscription: optional_params: dict, litellm_params: dict, model_response: TranscriptionResponse, - timeout: float, + timeout: float | str | httpx.Timeout | None, logging_obj: "LiteLLMLoggingObj", api_key: str | None, api_base: str | None, @@ -183,7 +234,7 @@ class NvidiaRivaAudioTranscription: optional_params: dict, litellm_params: dict, model_response: TranscriptionResponse, - timeout: float, + timeout: float | str | httpx.Timeout | None, logging_obj: "LiteLLMLoggingObj", api_key: str | None, api_base: str | None, @@ -213,7 +264,7 @@ class NvidiaRivaAudioTranscription: optional_params: dict, litellm_params: dict, model_response: TranscriptionResponse, - timeout: float, + timeout: float | str | httpx.Timeout | None, logging_obj: "LiteLLMLoggingObj", api_key: str | None, api_base: str | None, @@ -259,6 +310,10 @@ class NvidiaRivaAudioTranscription: response_format: Final = request_payload.get("response_format") or "json" timestamp_granularities: Final = request_payload.get("timestamp_granularities") + drop_params: Final = ( + "drop_params" in litellm_params and litellm_params["drop_params"] is True + ) or litellm.drop_params is True + riva_offline: Final = _parse_riva_offline(optional_params.get("riva_offline"), drop_params=drop_params) riva_module, riva_asr_module = _import_riva() auth_obj: Final = self._construct_auth( @@ -272,10 +327,6 @@ class NvidiaRivaAudioTranscription: riva_asr_module=riva_asr_module, recognition_config_dict=recognition_config_dict, ) - streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig( - config=recognition_config, interim_results=False - ) - logging_obj.pre_call( input=None, api_key=api_key, @@ -286,24 +337,31 @@ class NvidiaRivaAudioTranscription: "recognition_config": recognition_config_dict, "nvcf_function_id_set": bool(optional_params.get("nvcf_function_id")), "use_ssl": optional_params.get("use_ssl"), + "riva_offline": riva_offline, }, }, ) try: asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj) - audio_chunks: Final = self._iter_audio_chunks(resampled.pcm_bytes) - stream_kwargs: Final[dict[str, object]] = { - "audio_chunks": audio_chunks, - "streaming_config": streaming_config, - } - # Forward the deadline so the stream cannot block forever if the - # server stalls. Older riva-client versions do not accept a - # ``timeout`` kwarg, so pass it only when supported. - if timeout is not None and self._supports_timeout_kwarg(asr_service.streaming_response_generator): - stream_kwargs["timeout"] = float(timeout) - stream: Final = asr_service.streaming_response_generator(**stream_kwargs) - final_results: Final = self._collect_final_results(stream) + responses: Final = ( + self._offline_responses( + asr_service=asr_service, + pcm_bytes=resampled.pcm_bytes, + recognition_config=recognition_config, + timeout=None if timeout is None else _deadline_seconds(timeout, drop_params), + ) + if riva_offline + else self._streaming_responses( + asr_service=asr_service, + riva_asr_module=riva_asr_module, + pcm_bytes=resampled.pcm_bytes, + recognition_config=recognition_config, + timeout=timeout, + drop_params=drop_params, + ) + ) + final_results: Final = self._collect_final_results(responses, require_final=not riva_offline) except NvidiaRivaException: raise except Exception as e: @@ -409,6 +467,54 @@ class NvidiaRivaAudioTranscription: return config + def _streaming_responses( + self, + asr_service: _AsrService, + riva_asr_module: _RivaAsrModule, + pcm_bytes: bytes, + recognition_config: _RecognitionConfig, + timeout: float | str | httpx.Timeout | None, + drop_params: bool, + ) -> Iterable[object]: + streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig( + config=recognition_config, interim_results=False + ) + return asr_service.streaming_response_generator( + audio_chunks=self._iter_audio_chunks(pcm_bytes), + streaming_config=streaming_config, + **self._timeout_kwargs(asr_service.streaming_response_generator, timeout, drop_params), + ) + + def _offline_responses( + self, + asr_service: _AsrService, + pcm_bytes: bytes, + recognition_config: _RecognitionConfig, + timeout: float | None, + ) -> Iterable[object]: + pending: Final = asr_service.offline_recognize(audio_bytes=pcm_bytes, config=recognition_config, future=True) + try: + return (pending.result(timeout=timeout),) + except Exception as e: + if pending.done(): + raise + pending.cancel() + raise NvidiaRivaException( + status_code=408, + message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds", + ) from e + + @classmethod + def _timeout_kwargs( + cls, rpc: Callable[..., object], timeout: float | str | httpx.Timeout | None, drop_params: bool + ) -> Mapping[str, float]: + # Forward the deadline so the call cannot block forever if the + # server stalls. Older riva-client versions do not accept a + # ``timeout`` kwarg, so pass it only when supported. + if timeout is None or not cls._supports_timeout_kwarg(rpc): + return MappingProxyType({}) + return MappingProxyType({"timeout": _deadline_seconds(timeout, drop_params)}) + @staticmethod def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool: try: @@ -429,9 +535,9 @@ class NvidiaRivaAudioTranscription: yield chunk @staticmethod - def _collect_final_results(stream) -> list[dict[str, object]]: + def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> list[dict[str, object]]: """ - Walk the gRPC stream, ignore empty / non-final chunks, and return a + Walk the gRPC responses, ignore empty / non-final chunks, and return a list of normalized final-result dicts. Matching the user's note: the ``id`` blocks with no ``results`` are streaming heartbeats and must be skipped. @@ -440,7 +546,7 @@ class NvidiaRivaAudioTranscription: for response in stream: results = getattr(response, "results", None) or [] for result in results: - if not getattr(result, "is_final", False): + if require_final and not getattr(result, "is_final", False): continue alternatives = getattr(result, "alternatives", None) or [] if not alternatives: diff --git a/litellm/llms/nvidia_riva/common_utils.py b/litellm/llms/nvidia_riva/common_utils.py index b654e725dc8..c671b14e742 100644 --- a/litellm/llms/nvidia_riva/common_utils.py +++ b/litellm/llms/nvidia_riva/common_utils.py @@ -41,6 +41,9 @@ _GRPC_STATUS_CODE_TO_HTTP: Final[dict] = { } +_GRPC_MESSAGE_TOO_LARGE: Final = "message larger than max" + + def _extract_grpc_status_name(error: object) -> str | None: """ Best-effort extraction of a gRPC StatusCode name from an arbitrary error. @@ -79,8 +82,16 @@ def grpc_error_to_litellm_exception(error: Exception) -> NvidiaRivaException: a ``NvidiaRivaException`` with an appropriate HTTP-equivalent status code. """ status_name: Final = _extract_grpc_status_name(error) - http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500) - detail: Final = _extract_grpc_details(error) or str(error) + if status_name == "RESOURCE_EXHAUSTED" and _GRPC_MESSAGE_TOO_LARGE in detail: + return NvidiaRivaException( + status_code=413, + message=( + f"NVIDIA Riva gRPC error ({status_name}): {detail}. The audio is too large for one offline " + "request to this Riva server: send shorter audio, or drop riva_offline on models that support streaming." + ), + ) + + http_status: Final = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500) message = f"NVIDIA Riva gRPC error ({status_name}): {detail}" if status_name else f"NVIDIA Riva error: {detail}" return NvidiaRivaException(status_code=http_status, message=message) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index ea71e5fef50..176569f5d8d 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -1107,14 +1107,18 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ wrapper to call VertexGeminiConfig.map_openai_params """ - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - config: Final = VertexGeminiConfig() - _model: Final = openai_request_body.get("model", "") + raw_model: Final = openai_request_body.get("model", "") + model: Final = raw_model if isinstance(raw_model, str) else "" + sampling_params: Final = tuple( + p for p in ("temperature", "top_p", "top_k") if openai_request_body.get(p) is not None + ) + if sampling_params and not VertexGeminiConfig.supports_sampling_params(model): + VertexGeminiConfig.reject_unsupported_sampling_param( + model=model, param=", ".join(sampling_params), drop_params=False + ) vertex_params: Final = config.map_openai_params( - model=_model, + model=model, non_default_params=openai_request_body, optional_params={}, drop_params=False, diff --git a/litellm/llms/vertex_ai/gemini/long_decode.py b/litellm/llms/vertex_ai/gemini/long_decode.py index 81d6c04d0e8..571fa94d4a5 100644 --- a/litellm/llms/vertex_ai/gemini/long_decode.py +++ b/litellm/llms/vertex_ai/gemini/long_decode.py @@ -2,17 +2,20 @@ import json from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from itertools import accumulate, chain -from typing import Final, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys +from typing import Final, Literal, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys import httpx +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.exceptions import Timeout +from litellm.litellm_core_utils.env_utils import get_env_int_or_none from litellm.types.llms.vertex_ai import ( Candidates, Citation, CitationMetadata, GenerateContentResponseBody, + GenerationConfig, GroundingMetadata, GroundingSegment, GroundingSupport, @@ -48,8 +51,14 @@ HOP_FAILURES: Final = (VertexAIError, httpx.HTTPError, Timeout) HOP_END_FINISH_REASON: Final = "CONTINUATION" CUT_OFF_FINISH_REASONS: Final = frozenset({"MAX_TOKENS", HOP_END_FINISH_REASON}) +DEFAULT_OUTPUT_BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" +_OUTPUT_LIMIT_KEYS: Final = frozenset({"max_output_tokens", "maxOutputTokens"}) +_REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET: Final = frozenset({"generation_config", "continuationToken"}) +_TEXT_ONLY_MODALITIES: Final = TypeAdapter(list[Literal["TEXT"]]) + _VERTEX_V1_PATH: Final = "/v1/projects/" _VERTEX_PATH_ACCEPTING_CONTINUATION_TOKEN: Final = "/v1beta1/projects/" +_CANDIDATE_COUNT_KEYS: Final = ("candidate_count", "candidateCount") LongDecodeHops: TypeAlias = tuple[GenerateContentResponseBody, ...] AsyncHopPost: TypeAlias = Callable[[RequestBody], Awaitable[httpx.Response]] @@ -76,15 +85,53 @@ def _int_or_none(value: object) -> int | None: return value if isinstance(value, int) else None +def _asks_for_one_candidate(generation_config: Mapping[str, object]) -> bool: + return all( + _int_or_none(generation_config[key] or 1) == 1 for key in _CANDIDATE_COUNT_KEYS if key in generation_config + ) + + def continuation_cap_of(request_body: RequestBody) -> int | None: generation_config: Final = request_body.get("generationConfig") if generation_config is None or "max_output_tokens" not in generation_config: return None - if "candidate_count" in generation_config and _int_or_none(generation_config["candidate_count"]) != 1: + if not _asks_for_one_candidate(generation_config): return None return _int_or_none(generation_config["max_output_tokens"]) +def _asks_for_text_only(generation_config: Mapping[str, object]) -> bool: + if "response_modalities" in generation_config: + return False + try: + _ = _TEXT_ONLY_MODALITIES.validate_python(generation_config.get("responseModalities", [])) + except ValidationError: + return False + return True + + +def with_default_output_budget( + request_body: RequestBody, model_max_output_tokens: Callable[[], int | None] +) -> RequestBody: + generation_config: Final = request_body.get("generationConfig", GenerationConfig()) + request_json: Final[Mapping[str, object]] = request_body + budget: Final = get_env_int_or_none(DEFAULT_OUTPUT_BUDGET_ENV) + if budget is None or budget <= 0: + return request_body + if not _REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET.isdisjoint(request_body): + return request_body + if not isinstance(request_json.get("generationConfig", {}), dict): + return request_body + if not _OUTPUT_LIMIT_KEYS.isdisjoint(generation_config): + return request_body + if not _asks_for_text_only(generation_config): + return request_body + model_max: Final = model_max_output_tokens() + if model_max is None: + return request_body + return {**request_body, "generationConfig": {**generation_config, "max_output_tokens": min(budget, model_max)}} + + def client_continuation_token(candidate: Candidates) -> str | None: if candidate.get("finishReason") not in CUT_OFF_FINISH_REASONS: return None diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 7b832b321da..71bb05ee155 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,15 +7,19 @@ Why separate file? Make it easy to see how transformation works import os import re from collections.abc import Callable, Mapping +from functools import partial from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote import httpx -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, PositiveInt, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_image_mime_type_from_url, ) @@ -64,6 +68,7 @@ from ..common_utils import ( get_supports_response_schema, get_supports_system_message, ) +from .long_decode import with_default_output_budget # VertexBase is imported only for type checking to avoid introducing a # module-load-time cyclic import to vertex_llm_base. The instance is lazily @@ -868,7 +873,7 @@ def gemini_convert_messages_with_history( """ has_text_in_content = check_text_in_content(user_content) if has_text_in_content is False: - verbose_logger.warning( + verbose_logger.debug( "No text in user content. Adding a blank text to user content, to ensure Gemini doesn't fail the request. Relevant Issue - https://github.com/BerriAI/litellm/issues/5515" ) user_content.append( @@ -1189,7 +1194,18 @@ def transform_request_body( # labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata) labels: Final = pop_vertex_request_labels(optional_params, litellm_params) - filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, set(config_fields))} + from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + drops_top_k: Final = "top_k" in optional_params and not VertexGeminiConfig.supports_sampling_params(model) + if drops_top_k: + request_params: Final = cast(Mapping[str, object], litellm_params) # cast-ok: untyped shared signature + VertexGeminiConfig.reject_unsupported_sampling_param( + model=model, + param="top_k", + drop_params=drops_unsupported_params(request_params.get("drop_params")), + ) + config_keys: Final = set(config_fields) - {"top_k"} if drops_top_k else set(config_fields) + filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, config_keys)} generation_config: Final[GenerationConfig | None] = GenerationConfig(**filtered_params) @@ -1292,13 +1308,75 @@ def sync_transform_request_body( vertex_auth_header=vertex_auth_header, ) - return transform_request_body( - messages=messages, - model=model, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - cached_content=cached_content, - optional_params=optional_params, + return _with_default_output_budget( + transform_request_body( + messages=messages, + model=model, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + cached_content=cached_content, + optional_params=optional_params, + ), + model, + custom_llm_provider, + litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict + ) + + +class _CostMapOutputModalities(TypedDict): + supported_output_modalities: NotRequired[ReadOnly[tuple[str, ...]]] + + +_COST_MAP_OUTPUT_MODALITIES: Final = TypeAdapter(_CostMapOutputModalities) + + +def _answers_text(cost_map_key: str) -> bool: + try: + row: Final = _COST_MAP_OUTPUT_MODALITIES.validate_python(litellm.model_cost.get(cost_map_key)) # pyright: ignore[reportUnknownMemberType] # cost map rows are untyped JSON + except ValidationError: + return False + return "text" in row.get("supported_output_modalities", ("text",)) + + +class _OutputLimit(TypedDict): + max_output_tokens: NotRequired[ReadOnly[PositiveInt | None]] + + +_OUTPUT_LIMIT: Final = TypeAdapter(_OutputLimit) + + +def _output_limit(fields: object) -> int | None: + try: + return _OUTPUT_LIMIT.validate_python(fields).get("max_output_tokens") + except ValidationError: + return None + + +def _catalog_capped_output_limit(cost_map_key: str, override: object) -> int | None: + limits: Final = ( + _output_limit(GetModelCostMap.loaded_model_cost_map().get(cost_map_key)), + _output_limit(override), + ) + return min((limit for limit in limits if limit is not None), default=None) + + +def _text_chat_model_max_output_tokens( + model: str, custom_llm_provider: str, deployment_model_info: object +) -> int | None: + info: Final = litellm.utils.lookup_model_info_for_registration(model, custom_llm_provider) + if isinstance(info, litellm.utils.ModelInfoLookupMiss) or info.get("mode") != "chat": + return None + if not _answers_text(info["key"]): + return None + return _catalog_capped_output_limit(info["key"], deployment_model_info or info) + + +def _with_default_output_budget( + request_body: RequestBody, model: str, custom_llm_provider: str, deployment_model_info: object +) -> RequestBody: + return with_default_output_budget( + request_body, + partial(_text_chat_model_max_output_tokens, model, custom_llm_provider, deployment_model_info), ) @@ -1368,21 +1446,8 @@ async def async_transform_request_body( messages, should_inline=should_inline_remote_media(custom_llm_provider) ) - if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): - # _transform_request_body may issue a sync httpx.get (up to 5s timeout) - # via _get_gcs_object_content_type to fetch GCS object metadata. Run the - # whole sync transformation on a worker thread so it does not block the - # async event loop. - return await asyncify(transform_request_body)( - messages=inlined_messages, - model=model, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - cached_content=cached_content, - optional_params=optional_params, - ) - - return transform_request_body( + transform: Final = partial( + transform_request_body, messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, @@ -1390,6 +1455,21 @@ async def async_transform_request_body( cached_content=cached_content, optional_params=optional_params, ) + # _transform_request_body may issue a sync httpx.get (up to 5s timeout) + # via _get_gcs_object_content_type to fetch GCS object metadata. Run the + # whole sync transformation on a worker thread so it does not block the + # async event loop. + request_body: Final[RequestBody] = ( + await asyncify(transform)() + if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages) + else transform() + ) + return _with_default_output_budget( + request_body, + model, + custom_llm_provider, + litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict + ) def _default_user_message_when_system_message_passed() -> ChatCompletionUserMessage: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index ec40cdfd11f..241b2d2a3a0 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3,10 +3,11 @@ ## Initial implementation - covers gemini + image gen calls import json import re +import threading import time from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy -from functools import partial +from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args @@ -28,6 +29,7 @@ from litellm.constants import ( ) from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator from litellm.litellm_core_utils.prompt_templates.factory import ( encode_tool_call_id_with_signature, @@ -196,6 +198,29 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr ) +_SAMPLING_DEPRECATION_UNMAPPED_MODELS: Final = "models outside litellm's model map" +_sampling_deprecation_lock: Final = threading.Lock() + + +@lru_cache(maxsize=None) +def _log_gemini_sampling_params_deprecated(model: str) -> None: + verbose_logger.warning( + "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.", + model, + ) + + +def _sampling_deprecation_notice_key(model: str) -> str: + catalog: Final = GetModelCostMap.loaded_model_cost_map() + is_cataloged: Final = any(name in catalog for name in (model, f"vertex_ai/{model}", f"gemini/{model}")) + return model if is_cataloged else _SAMPLING_DEPRECATION_UNMAPPED_MODELS + + +def _warn_gemini_sampling_params_deprecated(model: str) -> None: + with _sampling_deprecation_lock: + _log_gemini_sampling_params_deprecated(_sampling_deprecation_notice_key(model)) + + def _cache_creation_tokens(logging_obj: LoggingClass) -> int: recorded: Final[object] = getattr(logging_obj, "vertex_ai_cache_creation_input_tokens", 0) return recorded if isinstance(recorded, int) and recorded > 0 else 0 @@ -404,11 +429,43 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return False return True + @staticmethod + def supports_sampling_params(model: str) -> bool: + """Whether Google still honors ``temperature`` / ``top_p`` / ``top_k`` for ``model``. + + Driven by ``supports_sampling_params`` in the model map: Gemini 3.6 Flash, 3.5 Flash-Lite + and later accept the fields but ignore them, and a later generation will 400 on them. + A model the map says nothing about keeps forwarding them. ``temperature`` and ``top_p`` + leave ``get_supported_openai_params`` when the flag is off, so the shared unsupported-param + check raises or drops them and ``allowed_openai_params`` still forwards them verbatim; + ``top_k`` is a provider-specific kwarg that bypasses that check and is gated while the + request body is built.""" + rows: Final = cast(Mapping[str, Mapping[str, object]], litellm.model_cost) # cast-ok: JSON rows by model id + bare: Final = model.split("/")[-1] + for candidate in (model, bare, f"vertex_ai/{bare}", f"gemini/{bare}"): + if isinstance(flag := rows.get(candidate, {}).get("supports_sampling_params"), bool): + return flag + return True + + @staticmethod + def reject_unsupported_sampling_param(model: str, param: str, drop_params: bool) -> None: + """Return under ``drop_params`` so the caller drops ``param``, else raise the client-side 400 + that names the escape hatch.""" + if litellm.drop_params or drop_params: + return + raise UnsupportedParamsError( + message=( + f"{model} does not support {param}: the model map flags it `supports_sampling_params: false` " + "because Google ignores sampling params on this model. " + "To drop unsupported params, set `litellm.drop_params = True`." + ), + status_code=400, + ) + def get_supported_openai_params(self, model: str) -> list[str]: + sampling_params: Final = ("temperature", "top_p", "top_k") if self.supports_sampling_params(model) else () supported_params: Final = [ - "temperature", - "top_p", - "top_k", + *sampling_params, "max_tokens", "max_completion_tokens", "stream", @@ -1266,8 +1323,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): allow_integer_enums: bool = True, ) -> dict: self._apply_include_server_side_tool_invocations(non_default_params, optional_params) - gemini_sampling_params_warned: bool = False + supports_sampling_params: Final = self.supports_sampling_params(model) for param, value in non_default_params.items(): + if param in ("temperature", "top_p", "top_k") and not supports_sampling_params: + continue if param == "temperature": if VertexGeminiConfig._is_gemini_3_or_newer(model): if value is not None and value < 1.0: @@ -1275,28 +1334,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "Warning: Setting temperature < 1.0 for Gemini 3 models (%s) can cause infinite loops, degraded reasoning performance, and failure on complex tasks. Strongly recommended to use temperature = 1.0 (default).", model, ) - if not gemini_sampling_params_warned: - verbose_logger.warning( - "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.", - model, - ) - gemini_sampling_params_warned = True + _warn_gemini_sampling_params_deprecated(model) optional_params["temperature"] = value elif param == "top_p": - if VertexGeminiConfig._is_gemini_3_or_newer(model) and not gemini_sampling_params_warned: - verbose_logger.warning( - "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.", - model, - ) - gemini_sampling_params_warned = True + if VertexGeminiConfig._is_gemini_3_or_newer(model): + _warn_gemini_sampling_params_deprecated(model) optional_params["top_p"] = value elif param == "top_k": - if VertexGeminiConfig._is_gemini_3_or_newer(model) and not gemini_sampling_params_warned: - verbose_logger.warning( - "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to function for Gemini 3+ (%s) but are planned for removal in a future release. Move sampling guidance into the `system` instructions instead.", - model, - ) - gemini_sampling_params_warned = True + if VertexGeminiConfig._is_gemini_3_or_newer(model): + _warn_gemini_sampling_params_deprecated(model) optional_params["top_k"] = value elif ( param == "stream" and value is True diff --git a/litellm/main.py b/litellm/main.py index 63375292077..d647cab24e3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5988,6 +5988,7 @@ def completion( *OPENAI_WIF_KWARGS_KEYS, *OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS, "github_copilot_auth_type", + "github_copilot_host", "github_copilot_user_session", PROVIDER_AFFINITY_HEADER_KWARG_KEY, "fireworks_forward_user_id", diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index cb6ed8d5c75..cdb9a925403 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -1,18 +1,21 @@ import inspect from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping +from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.anthropic.pass_through.messages import handler as main -from litellm.llms.anthropic.pass_through.safeguards import native_route_classifies -from litellm.rust_bridge.catalog import Route, RouteContext +from litellm.rust_bridge import detectors +from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target from litellm.rust_bridge.dispatch import PublicDispatch +from litellm.rust_bridge.messages import detectors as route_detectors from litellm.rust_bridge.messages.entrypoints import ( NATIVE_AMESSAGES, NATIVE_MESSAGES, ) from litellm.rust_bridge.public_call import ( + AsyncFlag, NativeCall, bind, native_call, @@ -46,6 +49,7 @@ def _python_amessages() -> PythonAmessages: _PYTHON_MESSAGES: Final = _python_messages() _MESSAGES: Final = signature(_PYTHON_MESSAGES) +_FLAG: Final = AsyncFlag("is_async", sync=_MESSAGES, is_set=bool) _PYTHON_AMESSAGES: Final = _python_amessages() _AMESSAGES: Final = signature(_PYTHON_AMESSAGES) @@ -71,45 +75,76 @@ def _public_request( return native_call(legacy, args, _native_kwargs(fields, kwargs)) -def _resolved_provider(request: NativeCall) -> str | None: - try: - return get_llm_provider( - str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider")) - )[1] - except BadRequestError: - return optional_str(request.resolved.get("custom_llm_provider")) - - -def _python_serves(request: NativeCall) -> bool: - """Vertex AI serves only Claude through the Anthropic Messages config, as - ``get_provider_anthropic_messages_config`` does; every other Vertex model keeps Python's - translation through Chat Completions. A request Python would answer with classifier - ``safeguard_results`` stays on Python too.""" +def _resolved(request: NativeCall) -> tuple[str, str | None]: model: Final = str(request.resolved["model"]) - vertex_without_claude: Final = _resolved_provider(request) == "vertex_ai" and "claude" not in model.lower() - return vertex_without_claude or native_route_classifies(model, request.kwargs.get("safeguards")) + provider: Final = optional_str(request.resolved.get("custom_llm_provider")) + try: + resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider) + except BadRequestError: + return model, provider + return resolved_model, resolved_provider -def _context(request: NativeCall) -> RouteContext: +def _target(request: NativeCall, model: str, provider: str | None) -> Target: + requested_model: Final = str(request.resolved["model"]) + model_info: Final = request.resolved.get("model_info") + api_base: Final = request.resolved.get("api_base") + native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup + model, provider, model_info, api_base if isinstance(api_base, str) else None + ) + if native_config is not None: + return Target.MESSAGES + compacting: Final = request.resolved.get("compaction") is not None + if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice + return Target.RESPONSES + return Target.CHAT_COMPLETIONS + + +_FEATURES: Final[detectors.Detectors] = MappingProxyType( + { + **detectors.COMMON, + Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, + Feature.RESPONSE_CACHE: detectors.response_cache, + Feature.MODIFY_PARAMS: detectors.modify_params, + Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, + Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, + Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection, + Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry, + } +) + +_ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType( + { + **_FEATURES, + Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook, + Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception, + Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted, + } +) + + +def _context(request: NativeCall, features: detectors.Detectors) -> RouteContext: + requested_model: Final = str(request.resolved["model"]) + model, provider = _resolved(request) return RouteContext( Route.MESSAGES, - provider=_resolved_provider(request), - model=str(request.resolved["model"]), + provider=provider, + model=requested_model, + target=_target(request, model, provider), + features=detectors.detected(features, request), ) _DISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs), - context=_context, - bypass=lambda request: request.kwargs.get("is_async") is True or _python_serves(request), + context=lambda request: _context(request, _FEATURES), ) _ADISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs), - context=_context, - bypass=_python_serves, + context=lambda request: _context(request, _ASYNC_FEATURES), ) @@ -118,9 +153,11 @@ def anthropic_messages_handler( **kwargs: object, # kwargs-ok: preserve the public Anthropic Messages call shape ) -> MessagesResult | Coroutine[object, object, MessagesResult]: python: Final = _PYTHON_MESSAGES + if _FLAG.requested(kwargs): + return anthropic_messages(**_FLAG.keywords(args, kwargs)) return _DISPATCH.run( args, - kwargs, + _FLAG.dropped(kwargs), python=python, binding=NATIVE_MESSAGES, native=native_call_hook, @@ -131,7 +168,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult: python: Final = _PYTHON_AMESSAGES return await _ADISPATCH.arun( args, - kwargs, + _FLAG.dropped(kwargs), python=python, binding=NATIVE_AMESSAGES, native=native_call_hook, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 15eba7f8af8..3d89af9d691 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16580,6 +16580,24 @@ "source": "https://developers.cloudflare.com/workers-ai/platform/pricing/", "supports_vision": true }, + "cloudflare/clairic": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 65536, + "mode": "extraction", + "output_cost_per_token": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/clairic/", + "supports_vision": true + }, + "cloudflare/@cf/cloudflare/clairic": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 65536, + "mode": "extraction", + "output_cost_per_token": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/clairic/", + "supports_vision": true + }, "cloudflare/@cf/meta/llama-2-7b-chat-fp16": { "input_cost_per_token": 1.923e-06, "litellm_provider": "cloudflare", @@ -21903,6 +21921,7 @@ ], "supports_function_calling": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true }, @@ -21927,6 +21946,7 @@ ], "supports_function_calling": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true }, @@ -21958,6 +21978,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true }, @@ -22021,6 +22042,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true }, @@ -29417,6 +29439,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -30139,6 +30162,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -30200,6 +30224,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -30268,6 +30293,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -30317,6 +30343,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": false, "supports_url_context": false, @@ -31743,6 +31770,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -31927,6 +31955,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -31989,6 +32018,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32051,6 +32081,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32420,6 +32451,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32481,6 +32513,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32549,6 +32582,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32598,6 +32632,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": false, "supports_url_context": false, @@ -33192,6 +33227,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_vision": true }, "github_copilot/gemini-3.8-flash": { @@ -33210,6 +33246,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_vision": true }, "github_copilot/gpt-3.5-turbo": { @@ -52605,6 +52642,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -65823,6 +65861,7 @@ "supports_tool_choice": true, "supports_response_schema": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_vision": true, "source": "https://deepinfra.com/pricing", "supports_audio_input": true @@ -69264,6 +69303,7 @@ "perplexity/google/gemini-3.5-flash-lite": { "litellm_provider": "perplexity", "mode": "responses", + "supports_sampling_params": false, "supports_web_search": true, "supports_function_calling": true, "input_cost_per_token": 3e-07, @@ -69274,6 +69314,7 @@ "perplexity/google/gemini-3.6-flash": { "litellm_provider": "perplexity", "mode": "responses", + "supports_sampling_params": false, "supports_web_search": true, "supports_function_calling": true, "input_cost_per_token": 1.5e-06, @@ -69284,6 +69325,7 @@ "perplexity/google/gemini-3.7-flash": { "litellm_provider": "perplexity", "mode": "responses", + "supports_sampling_params": false, "supports_web_search": true, "supports_function_calling": true, "input_cost_per_token": 7.5e-07, @@ -71375,6 +71417,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_vision": true, "supports_web_search": true }, @@ -71392,6 +71435,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_vision": true, "supports_web_search": true }, @@ -74322,6 +74366,7 @@ "perplexity/google/gemini-3.8-flash": { "litellm_provider": "perplexity", "mode": "responses", + "supports_sampling_params": false, "input_cost_per_token": 7.5e-07, "output_cost_per_token": 3.75e-06, "cache_read_input_token_cost": 7.5e-08, @@ -75064,12 +75109,10 @@ "gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, "cache_read_input_token_cost_batches": 7.5e-08, - "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_image": 0.00168, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, - "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, @@ -75080,7 +75123,6 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, - "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", @@ -75113,12 +75155,10 @@ "vertex_ai/gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, "cache_read_input_token_cost_batches": 7.5e-08, - "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_image": 0.00168, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, - "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, @@ -75129,7 +75169,6 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, - "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", @@ -76929,5 +76968,55 @@ "supported_output_modalities": [ "text" ] + }, + "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash-1017-us": { + "cache_read_input_token_cost": 9e-09, + "cache_read_input_token_cost_priority": 1.125e-08, + "input_cost_per_token": 4.5e-07, + "input_cost_per_token_priority": 5.625e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.8e-06, + "output_cost_per_token_priority": 2.25e-06, + "source": "https://docs.fireworks.ai/serverless/us-only-serverless", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "wandb/reflection/Beam-1.0-501B-A23B": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "wandb/zai-org/GLM-5.3": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 1.2e-06, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false } } diff --git a/litellm/models/credentials.py b/litellm/models/credentials.py index d569ea12617..e6faf461104 100644 --- a/litellm/models/credentials.py +++ b/litellm/models/credentials.py @@ -10,6 +10,7 @@ from typing import Literal, TypeAlias from pydantic import ConfigDict, Field, model_validator +from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST from litellm.types.llms.base import LiteLLMBaseModel CredentialSource: TypeAlias = Literal["db", "config"] @@ -101,6 +102,7 @@ class UserConnectionFlowHandle(LiteLLMBaseModel): user_id: str credential_name: str + github_host: str = GITHUB_COPILOT_DEFAULT_HOST device_code: str = Field(repr=False) interval: int expires_at: float diff --git a/litellm/ocr/dispatch.py b/litellm/ocr/dispatch.py index 4fb4a6c709a..529c3df2bb9 100644 --- a/litellm/ocr/dispatch.py +++ b/litellm/ocr/dispatch.py @@ -9,7 +9,14 @@ from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR -from litellm.rust_bridge.public_call import NativeCall, native_call, native_call_hook, optional_str, signature +from litellm.rust_bridge.public_call import ( + AsyncFlag, + NativeCall, + native_call, + native_call_hook, + optional_str, + signature, +) __all__ = ("aocr", "ocr") @@ -39,6 +46,7 @@ def _bind_request( _OCR: Final = signature(_bind_request) +_FLAG: Final = AsyncFlag("aocr", sync=_OCR) def _public_request(name: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NativeCall: @@ -59,7 +67,6 @@ _DISPATCH: Final = PublicDispatch( route=Route.OCR, request=lambda args, kwargs: _public_request("ocr", args, kwargs), context=_context, - bypass=lambda request: request.kwargs.get("aocr") is True, ) _ADISPATCH: Final = PublicDispatch( @@ -73,9 +80,11 @@ def ocr( *args: object, **kwargs: object, # kwargs-ok: preserve the public OCR call shape ) -> OCRResponse | Coroutine[object, object, OCRResponse]: + if _FLAG.requested(kwargs): + return aocr(**_FLAG.keywords(args, kwargs)) return _DISPATCH.run( args, - kwargs, + _FLAG.dropped(kwargs), python=runtime.NO_PYTHON, binding=NATIVE_OCR, native=native_call_hook, @@ -85,7 +94,7 @@ def ocr( async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape return await _ADISPATCH.arun( args, - kwargs, + _FLAG.dropped(kwargs), python=runtime.NO_PYTHON, binding=NATIVE_AOCR, native=native_call_hook, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 744f9e1aa4c..4e94b1bb910 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -437,6 +437,7 @@ def llm_passthrough_route( params: QueryParamTypes | None = None, cookies: CookieTypes | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, + guardrailConfig: object = None, **kwargs, ) -> ( httpx.Response @@ -547,7 +548,7 @@ def llm_passthrough_route( _request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None) headers, signed_json_body = provider_config.sign_request( headers=headers, - litellm_params=litellm_params_dict, + litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig}, request_data=_request_data, api_base=str(updated_url), model=model, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index f65310a77e3..7eddf592262 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1197,6 +1197,8 @@ def should_strip_caller_authorization( admission credential — a virtual key, an IdP JWT, or an SSO / OIDC / session token whose ``api_key`` is ``None`` — and must never reach the upstream, so it is stripped regardless of the ``api_key`` value. + - **Every other auth-enabled mode**: strip when ``Authorization`` is the + credential admission validated. """ if _authorization_is_admission_jwt(raw_headers, user_api_key_auth): return True @@ -1216,7 +1218,21 @@ def should_strip_caller_authorization( mcp_server.auth_type == MCPAuth.oauth2 and mcp_server.delegate_auth_to_upstream ) if not (mcp_server.is_oauth_passthrough or is_delegated_oauth): - return False + from litellm.proxy.auth.user_api_key_auth import is_no_auth_dev_mode + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict + master_key, + ) + + proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings) + admission_ran: Final = not is_no_auth_dev_mode(master_key, proxy_general_settings) or ( + user_api_key_auth is not None and user_api_key_auth.authenticated_by_custom_auth + ) + return admission_ran and ( + _authorization_is_litellm_admission_credential_with_settings( + raw_headers, user_api_key_auth, proxy_general_settings + ) + ) has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers) if is_delegated_oauth: @@ -1243,16 +1259,39 @@ def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None return bool(_raw_header_value(raw_headers, "x-litellm-api-key")) +_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + def _authorization_is_litellm_admission_credential( raw_headers: Mapping[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, +) -> bool: + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict + ) + + proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings) + return _authorization_is_litellm_admission_credential_with_settings( + raw_headers, user_api_key_auth, proxy_general_settings + ) + + +def _authorization_is_litellm_admission_credential_with_settings( + raw_headers: Mapping[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + proxy_general_settings: Mapping[str, object], ) -> bool: """True when ``Authorization`` carries the LiteLLM key admission validated. - That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the - same key in both headers. + That is the case when no usable admission header was sent, or when the client repeated the same key in both headers. """ - admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key") + custom_header_name: Final[object | None] = proxy_general_settings.get("litellm_key_header_name") + custom_admission_header: Final = ( + _raw_header_value(raw_headers, custom_header_name.lower()) + if isinstance(custom_header_name, str) and custom_header_name + else None + ) + admission_header: Final = custom_admission_header or _raw_header_value(raw_headers, "x-litellm-api-key") authorization: Final = _raw_header_value(raw_headers, "authorization") if admission_header and authorization: return strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(admission_header, "Bearer") @@ -4322,6 +4361,7 @@ class MCPServerManager: raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, protocol_version_override: MCPUpstreamProtocol | None = None, + for_tool_listing: bool = False, ) -> MCPClient: self.catalog.assert_current(server) record_auth_resolution(server.server_id, AuthResolution.unresolved) @@ -4355,6 +4395,7 @@ class MCPServerManager: else None ), upstream_transport=self._upstream_transport, + for_tool_listing=for_tool_listing, ) _create_mcp_client = create_mcp_client @@ -4501,6 +4542,7 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, raw_headers=raw_headers, client_ip=client_ip, + for_tool_listing=True, ) ## HANDLE OPENAPI TOOLS diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index f217072bd39..37404ceece8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -460,6 +460,8 @@ def validate_static_credential( headers: Mapping[str, str], upstream_token_header: str | None = None, static_header_names: Iterable[str] = (), + *, + allow_absent: bool = False, ) -> Result[None, CredError]: if auth_type not in _STATIC_MODES: return Ok(None) @@ -475,17 +477,23 @@ def validate_static_credential( ) ) values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) + if allow_absent and not values: + return Ok(None) if any(_usable_credential_value(auth_type, name, value) for name, value in values): return Ok(None) return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential")) -async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: +async def prepare_mcp_client(server: MCPServer, client: MCPClient, *, for_tool_listing: bool = False) -> MCPClient: if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: return client request: Final = await client.prepare_request_auth() match validate_static_credential( - server.auth_type, request.headers, server.upstream_token_header, server.static_headers or () + server.auth_type, + request.headers, + server.upstream_token_header, + server.static_headers or (), + allow_absent=for_tool_listing and server.auth_type == MCPAuth.authorization, ): case Error(error): raise_public(error) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 3b7cf904b6a..1e39c068806 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -208,6 +208,7 @@ async def prepare_upstream_client( sampling_callback: SamplingCallback | None = None, elicitation_callback: ElicitationCallback | None = None, upstream_transport: httpx2.AsyncBaseTransport | None = None, + for_tool_listing: bool = False, ) -> MCPClient: transport: Final[MCPTransportType] = server.transport or MCPTransport.sse server_spec: Final = None if transport == MCPTransport.stdio else to_server_spec_fail_closed(server) @@ -272,6 +273,7 @@ async def prepare_upstream_client( sampling_callback=sampling_callback, elicitation_callback=elicitation_callback, ), + for_tool_listing=for_tool_listing, ) aws_auth: Final = ( MCPSigV4Auth( @@ -316,4 +318,5 @@ async def prepare_upstream_client( sampling_callback=sampling_callback, elicitation_callback=elicitation_callback, ), + for_tool_listing=for_tool_listing, ) diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index a7adb5c1bf5..0e2eed73364 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -275,7 +275,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( LazyFeature( name="decisions", module_path="litellm.proxy.decisions_endpoints.endpoints", - path_prefixes=("/v1/decisions", "/decisions", "/v1/systemone", "/systemone"), + path_prefixes=("/v1/decisions", "/decisions", "/v1/systemone", "/systemone", "/v1/extraction", "/extraction"), ), LazyFeature( name="claude_code_marketplace", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7198d74909c..545b317055a 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -9829,6 +9829,249 @@ "title": "DecisionsUsage", "type": "object" }, + "ExtractionAnswer": { + "additionalProperties": true, + "properties": { + "confidence": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Confidence" + }, + "guess": { + "$ref": "#/components/schemas/JsonValue" + }, + "mean": { + "$ref": "#/components/schemas/JsonValue" + }, + "median": { + "$ref": "#/components/schemas/JsonValue" + }, + "p10": { + "$ref": "#/components/schemas/JsonValue" + }, + "p90": { + "$ref": "#/components/schemas/JsonValue" + }, + "p_above": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "P Above" + }, + "p_below": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "P Below" + }, + "p_null": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "P Null" + }, + "quantiles": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/JsonValue" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Quantiles" + }, + "span": { + "anyOf": [ + { + "items": { + "type": "integer" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Span" + }, + "spans": { + "anyOf": [ + { + "items": { + "items": { + "type": "integer" + }, + "type": "array" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Spans" + }, + "steps": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/JsonValue" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Steps" + }, + "type": { + "title": "Type", + "type": "string" + }, + "value": { + "$ref": "#/components/schemas/JsonValue", + "description": "The answer, null when the input does not determine it" + } + }, + "required": [ + "type", + "value" + ], + "title": "ExtractionAnswer", + "type": "object" + }, + "ExtractionResponse": { + "additionalProperties": true, + "properties": { + "answers": { + "additionalProperties": { + "$ref": "#/components/schemas/ExtractionAnswer" + }, + "title": "Answers", + "type": "object" + }, + "model": { + "title": "Model", + "type": "string" + }, + "state_text": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "State Text" + }, + "truncated": { + "title": "Truncated", + "type": "boolean" + }, + "usage": { + "$ref": "#/components/schemas/ExtractionUsage" + } + }, + "required": [ + "model", + "answers", + "truncated", + "usage" + ], + "title": "ExtractionResponse", + "type": "object" + }, + "ExtractionUsage": { + "additionalProperties": true, + "properties": { + "image_tokens": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Image Tokens" + }, + "input_tokens": { + "title": "Input Tokens", + "type": "integer" + }, + "output_tokens": { + "title": "Output Tokens", + "type": "integer" + }, + "passes": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Passes" + }, + "schema_tokens": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Schema Tokens" + }, + "state_tokens": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "State Tokens" + } + }, + "required": [ + "input_tokens", + "output_tokens" + ], + "title": "ExtractionUsage", + "type": "object" + }, + "JsonValue": {}, "NoulAnswer": { "additionalProperties": true, "properties": { @@ -10765,6 +11008,403 @@ ] } }, + "/extraction": { + "post": { + "description": "Read a `state` text and up to 4 images and get one typed answer per named field in `fields`.\n\nField types are `text`, `number`, `date`, `time`, `datetime`, `duration`, `list` (set `items`), `estimate`\n(set `min` and `max`) and `forecast` (set `min`, `max` and `horizon`). Each answer carries its `value`, which is\nnull when the input does not determine it, plus `confidence`, `p_null`, text `span`s into the state and, for\nestimates, `mean`, `median`, `p10` and `p90`. Images are data URLs or `{content_type, base64}` objects, remote\nURLs are not accepted. An unknown field type, a field name that is not 1 to 64 letters, digits, `_`, `.` or `-`\nstarting with a letter or `_`, a remote image URL or more than 4 images is rejected with a 400 before any\nprovider is called. Any other top-level key, `stream: true` and OpenAI generation settings such as `temperature`\nincluded, is rejected with a 400 naming it unless `drop_params` or the deployment's `additional_drop_params`\ndrops it. Settings inside `fields` and `options` go to the provider as sent, so a field missing a setting its\ntype needs or an empty `state` with no images comes back with the provider's own error. `model` is any\nextraction model in the proxy model list, such as Cloudflare Clairic. The request and the answers keep the\nprovider's format.", + "operationId": "extraction_extraction_post", + "requestBody": { + "content": { + "application/json": { + "example": { + "fields": { + "days_to_pay": { + "max": 90, + "min": 0, + "type": "estimate" + }, + "invoice_date": { + "type": "date" + }, + "items": { + "instructions": "What was bought?", + "items": "text", + "type": "list" + }, + "total": { + "type": "number", + "unit": "USD" + }, + "vendor": { + "instructions": "Who sent the invoice?", + "type": "text" + } + }, + "model": "clairic", + "state": "Invoice #4471 from Acme Tools, dated 2026-03-03. 12 hammers. Total $174. Payment due in 30 days." + }, + "schema": { + "additionalProperties": true, + "properties": { + "fields": { + "description": "Field name to typed field spec. Each name becomes a key in the answers", + "minProperties": 1, + "patternProperties": { + "^[A-Za-z_][A-Za-z0-9_.\\-]{0,63}$": { + "additionalProperties": true, + "properties": { + "decimals": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Numbers, durations and number lists only, 0 to 6", + "title": "Decimals" + }, + "horizon": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Forecast only, number of steps", + "title": "Horizon" + }, + "instructions": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "What the field should contain", + "title": "Instructions" + }, + "integer": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Numbers and durations only", + "title": "Integer" + }, + "interval": { + "anyOf": [ + { + "enum": [ + "step", + "minute", + "hour", + "day", + "week", + "month", + "quarter", + "year" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Interval" + }, + "items": { + "anyOf": [ + { + "enum": [ + "text", + "number", + "date" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Item type, required for list fields", + "title": "Items" + }, + "max": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Estimate and forecast upper bound", + "title": "Max" + }, + "min": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Estimate and forecast lower bound", + "title": "Min" + }, + "null_threshold": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Probability of null above which the value is null", + "title": "Null Threshold" + }, + "scale": { + "anyOf": [ + { + "enum": [ + "linear", + "log" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Scale" + }, + "step": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Estimate grid step, or month or year", + "title": "Step" + }, + "type": { + "description": "The answer type for this field", + "enum": [ + "text", + "number", + "date", + "time", + "datetime", + "duration", + "list", + "estimate", + "forecast" + ], + "title": "Type", + "type": "string" + }, + "unit": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Unit label. For durations, seconds through years", + "title": "Unit" + } + }, + "required": [ + "type" + ], + "title": "ExtractionField", + "type": "object" + } + }, + "title": "Fields", + "type": "object" + }, + "images": { + "anyOf": [ + { + "items": { + "anyOf": [ + { + "pattern": "^[Dd][Aa][Tt][Aa]:", + "type": "string" + }, + { + "additionalProperties": true, + "properties": { + "base64": { + "description": "Base64-encoded image bytes", + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "The image media type, such as image/png", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "ExtractionBase64Image", + "type": "object" + } + ] + }, + "maxItems": 4, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "Up to 4 images as data URLs or content_type plus base64, no remote URLs", + "title": "Images" + }, + "model": { + "description": "An extraction model from the proxy model_list", + "title": "Model", + "type": "string" + }, + "options": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "include_quantiles": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Add 99 percentiles to estimates", + "title": "Include Quantiles" + }, + "return_state_text": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Return the text that spans index into", + "title": "Return State Text" + } + }, + "title": "ExtractionOptions", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "reference_date": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "ISO date that anchors dates written without a year", + "title": "Reference Date" + }, + "state": { + "default": "", + "description": "The text to read. Must be non-empty unless images are sent. Text answers are spans of it", + "title": "State", + "type": "string" + } + }, + "required": [ + "fields", + "model" + ], + "title": "ExtractionRequest", + "type": "object" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ExtractionResponse" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Extract typed fields from a state", + "tags": [ + "decisions" + ] + } + }, "/systemone": { "post": { "description": "Judge a `state` against 1 to 128 named `questions` and get one calibrated answer per question name.\n\nQuestion types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with\n`confidence` and `probabilities`), `score` (an index into the `criteria` levels, with `confidence`,\n`legend` and `probabilities`) and `box` (the image region the `instructions` describe, with `confidence`).\nTypeSafe models also take images in `assets`, each one referenced by its key in `state`. Requests that break\nthese rules are rejected with a 400 before any provider is called. `model` is any decision model in the proxy model list. OpenAI decision models accept this format\ntoo, LiteLLM translates it and keeps the answers keyed by question name. Streaming is not supported.\n\n[Docs](https://docs.litellm.ai/docs/decisions)", @@ -11775,6 +12415,403 @@ ] } }, + "/v1/extraction": { + "post": { + "description": "Read a `state` text and up to 4 images and get one typed answer per named field in `fields`.\n\nField types are `text`, `number`, `date`, `time`, `datetime`, `duration`, `list` (set `items`), `estimate`\n(set `min` and `max`) and `forecast` (set `min`, `max` and `horizon`). Each answer carries its `value`, which is\nnull when the input does not determine it, plus `confidence`, `p_null`, text `span`s into the state and, for\nestimates, `mean`, `median`, `p10` and `p90`. Images are data URLs or `{content_type, base64}` objects, remote\nURLs are not accepted. An unknown field type, a field name that is not 1 to 64 letters, digits, `_`, `.` or `-`\nstarting with a letter or `_`, a remote image URL or more than 4 images is rejected with a 400 before any\nprovider is called. Any other top-level key, `stream: true` and OpenAI generation settings such as `temperature`\nincluded, is rejected with a 400 naming it unless `drop_params` or the deployment's `additional_drop_params`\ndrops it. Settings inside `fields` and `options` go to the provider as sent, so a field missing a setting its\ntype needs or an empty `state` with no images comes back with the provider's own error. `model` is any\nextraction model in the proxy model list, such as Cloudflare Clairic. The request and the answers keep the\nprovider's format.", + "operationId": "extraction_v1_extraction_post", + "requestBody": { + "content": { + "application/json": { + "example": { + "fields": { + "days_to_pay": { + "max": 90, + "min": 0, + "type": "estimate" + }, + "invoice_date": { + "type": "date" + }, + "items": { + "instructions": "What was bought?", + "items": "text", + "type": "list" + }, + "total": { + "type": "number", + "unit": "USD" + }, + "vendor": { + "instructions": "Who sent the invoice?", + "type": "text" + } + }, + "model": "clairic", + "state": "Invoice #4471 from Acme Tools, dated 2026-03-03. 12 hammers. Total $174. Payment due in 30 days." + }, + "schema": { + "additionalProperties": true, + "properties": { + "fields": { + "description": "Field name to typed field spec. Each name becomes a key in the answers", + "minProperties": 1, + "patternProperties": { + "^[A-Za-z_][A-Za-z0-9_.\\-]{0,63}$": { + "additionalProperties": true, + "properties": { + "decimals": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Numbers, durations and number lists only, 0 to 6", + "title": "Decimals" + }, + "horizon": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Forecast only, number of steps", + "title": "Horizon" + }, + "instructions": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "What the field should contain", + "title": "Instructions" + }, + "integer": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Numbers and durations only", + "title": "Integer" + }, + "interval": { + "anyOf": [ + { + "enum": [ + "step", + "minute", + "hour", + "day", + "week", + "month", + "quarter", + "year" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Interval" + }, + "items": { + "anyOf": [ + { + "enum": [ + "text", + "number", + "date" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Item type, required for list fields", + "title": "Items" + }, + "max": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Estimate and forecast upper bound", + "title": "Max" + }, + "min": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Estimate and forecast lower bound", + "title": "Min" + }, + "null_threshold": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Probability of null above which the value is null", + "title": "Null Threshold" + }, + "scale": { + "anyOf": [ + { + "enum": [ + "linear", + "log" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Scale" + }, + "step": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Estimate grid step, or month or year", + "title": "Step" + }, + "type": { + "description": "The answer type for this field", + "enum": [ + "text", + "number", + "date", + "time", + "datetime", + "duration", + "list", + "estimate", + "forecast" + ], + "title": "Type", + "type": "string" + }, + "unit": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Unit label. For durations, seconds through years", + "title": "Unit" + } + }, + "required": [ + "type" + ], + "title": "ExtractionField", + "type": "object" + } + }, + "title": "Fields", + "type": "object" + }, + "images": { + "anyOf": [ + { + "items": { + "anyOf": [ + { + "pattern": "^[Dd][Aa][Tt][Aa]:", + "type": "string" + }, + { + "additionalProperties": true, + "properties": { + "base64": { + "description": "Base64-encoded image bytes", + "title": "Base64", + "type": "string" + }, + "content_type": { + "description": "The image media type, such as image/png", + "title": "Content Type", + "type": "string" + } + }, + "required": [ + "content_type", + "base64" + ], + "title": "ExtractionBase64Image", + "type": "object" + } + ] + }, + "maxItems": 4, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "Up to 4 images as data URLs or content_type plus base64, no remote URLs", + "title": "Images" + }, + "model": { + "description": "An extraction model from the proxy model_list", + "title": "Model", + "type": "string" + }, + "options": { + "anyOf": [ + { + "additionalProperties": true, + "properties": { + "include_quantiles": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Add 99 percentiles to estimates", + "title": "Include Quantiles" + }, + "return_state_text": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Return the text that spans index into", + "title": "Return State Text" + } + }, + "title": "ExtractionOptions", + "type": "object" + }, + { + "type": "null" + } + ] + }, + "reference_date": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "ISO date that anchors dates written without a year", + "title": "Reference Date" + }, + "state": { + "default": "", + "description": "The text to read. Must be non-empty unless images are sent. Text answers are spans of it", + "title": "State", + "type": "string" + } + }, + "required": [ + "fields", + "model" + ], + "title": "ExtractionRequest", + "type": "object" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ExtractionResponse" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Extract typed fields from a state", + "tags": [ + "decisions" + ] + } + }, "/v1/systemone": { "post": { "description": "Judge a `state` against 1 to 128 named `questions` and get one calibrated answer per question name.\n\nQuestion types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with\n`confidence` and `probabilities`), `score` (an index into the `criteria` levels, with `confidence`,\n`legend` and `probabilities`) and `box` (the image region the `instructions` describe, with `confidence`).\nTypeSafe models also take images in `assets`, each one referenced by its key in `state`. Requests that break\nthese rules are rejected with a 400 before any provider is called. `model` is any decision model in the proxy model list. OpenAI decision models accept this format\ntoo, LiteLLM translates it and keeps the answers keyed by question name. Streaming is not supported.\n\n[Docs](https://docs.litellm.ai/docs/decisions)", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 315533d6d0f..54acf59ef80 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -377,6 +377,8 @@ class LiteLLMRoutes(enum.Enum): "/v1/decisions", "/systemone", "/v1/systemone", + "/extraction", + "/v1/extraction", # OCR "/ocr", "/v1/ocr", @@ -3729,6 +3731,7 @@ class SpendLogsMetadata(TypedDict): routing_decision: StandardLoggingRoutingDecision | None kubernetes_pod_routing: ReadOnly[StandardLoggingKubernetesPodRouting | None] internal_call_origin: InternalCallOrigin | None + billed_retry_attempt: ReadOnly[NotRequired[bool | None]] litellm_roi_estimator: ReadOnly[NotRequired[bool | None]] guardrail_information: list[StandardLoggingGuardrailInformation] | None eval_information: Any | None diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py index cbbd1548301..cdbbd6029fa 100644 --- a/litellm/proxy/agent_endpoints/auth/managed_authorization.py +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -32,6 +32,7 @@ _MANAGED_MODEL_ROUTES: Final = frozenset( "rerank", "decisions", "systemone", + "extraction", "ocr", ), ) @@ -77,6 +78,7 @@ _MODEL_ROUTE_KINDS: Final[ "/rerank": "body", "/decisions": "body", "/systemone": "body", + "/extraction": "body", "/messages/count_tokens": "body", ":countTokens": "path", } diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 996c28bebaa..7a186d84aa5 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -446,6 +446,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # the request away from the admin's pinned configuration. "nvcf_function_id", "use_ssl", + "riva_offline", # Per-deployment opt-in that hands the whole call to the Rust core. It is a # deployment decision, not a request one: the Rust path uses its own client # rather than the one the deployment configured, and reports no post_call, @@ -462,6 +463,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # stored credential's values and the caller's connection by the proxy's own # secret_fields, so a body-supplied value could only spoof either. "github_copilot_auth_type", + "github_copilot_host", "user_provider_credentials", "github_copilot_user_session", # Observability credentials, hosts, and project identifiers: derived diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 842705ad5a7..4038ff6f582 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -494,7 +494,7 @@ Each command resolves your LiteLLM key (logging in via SSO when none is stored a The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. Codex gets the same list as a catalog file, `$CODEX_HOME/litellm-models.json` (default `~/.codex/`), passed as `-c model_catalog_json=` so `/model` lists exactly the proxy's chat models; a proxy model the installed Codex already knows (`gpt-5.5`, say) keeps that Codex's own entry, reasoning levels and prompt included, and only its place in the picker comes from the proxy, while a model Codex does not know gets the plain entry Codex uses for an unknown `-m` slug. Before launching, `lite codex` asks the installed Codex for its own list and then has it read the written file back (both through `codex debug models`), and when the fetch, either of those or the write fails (Codex releases older than 0.130 have no such command) it says so on stderr and launches with Codex's built-in catalog, leaving a rejected file in place. -pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. +pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. The entry also turns on pi's session header (`x-session-id`, sent by pi 0.80.7 and later), so every chat turn of one pi session shares one session id in the logs and sticks to one deployment when session affinity is on. A custom `pi --session-id` must be at least 8 letters, digits, `_`, or `-` for the proxy to read it. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. Options (these belong to the wrapper, so put them before the agent's own flags): diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index b228617699c..357b82bdf67 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -191,7 +191,7 @@ def provider_block( Real contextWindow/maxTokens matter: pi otherwise assumes 128k/16384, which breaks compaction thresholds and over-asks models with smaller output caps. pi sniffs compat from the base URL, and one gateway URL fronts models with - different capabilities, so both flags are pinned off. + different capabilities, so the store and long-cache flags are pinned off. """ return { "baseUrl": base_url.rstrip("/") + "/v1", @@ -199,6 +199,8 @@ def provider_block( "compat": { "supportsStore": False, "supportsLongCacheRetention": False, + "sendSessionAffinityHeaders": True, + "sessionAffinityFormat": "openrouter", }, "apiKey": f"${LITELLM_PROXY_API_KEY_ENV}", "models": [_model_entry(model_id, limits) for model_id in model_ids], diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6f8ef830ee6..02649c52644 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -68,6 +68,7 @@ from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.litellm_core_utils.llm_response_utils.get_headers import ( @@ -134,12 +135,14 @@ 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 from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model +from litellm.types.extraction import ExtractionResponse from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError from litellm.types.router_weights import validate_router_weights _LateResponseT = TypeVar("_LateResponseT", bound=Response) _LlmCallT = TypeVar("_LlmCallT") +_ResponseT: Final = TypeVar("_ResponseT") _ChunkT: Final = TypeVar("_ChunkT") KNOWN_PROXY_ROUTES: Final = frozenset( @@ -191,6 +194,7 @@ ProxyRouteType: TypeAlias = Literal[ "asearch", "adecisions", "asystemone", + "aextraction", "avideo_generation", "avideo_list", "avideo_status", @@ -263,6 +267,7 @@ from litellm.proxy.litellm_pre_call_utils import ( sanitize_for_log, ) from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -549,22 +554,29 @@ def model_deployments_support_stream_options( return all(_litellm_model_supports_stream_options(m) for m in candidate_models) +_STREAM_USAGE_FLAGS: Final = ("_litellm_strip_stream_usage", STREAM_USAGE_INJECTED_FIELD) + + def stream_usage_tracking_updates( data: Mapping[str, object], general_settings: Mapping[str, object], route_type: str, supports_stream_options: Callable[[], bool], ) -> Mapping[str, object]: - scrub: Final = {"_litellm_strip_stream_usage": False} if "_litellm_strip_stream_usage" in data else {} + scrub: Final = {flag: False for flag in _STREAM_USAGE_FLAGS if flag in data} if data.get("stream", False) is not True: return scrub always_include: Final = general_settings.get("always_include_stream_usage") stream_options: Final = data.get("stream_options") if always_include is True: if "stream_options" not in data: - return {**scrub, "stream_options": {"include_usage": True}} + return {**scrub, "stream_options": {"include_usage": True}, STREAM_USAGE_INJECTED_FIELD: True} if isinstance(stream_options, dict) and "include_usage" not in stream_options: - return {**scrub, "stream_options": {**stream_options, "include_usage": True}} + return { + **scrub, + "stream_options": {**stream_options, "include_usage": True}, + STREAM_USAGE_INJECTED_FIELD: True, + } return scrub if always_include is False and route_type == "agenerate_content_stream" and "stream_options" not in data: return {**scrub, "stream_options": {"include_usage": False}} @@ -578,6 +590,7 @@ def stream_usage_tracking_updates( return { "stream_options": {**merged_stream_options, "include_usage": True}, "_litellm_strip_stream_usage": True, + STREAM_USAGE_INJECTED_FIELD: True, } @@ -1482,11 +1495,11 @@ _is_azure_model_router_request: Final = is_azure_model_router_request def _override_openai_response_model( *, - response_obj: object, + response_obj: _ResponseT, requested_model: str, log_context: str, return_raw_model_name: bool = False, -) -> None: +) -> _ResponseT: """ Force the OpenAI-compatible `model` field in the response to match what the client requested. @@ -1510,7 +1523,7 @@ def _override_openai_response_model( model group name instead of the comma-separated list the client sent. """ if return_raw_model_name or not requested_model: - return + return response_obj hidden_params: Final = get_hidden_params_dict(response_obj) if isinstance(hidden_params, dict): @@ -1523,7 +1536,7 @@ def _override_openai_response_model( log_context, attempted_fallbacks, ) - return + return response_obj # For fastest_response batch completions, use the winning model's group # name rather than the comma-separated list the client sent. @@ -1542,7 +1555,7 @@ def _override_openai_response_model( "%s: fastest_response detected but no model group header found, preserving actual model from response.", log_context, ) - return + return response_obj # Check if this is an Azure Model Router request - if so, preserve the actual model used if is_azure_model_router_request(requested_model, hidden_params): @@ -1550,11 +1563,12 @@ def _override_openai_response_model( "%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.", log_context, ) - return + return response_obj + unchanged_response: Final = response_obj if isinstance(response_obj, dict): if "model" not in response_obj: - return + return unchanged_response downstream_model = response_obj.get("model") if downstream_model != requested_model: verbose_proxy_logger.debug( @@ -1564,10 +1578,10 @@ def _override_openai_response_model( downstream_model, ) response_obj["model"] = requested_model - return + return unchanged_response if not hasattr(response_obj, "model"): - return + return response_obj downstream_model = getattr(response_obj, "model", None) if downstream_model != requested_model: @@ -1578,6 +1592,9 @@ def _override_openai_response_model( downstream_model, ) + if isinstance(response_obj, ExtractionResponse): + return response_obj.model_copy(update={"model": requested_model}) + try: setattr(response_obj, "model", requested_model) except Exception as e: @@ -1589,6 +1606,7 @@ def _override_openai_response_model( str(e), exc_info=True, ) + return response_obj _METADATA_BUCKET_KEYS: Final = ("metadata", "litellm_metadata") @@ -2244,6 +2262,7 @@ class ProxyBaseLLMRequestProcessing: "asearch", "adecisions", "asystemone", + "aextraction", "avideo_generation", "avideo_list", "avideo_status", @@ -3270,7 +3289,7 @@ class ProxyBaseLLMRequestProcessing: # Always return the client-requested model name (not provider-prefixed internal identifiers) # for OpenAI-compatible responses. if requested_model_from_client: - _override_openai_response_model( + response = _override_openai_response_model( # rebind-ok: a frozen response comes back as a copy response_obj=response, requested_model=requested_model_from_client, log_context=f"litellm_call_id={logging_obj.litellm_call_id}", @@ -4024,6 +4043,8 @@ class ProxyBaseLLMRequestProcessing: async def _on_deferred_native_stream_complete(*args: object) -> None: match args: + case (BilledRoundsDispatch(dispatch=logging_coroutine),): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (assembled_response, cache_hit): diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 953c944bfb2..d20b5602186 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -19,6 +19,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.constants import ( GITHUB_COPILOT_AUTH_TYPE_KEY, + GITHUB_COPILOT_HOST_KEY, GITHUB_COPILOT_PER_USER_AUTH_TYPE, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor @@ -84,6 +85,22 @@ _DISPLAY_NAME_MAX_LENGTH: Final = 255 _DISPLAY_NAME_FORBIDDEN_CATEGORIES: Final = frozenset({"Cc", "Cf", "Cs", "Zl", "Zp"}) +def _normalized_github_copilot_values( + credential_values: Mapping[str, object], + credential_info: Mapping[str, object], +) -> Mapping[str, object]: + provider: Final = stored_credential_provider(credential_info.get("custom_llm_provider")) + if provider != _GITHUB_COPILOT_PROVIDER or GITHUB_COPILOT_HOST_KEY not in credential_values: + return credential_values + + from litellm.llms.github_copilot.github_host import normalized_github_host + + host: Final = normalized_github_host(credential_values[GITHUB_COPILOT_HOST_KEY]) + if host is None: + raise HTTPException(status_code=400, detail=f"Invalid {GITHUB_COPILOT_HOST_KEY}") + return {**credential_values, GITHUB_COPILOT_HOST_KEY: host} + + def _reject_non_admin_wif_fields( wif_fields: tuple[str, ...], user_api_key_dict: UserAPIKeyAuth, @@ -318,7 +335,14 @@ async def create_credential( status_code=400, detail="Credential values are required. Unable to infer credential values from model ID.", ) - _reject_non_admin_wif_fields(server_owned_wif_fields_written(credential_values), user_api_key_dict) + credential_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field + ) + normalized_values: Final = _normalized_github_copilot_values( + _CREDENTIAL_DICT_ADAPTER.validate_python(credential_values), + credential_info, + ) + _reject_non_admin_wif_fields(server_owned_wif_fields_written(normalized_values), user_api_key_dict) _reject_non_admin_wif_fields( await named_credential_wif_fields(credential.credential_name, prisma_client), user_api_key_dict ) @@ -326,8 +350,8 @@ async def create_credential( processed_credential: Final = CredentialItem( credential_name=credential.credential_name, display_name=_normalized_display_name(credential.display_name), - credential_values=_without_null_values(_CREDENTIAL_DICT_ADAPTER.validate_python(credential_values)), - credential_info=credential.credential_info, + credential_values=_without_null_values(normalized_values), + credential_info=credential_info, ) encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential) credentials_dict: Final = encrypted_credential.model_dump(exclude_none=True) @@ -614,22 +638,24 @@ async def start_user_connection( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI resolves the dependency from the default ) -> UserConnectionStartResponse: """Begin a GitHub device flow for the calling user's connection to a per-user credential.""" - from litellm.llms.github_copilot.per_user_auth import astart_device_flow + from litellm.llms.github_copilot.per_user_auth import astart_device_flow, github_copilot_host_endpoints from litellm.proxy.proxy_server import prisma_client try: user_id: Final = _authenticated_user_id(user_api_key_dict) _per_user_credential_or_404(credential_name) + endpoints: Final = github_copilot_host_endpoints(credential_name) if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - flow: Final = await astart_device_flow() + flow: Final = await astart_device_flow(endpoints) handle: Final = encrypt_value_helper( UserConnectionFlowHandle( user_id=user_id, credential_name=credential_name, + github_host=endpoints.host, device_code=flow.device_code, interval=flow.interval, expires_at=time.time() + flow.expires_in, @@ -665,12 +691,14 @@ async def poll_user_connection( acheck_copilot_seat, afetch_github_login, apoll_device_flow, + github_copilot_host_endpoints, ) from litellm.proxy.proxy_server import prisma_client, user_api_key_cache try: user_id: Final = _authenticated_user_id(user_api_key_dict) _per_user_credential_or_404(credential_name) + endpoints: Final = github_copilot_host_endpoints(credential_name) if prisma_client is None: raise HTTPException( status_code=500, @@ -681,11 +709,12 @@ async def poll_user_connection( handle is None or handle.user_id != user_id or handle.credential_name != credential_name + or handle.github_host != endpoints.host or handle.expires_at <= time.time() ): raise HTTPException(status_code=400, detail="invalid or expired flow_handle") - poll: Final = await apoll_device_flow(handle.device_code) + poll: Final = await apoll_device_flow(handle.device_code, endpoints) if poll.status in ("expired", "denied", "pending", "slow_down"): return UserConnectionPollResponse(status=poll.status, interval=poll.interval) @@ -693,12 +722,19 @@ async def poll_user_connection( if not github_token: raise HTTPException(status_code=502, detail="GitHub device flow completed without a token") try: - await acheck_copilot_seat(user_id=user_id, github_token=github_token, credential_name=credential_name) + await acheck_copilot_seat( + user_id=user_id, + github_token=github_token, + credential_name=credential_name, + endpoints=endpoints, + ) except litellm.CallerCredentialAuthenticationError: return UserConnectionPollResponse(status="no_copilot_seat") - github_login: Final = await afetch_github_login(github_token) + github_login: Final = await afetch_github_login(github_token, endpoints) connection_payload: Final = GithubCopilotUserConnectionPayload( - access_token=github_token, github_login=github_login + access_token=github_token, + github_login=github_login, + github_host=handle.github_host, ) await upsert_user_provider_credential( prisma_client, @@ -760,7 +796,7 @@ async def delete_user_connection( ) prior: Final = await delete_user_provider_credential(prisma_client, user_id, credential_name) if prior is not None: - evict_copilot_user_session(user_id, prior.access_token) + evict_copilot_user_session(user_id, credential_name, prior.github_host) return UserConnectionDeleteResponse(status="disconnected") except Exception as e: # noqa: BLE001 # endpoint boundary: every failure becomes the proxy error contract raise handle_exception_on_proxy(e) @@ -785,16 +821,15 @@ async def _purge_user_connections_for_credential(credential_name: str) -> None: "_purge_user_connections_for_credential: failed to drop user connections for %s", credential_name ) return - access_tokens: Final[dict[str, str]] = {} # mutable-ok: accumulates one entry per deleted row - for row in prior_rows: - decoded = decode_user_provider_credential(row.credential_b64) - if decoded is not None: - access_tokens[row.user_id] = decoded.access_token + github_hosts: Final = { + row.user_id: decoded.github_host + for row in prior_rows + if (decoded := decode_user_provider_credential(row.credential_b64)) is not None + } for user_id in user_ids: await invalidate_user_provider_credential_cache(user_api_key_cache, user_id, credential_name) - token = access_tokens.get(user_id) - if token: - evict_copilot_user_session(user_id, token) + if user_id in github_hosts: + evict_copilot_user_session(user_id, credential_name, github_hosts[user_id]) @router.delete( @@ -926,7 +961,25 @@ async def update_credential( credential_values=incoming_values, credential_values_to_delete=credential.credential_values_to_delete, ) - merged_credential: Final = update_db_credential(db_credential, patch) + decrypted_stored_credential: Final = decrypted_credential(db_credential) + decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + decrypted_stored_credential.credential_values + ) + stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field + ) + patch_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + patch.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field + ) + effective_info: Final = {**stored_info, **patch_info} + normalized_patch: Final = ( + patch.model_copy( + update={"credential_values": _normalized_github_copilot_values(incoming_values, effective_info)} + ) + if incoming_values.get(GITHUB_COPILOT_HOST_KEY) is not None + else patch + ) + merged_credential: Final = update_db_credential(db_credential, normalized_patch) credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str "dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True)) ) @@ -940,13 +993,20 @@ async def update_credential( ) # Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted) - _sync_in_memory_credential(patch, credential_name) + _sync_in_memory_credential(normalized_patch, credential_name) - merged_values: Final = cast( # cast-ok: credential_values is a plain dict at runtime - Mapping[object, object], merged_credential.credential_values + from litellm.llms.github_copilot.github_host import normalized_github_host + + new_decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + decrypted_credential(merged_credential).credential_values ) - still_per_user: Final = merged_values.get(GITHUB_COPILOT_AUTH_TYPE_KEY) == GITHUB_COPILOT_PER_USER_AUTH_TYPE - if not still_per_user: + still_per_user: Final = ( + new_decrypted_values.get(GITHUB_COPILOT_AUTH_TYPE_KEY) == GITHUB_COPILOT_PER_USER_AUTH_TYPE + ) + old_host: Final = normalized_github_host(decrypted_values.get(GITHUB_COPILOT_HOST_KEY)) + new_host: Final = normalized_github_host(new_decrypted_values.get(GITHUB_COPILOT_HOST_KEY)) + host_changed: Final = old_host is None or old_host != new_host + if not still_per_user or host_changed: await _purge_user_connections_for_credential(credential_name) return {"success": True, "message": "Credential updated successfully"} diff --git a/litellm/proxy/credential_endpoints/user_provider_credentials.py b/litellm/proxy/credential_endpoints/user_provider_credentials.py index 4253483a228..22ee271410e 100644 --- a/litellm/proxy/credential_endpoints/user_provider_credentials.py +++ b/litellm/proxy/credential_endpoints/user_provider_credentials.py @@ -25,6 +25,7 @@ from litellm.caching.dual_cache import DualCache if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache from litellm.constants import ( + GITHUB_COPILOT_DEFAULT_HOST, GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS, USER_PROVIDER_CREDENTIAL_CACHE_PREFIX, USER_PROVIDER_CREDENTIAL_NOT_CONNECTED, @@ -50,6 +51,7 @@ class GithubCopilotUserConnectionPayload(BaseModel): access_token: str = Field(repr=False) github_login: str + github_host: str = GITHUB_COPILOT_DEFAULT_HOST class _UserProviderCredentialsRepository(PrismaTableRepository["prisma_models.LiteLLM_UserProviderCredentials"]): @@ -186,7 +188,7 @@ async def _try_cache_get(token_cache: "RedisCache", key: str) -> object: try: return await _string_cache(token_cache).async_get_cache(key) except Exception: # noqa: BLE001 # a Redis outage must fall back to the database, never reject a connected user - verbose_proxy_logger.warning("aget_user_provider_tokens: Redis get failed; falling back to the database") + verbose_proxy_logger.warning("aget_user_provider_credentials: Redis get failed; falling back to the database") return None @@ -196,7 +198,7 @@ async def _try_cache_set(token_cache: "RedisCache", key: str, value: str, nx: bo key, value, nx=nx, ttl=GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS ) except Exception: # noqa: BLE001 # caching is best-effort; a Redis outage is not worth failing the request - verbose_proxy_logger.warning("aget_user_provider_tokens: Redis set failed; skipping the cache write") + verbose_proxy_logger.warning("aget_user_provider_credentials: Redis set failed; skipping the cache write") @with_service_target("user_provider_connections") @@ -292,16 +294,16 @@ async def set_user_provider_credential_cache( @with_service_target("user_provider_connections") -async def aget_user_provider_tokens( +async def aget_user_provider_credentials( prisma_client: "PrismaClient", cache: DualCache, user_id: str, credential_names: Sequence[str], -) -> Mapping[str, str]: - """Map each per-user credential name to the caller's stored GitHub token. +) -> Mapping[str, Mapping[str, str]]: + """Map each per-user credential name to the caller's stored GitHub connection. Cached values are the stored ciphertext or the ``_NOT_CONNECTED`` marker; - plaintext tokens only ever live in the returned dict. Redis is the only + plaintext connection data only ever lives in the returned mapping. Redis is the only cache layer used: without it there is no shared invalidation, so a worker serving a stale local entry after a disconnect is worse than a DB read.""" token_cache: Final = cache.redis_cache @@ -320,9 +322,9 @@ async def aget_user_provider_tokens( misses.append(name) if not misses: return { - name: token + name: connection for name, ciphertext in cached.items() - if (token := _token_from_ciphertext(ciphertext)) is not None + if (connection := _connection_from_ciphertext(ciphertext)) is not None } try: @@ -333,7 +335,7 @@ async def aget_user_provider_tokens( where={"user_id": user_id}, ) except Exception: # noqa: BLE001 # a DB outage reads as not connected so the caller gets a 401, never a 500 - verbose_proxy_logger.exception("aget_user_provider_tokens: DB read failed for user_id=%s", user_id) + verbose_proxy_logger.exception("aget_user_provider_credentials: DB read failed") return {} found: Final = {row.credential_name: row.credential_b64 for row in rows} for name in misses: @@ -342,10 +344,14 @@ async def aget_user_provider_tokens( if name in found: cached[name] = found[name] return { - name: token for name, ciphertext in cached.items() if (token := _token_from_ciphertext(ciphertext)) is not None + name: connection + for name, ciphertext in cached.items() + if (connection := _connection_from_ciphertext(ciphertext)) is not None } -def _token_from_ciphertext(ciphertext: str) -> str | None: +def _connection_from_ciphertext(ciphertext: str) -> Mapping[str, str] | None: payload: Final = decode_user_provider_credential(ciphertext) - return payload.access_token if payload is not None else None + if payload is None: + return None + return {"access_token": payload.access_token, "github_host": payload.github_host} diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e7909cbe600..4e98823ba77 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -29,6 +29,7 @@ from litellm._internal_context import service_target, with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( + BILLED_RETRY_ATTEMPT_METADATA_KEY, DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -312,6 +313,11 @@ def _spend_tables_left_to_send( return cast(DBSpendUpdateTransactions, remaining) if any(remaining.values()) else None +def _is_billed_retry_attempt(metadata: SpendLogsMetadata) -> bool: + """An earlier billed retry attempt's own spend row, not a request the caller sent.""" + return BILLED_RETRY_ATTEMPT_METADATA_KEY in metadata and metadata[BILLED_RETRY_ATTEMPT_METADATA_KEY] is True + + def _timed_request_duration_ms( payload: dict | SpendLogsPayload, request_status: Literal["success", "failure"], @@ -1651,6 +1657,8 @@ class DBSpendUpdateWriter: window_spend_update_queue=self.window_spend_update_queue, ) + await self._flush_tool_discovery_queue(prisma_client=prisma_client) + # Only commit from redis to db if this pod is the leader if await self.pod_lock_manager.acquire_lock( cronjob_id=DB_SPEND_UPDATE_JOB_NAME, @@ -2803,11 +2811,12 @@ class DBSpendUpdateWriter: if not date: return metadata: Final[SpendLogsMetadata] = json.loads(payload["metadata"]) + internal_origin: Final = bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults request_error_accumulator.record( payload=payload, request_status=prisma_client.get_request_status(payload), # pyright: ignore[reportUnknownMemberType] # legacy payload union; this caller supplies a typed spend payload date=date.split("T")[0], - is_internal_call=bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)), # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults + is_internal_call=internal_origin or _is_billed_retry_attempt(metadata), ) async def _common_add_spend_log_transaction_to_daily_transaction( @@ -2872,6 +2881,7 @@ class DBSpendUpdateWriter: endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None) is_internal_call: Final = bool(_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) + caller_request: Final = not is_internal_call and not _is_billed_retry_attempt(_metadata) cache_read_input_tokens: Final = extract_cache_read_tokens(usage_obj) compression_saved_tokens: Final = extract_compression_saved_tokens(_metadata) savings_spend: Final = compute_savings_spend( @@ -2906,9 +2916,9 @@ class DBSpendUpdateWriter: # requests the caller made: counting them inflates request-volume # readers, and an auto-router savings figure computed on a shadow # duplicate credits savings for traffic no user sent. - api_requests=0 if is_internal_call else 1, - successful_requests=1 if not is_internal_call and request_status == "success" else 0, - failed_requests=1 if not is_internal_call and request_status != "success" else 0, + api_requests=1 if caller_request else 0, + successful_requests=1 if caller_request and request_status == "success" else 0, + failed_requests=1 if caller_request and request_status != "success" else 0, cache_read_input_tokens=cache_read_input_tokens, cache_creation_input_tokens=extract_cache_creation_tokens(usage_obj), compression_saved_tokens=compression_saved_tokens, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 898aee493df..54d20bd4507 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -168,9 +168,11 @@ class RedisUpdateBuffer: @with_service_target(SPEND_QUEUE_TARGET) async def _store_transactions_in_redis( self, - transactions: Mapping[str, BaseDailySpendTransaction] | None, + transactions: dict[str, BaseDailySpendTransaction] | None, # mutable-ok: add_update requires dict redis_key: str, service_type: ServiceTypes, + *, + source_queue: DailySpendUpdateQueue, ) -> None: """ Helper method to store transactions in Redis and emit an event @@ -179,6 +181,7 @@ class RedisUpdateBuffer: transactions: The transactions to store redis_key: The Redis key to store under service_type: The service type for event emission + source_queue: The in-memory queue to restore transactions to if Redis push fails """ if transactions is None or len(transactions) == 0: return @@ -191,6 +194,15 @@ class RedisUpdateBuffer: key=redis_key, values=list_of_transactions, ) + except Exception as e: + verbose_proxy_logger.error( + "Spend tracking - failed to push spend updates to Redis (redis_key=%s). " + "Restoring them to the in-memory queue for retry on next tick. Error: %s", + redis_key, + str(e), + ) + await source_queue.add_update(transactions) + else: verbose_proxy_logger.debug( "Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s", redis_key, @@ -200,12 +212,6 @@ class RedisUpdateBuffer: queue_size=current_redis_buffer_size, service=service_type, ) - except Exception as e: - verbose_proxy_logger.error( - "Spend tracking - failed to push spend updates to Redis (redis_key=%s). Error: %s", - redis_key, - str(e), - ) @with_service_target(SPEND_QUEUE_TARGET) async def store_in_memory_spend_updates_in_redis( @@ -793,6 +799,7 @@ class RedisUpdateBuffer: transactions=daily_tag_spend_update_transactions, redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, + source_queue=daily_tag_spend_update_queue, ) @with_service_target(SPEND_QUEUE_TARGET) diff --git a/litellm/proxy/db/health_check_latest.py b/litellm/proxy/db/health_check_latest.py index 7672ed89e9a..024be7603d6 100644 --- a/litellm/proxy/db/health_check_latest.py +++ b/litellm/proxy/db/health_check_latest.py @@ -10,8 +10,9 @@ worker on every call. ``SELECT DISTINCT ON`` keeps the transfer at one row per from __future__ import annotations import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol from pydantic import ConfigDict, JsonValue, TypeAdapter, field_validator @@ -113,3 +114,13 @@ async def fetch_latest_health_checks_for_models( except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err) return () + + +def latest_by_model_name(rows: Iterable[LatestHealthCheckRow]) -> Mapping[str, LatestHealthCheckRow]: + """The newest row per model name. + + The queries above return one row per (model_id, model_name), so a model group with several + deployments comes back as several rows, plus one without a model_id if it was ever checked by + name. They arrive in model_id order with NULL last, which says nothing about which check is newest. + """ + return MappingProxyType({row.model_name: row for row in sorted(rows, key=lambda row: row.checked_at)}) diff --git a/litellm/proxy/decisions_endpoints/endpoints.py b/litellm/proxy/decisions_endpoints/endpoints.py index 1477e9f9591..366842d23a6 100644 --- a/litellm/proxy/decisions_endpoints/endpoints.py +++ b/litellm/proxy/decisions_endpoints/endpoints.py @@ -1,6 +1,6 @@ from collections.abc import Mapping from types import MappingProxyType -from typing import Annotated, Final, Literal +from typing import Annotated, Final, Literal, TypeAlias from fastapi import APIRouter, Depends, Request, Response from fastapi.responses import ORJSONResponse # pyright: ignore[reportDeprecated] # required endpoint contract @@ -18,6 +18,7 @@ from litellm.types.decisions import ( OpenAIDecisionRequestBody, OpenAIDecisionResponse, ) +from litellm.types.extraction import ExtractionRequest, ExtractionRequestBody, ExtractionResponse router: Final = APIRouter() _REQUEST_DATA_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) @@ -25,6 +26,7 @@ _DECISIONS_REQUEST_BODY_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = Type _OPENAI_DECISION_REQUEST_BODY_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter( OpenAIDecisionRequestBody ) +_EXTRACTION_REQUEST_BODY_ADAPTER: Final[TypeAdapter[ExtractionRequestBody]] = TypeAdapter(ExtractionRequestBody) _GENERAL_SETTINGS_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) _OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) _OPTIONAL_FLOAT_ADAPTER: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) @@ -73,16 +75,35 @@ _DECISIONS_REQUEST_BODY: Final = inline_request_body( ], }, ) +_EXTRACTION_REQUEST_BODY: Final = inline_request_body( + ExtractionRequest, + { + "model": "clairic", + "state": "Invoice #4471 from Acme Tools, dated 2026-03-03. 12 hammers. Total $174. Payment due in 30 days.", + "fields": { + "vendor": {"type": "text", "instructions": "Who sent the invoice?"}, + "invoice_date": {"type": "date"}, + "total": {"type": "number", "unit": "USD"}, + "items": {"type": "list", "items": "text", "instructions": "What was bought?"}, + "days_to_pay": {"type": "estimate", "min": 0, "max": 90}, + }, + }, +) +_ApiName: TypeAlias = Literal["Decisions", "Extraction"] +_RouteType: TypeAlias = Literal["adecisions", "asystemone", "aextraction"] async def _invalid_request( - raw_data: Mapping[str, object], error: ValidationError, user_api_key_dict: UserAPIKeyAuth + raw_data: Mapping[str, object], + error: ValidationError, + user_api_key_dict: UserAPIKeyAuth, + api_name: _ApiName, ) -> Exception: from litellm.proxy.proxy_server import proxy_logging_obj, version return await ProxyBaseLLMRequestProcessing(data=dict(raw_data)).handle_llm_api_exception( e=BadRequestError( - message=f"Invalid Decisions request: {error}", + message=f"Invalid {api_name} request: {error}", model=str(raw_data.get("model", "")), llm_provider="", ), @@ -95,19 +116,21 @@ async def _invalid_request( _FIELDS_CHECKED_AFTER_ROUTING: Final = frozenset({"safety_identifier", "images", "assets", "audio", "videos"}) -def _fields_checked_before_routing(data: Mapping[str, object]) -> Mapping[str, object]: +def _fields_checked_before_routing(data: Mapping[str, object], route_type: _RouteType) -> Mapping[str, object]: + if route_type == "aextraction": + return data return MappingProxyType({key: value for key, value in data.items() if key not in _FIELDS_CHECKED_AFTER_ROUTING}) _DEPLOYMENT_ONLY_FIELDS: Final = frozenset({"model_info"}) -async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: +async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth, api_name: _ApiName) -> dict[str, object]: body: Final = await request.body() try: data: Final = _REQUEST_DATA_ADAPTER.validate_json(body) except ValidationError as error: - raise await _invalid_request(raw_data={}, error=error, user_api_key_dict=user_api_key_dict) + raise await _invalid_request(raw_data={}, error=error, user_api_key_dict=user_api_key_dict, api_name=api_name) return {key: value for key, value in data.items() if key not in _DEPLOYMENT_ONLY_FIELDS} @@ -115,8 +138,10 @@ async def _process_decisions( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, - body_adapter: TypeAdapter[DecisionsRequestBody] | TypeAdapter[OpenAIDecisionRequestBody], - route_type: Literal["adecisions", "asystemone"], + body_adapter: ( + TypeAdapter[DecisionsRequestBody] | TypeAdapter[OpenAIDecisionRequestBody] | TypeAdapter[ExtractionRequestBody] + ), + route_type: _RouteType, ) -> object: from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, @@ -139,15 +164,17 @@ async def _process_decisions( user_temperature as proxy_user_temperature, ) - data: Final = await _request_data(request, user_api_key_dict) + api_name: Final[_ApiName] = "Extraction" if route_type == "aextraction" else "Decisions" + data: Final = await _request_data(request, user_api_key_dict, api_name) try: - body_adapter.validate_python(_fields_checked_before_routing(data)) + body_adapter.validate_python(_fields_checked_before_routing(data, route_type)) except ValidationError as error: - raise await _invalid_request(raw_data=data, error=error, user_api_key_dict=user_api_key_dict) + raise await _invalid_request(raw_data=data, error=error, user_api_key_dict=user_api_key_dict, api_name=api_name) general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(proxy_general_settings) user_api_base: Final = _OPTIONAL_STRING_ADAPTER.validate_python(proxy_user_api_base) user_model: Final = _OPTIONAL_STRING_ADAPTER.validate_python(proxy_user_model) user_temperature: Final = _OPTIONAL_FLOAT_ADAPTER.validate_python(proxy_user_temperature) + takes_generation_settings: Final = route_type != "aextraction" processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -162,9 +189,9 @@ async def _process_decisions( select_data_generator=None, model=None, user_model=user_model, - user_temperature=user_temperature, + user_temperature=user_temperature if takes_generation_settings else None, user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, + user_max_tokens=user_max_tokens if takes_generation_settings else None, user_api_base=user_api_base, version=version, ) @@ -263,3 +290,53 @@ async def decisions( body_adapter=_OPENAI_DECISION_REQUEST_BODY_ADAPTER, route_type="adecisions", ) + + +@router.post( + "/v1/extraction", + dependencies=[Depends(user_api_key_auth)], + response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract + tags=["extraction"], + summary="Extract typed fields from a state", + responses={200: {"model": ExtractionResponse}}, + response_model=None, + openapi_extra={"requestBody": _EXTRACTION_REQUEST_BODY}, +) +@router.post( + "/extraction", + dependencies=[Depends(user_api_key_auth)], + response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract + tags=["extraction"], + summary="Extract typed fields from a state", + responses={200: {"model": ExtractionResponse}}, + response_model=None, + openapi_extra={"requestBody": _EXTRACTION_REQUEST_BODY}, +) +async def extraction( + request: Request, + fastapi_response: Response, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> object: + """ + Read a `state` text and up to 4 images and get one typed answer per named field in `fields`. + + Field types are `text`, `number`, `date`, `time`, `datetime`, `duration`, `list` (set `items`), `estimate` + (set `min` and `max`) and `forecast` (set `min`, `max` and `horizon`). Each answer carries its `value`, which is + null when the input does not determine it, plus `confidence`, `p_null`, text `span`s into the state and, for + estimates, `mean`, `median`, `p10` and `p90`. Images are data URLs or `{content_type, base64}` objects, remote + URLs are not accepted. An unknown field type, a field name that is not 1 to 64 letters, digits, `_`, `.` or `-` + starting with a letter or `_`, a remote image URL or more than 4 images is rejected with a 400 before any + provider is called. Any other top-level key, `stream: true` and OpenAI generation settings such as `temperature` + included, is rejected with a 400 naming it unless `drop_params` or the deployment's `additional_drop_params` + drops it. Settings inside `fields` and `options` go to the provider as sent, so a field missing a setting its + type needs or an empty `state` with no images comes back with the provider's own error. `model` is any + extraction model in the proxy model list, such as Cloudflare Clairic. The request and the answers keep the + provider's format. + """ + return await _process_decisions( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + body_adapter=_EXTRACTION_REQUEST_BODY_ADAPTER, + route_type="aextraction", + ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index 06ea3ce4b20..7fb87080cd1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -494,7 +494,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): sse: Final = AnthropicSSERestorer(step) events: Final = ResponsesStreamRestorer(step, rehydrate) carries: Final[CarryWindows] = {} - last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush. + last_chunk = None # rebind-ok: tracks the most recent chunk with choices for the final flush. async for chunk in response: if isinstance(chunk, (bytes, str)): @@ -506,7 +506,8 @@ class LLMShieldProxyGuardrail(CustomGuardrail): yield event continue restored_chunk = detached(chunk) - last_chunk = restored_chunk + if getattr(restored_chunk, "choices", None): + last_chunk = restored_chunk for choice in attribute_items(restored_chunk, "choices"): await self._restore_choice(choice, carries, session_id) yield restored_chunk diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1e7f67392e1..684f1619512 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -2155,6 +2155,7 @@ async def test_model_connection( "ocr", "evaluation", "decisions", + "extraction", ] | None = fastapi.Body( None, diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index ef60a800cdd..47c3869fb64 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -254,39 +254,45 @@ class KeyManagementEventHooks: litellm_changed_by: str | None = None, ) -> None: from litellm.proxy.management_helpers.audit_logs import ( - create_audit_log_for_update, + create_audit_logs_for_updates, get_audit_log_changed_by, is_audit_logging_enabled, + track_audit_task, ) from litellm.proxy.proxy_server import litellm_proxy_admin_name if not is_audit_logging_enabled(): return - for key in keys_being_deleted: - key_row = LiteLLM_VerificationToken.model_validate(key, from_attributes=True) - if key_row.token is None: - continue - _key_row = key_row.model_dump_json(exclude_none=True) - - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ), - changed_by_api_key=user_api_key_dict.token, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=key_row.token, - action="deleted", - updated_values="{}", - before_value=_key_row, - ) - ) + changed_by: Final = get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + key_rows: Final = ( + LiteLLM_VerificationToken.model_validate(key, from_attributes=True) for key in keys_being_deleted + ) + audit_rows: Final = tuple( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=changed_by, + changed_by_api_key=user_api_key_dict.token, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=key_row.token, + action="deleted", + updated_values="{}", + before_value=key_row.model_dump_json(exclude_none=True), ) + for key_row in key_rows + if key_row.token is not None + ) + if not audit_rows: + return + track_audit_task( + asyncio.create_task( + create_audit_logs_for_updates(audit_rows), name=f"{len(audit_rows)} deleted key audit logs" + ) + ) @staticmethod async def _store_virtual_key_in_secret_manager(secret_name: str, secret_token: str, team_id: str | None = None): diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 37cef41034b..bc9cee1cb1d 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -374,6 +374,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations = None + self._retired = False + self._flush_after_retire = None self.deployment_budget_config = None @with_service_target("model_budgets") diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d858eff3348..a8ba8d65785 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -7,6 +7,7 @@ This is currently in development and not yet ready for production. import asyncio import binascii import itertools +import json import logging import os import uuid @@ -42,6 +43,7 @@ from litellm.caching.redis_batch import ( ) 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.prompt_templates.common_utils import ( get_str_from_messages, @@ -78,6 +80,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.extraction import ExtractionResponse from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType @@ -490,6 +493,7 @@ RESPONSES_API_CALL_TYPES: Final = ("aresponses", "responses") EMBEDDING_API_CALL_TYPES: Final = ("aembedding", "embedding") TEXT_COMPLETION_API_CALL_TYPES: Final = ("atext_completion", "text_completion") RERANK_API_CALL_TYPES: Final = (CallTypes.rerank.value, CallTypes.arerank.value) +EXTRACTION_API_CALL_TYPES: Final = (CallTypes.extraction.value, CallTypes.aextraction.value) GOOGLE_GENAI_NATIVE_CALL_TYPES: Final = ( CallTypes.generate_content.value, CallTypes.agenerate_content.value, @@ -691,6 +695,7 @@ class RequestRateLimiterStash: owner_litellm_call_id: str | None = None rate_limit_response: RateLimitResponse | None = None parallel_slot: ParallelSlotAcquisition | None = None + parallel_slot_reused_at: datetime | None = None parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) reserved_tokens: int = 0 reserved_model: RateLimitedModel | None = None @@ -787,6 +792,19 @@ def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: return stash +def _ended_before_slot_reuse(ended_at: datetime | None, reused_at: datetime | None) -> bool: + if ended_at is None or reused_at is None or (ended_at.tzinfo is None) != (reused_at.tzinfo is None): + return False + return ended_at <= reused_at + + +def _callback_kwargs_mapping(kwargs: object) -> Mapping[str, object]: + try: + return _REQUEST_RATE_LIMIT_DATA.validate_python(kwargs) + except ValidationError: + return MappingProxyType({}) + + def get_request_stash_for_call(litellm_call_id: str | None) -> RequestRateLimiterStash | None: stash: Final = _request_stash.get() if stash is None: @@ -1094,6 +1112,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): or PROXY_MaxParallelRequestsHandler_v3._has_explicit_output_cap(data, call_type) or is_embedding or endpoint_type == EndpointType.DECISIONS + or call_type in EXTRACTION_API_CALL_TYPES ): return effective_cap: Final = max(capped_floor, configured_output_tokens or 0) @@ -1218,6 +1237,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES else (None, estimable_data.get("prompt"), None) if call_type in TEXT_COMPLETION_API_CALL_TYPES + else (None, self._extraction_input_text(_REQUEST_RATE_LIMIT_DATA.validate_python(estimable_data)), None) + if call_type in EXTRACTION_API_CALL_TYPES else (estimable_data.get("messages"), None, None) if call_type else ( @@ -1245,7 +1266,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type) - is_embedding: Final = self._is_embedding_request(data, call_type) + generates_no_output: Final = ( + self._is_embedding_request(data, call_type) or call_type in EXTRACTION_API_CALL_TYPES + ) base_output_floor: Final = self.no_max_tokens_output_floor(min_configured_tpm_limit) output_floor: Final = ( @@ -1255,7 +1278,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) max_tokens_estimate: Final = ( 0 - if is_embedding or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None) + if generates_no_output + or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None) else explicit_max_tokens if explicit_max_tokens is not None else configured_output_tokens @@ -1642,6 +1666,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): read_only: bool = False, skip_tpm_check: bool = False, parallel_slot_id: str | None = None, + skip_parallel_check: bool = False, ) -> RateLimitResponse: """ Check if any of the rate limit descriptors should be rate limited. @@ -1658,6 +1683,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ``reserve_tpm_tokens`` reservation path should set this to avoid the +1-per-key Lua / in-memory increment double-charging the tokens counter. + skip_parallel_check: If True, ``max_parallel_requests`` gauges are + neither checked nor acquired, for a call that runs inside a + request already holding its slot. ``max_parallel_requests`` descriptors are enforced by the dedicated concurrency-gauge path (``_check_parallel_request_gauges``), never by @@ -1750,7 +1778,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if windowed_response["overall_code"] == "OVER_LIMIT": return windowed_response - if not gauges: + if not gauges or skip_parallel_check: return windowed_response gauge_response: Final = await self._check_parallel_request_gauges( @@ -2013,12 +2041,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span: Span | None, *, in_logging_callback: bool = False, + round_ended_at: datetime | None = None, ) -> None: if stash is None: return async with stash.parallel_slot_release_lock: acquisition: Final = stash.parallel_slot - if acquisition is None: + if acquisition is None or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at): return deferred: Final = in_logging_callback and await self._defer_parallel_slot_release( acquisition, parent_otel_span @@ -3834,6 +3863,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if isinstance(part, (str, dict)) ) + @staticmethod + def _extraction_input_text(data: Mapping[str, object]) -> str: + state: Final = data.get("state") + return f"{state if isinstance(state, str) else ''}\n{json.dumps(data.get('fields'), default=str)}" + def _estimate_precise_input_tokens(self, data: object, model: str | None, call_type: str | None = None) -> int: """ Model-aware input token estimate for the project ITPM reservation, @@ -3892,6 +3926,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if call_type in RERANK_API_CALL_TYPES else (None, prompt, data.get("tools"), data.get("tool_choice")) if call_type in TEXT_COMPLETION_API_CALL_TYPES + else (None, self._extraction_input_text(_REQUEST_RATE_LIMIT_DATA.validate_python(data)), None, None) + if call_type in EXTRACTION_API_CALL_TYPES else (data.get("messages"), fallback_text, data.get("tools"), data.get("tool_choice")) ) messages, selected_text, countable_tools, countable_tool_choice = selected_inputs @@ -3930,6 +3966,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_reservation_amount: int, call_type: str | None = None, endpoint_type: EndpointType = EndpointType.GENERIC, + reuses_request_slot: bool = False, ) -> None: """ Reserve project-scoped ITPM/OTPM tokens (Bedrock Mantle-style @@ -4009,7 +4046,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.reservation_released = True - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=io_response, descriptors=descriptors, @@ -4170,6 +4208,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") stash: Final = claim_request_stash_for_data(data) + reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None + if reuses_request_slot: + stash.parallel_slot_reused_at = self._get_current_time() ######################################################### # Check if the call type has a specific rate limiter @@ -4214,11 +4255,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # then because _reserve_project_io_tokens_or_raise below charges # them unconditionally and counting them here too would # double-charge every request. - parallel_counter_keys: Final = [ - self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") - for d in descriptors - if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None - ] + parallel_counter_keys: Final = ( + [] + if reuses_request_slot + else [ + self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") + for d in descriptors + if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None + ] + ) parallel_slot_id: Final = uuid.uuid4().hex if parallel_counter_keys else None first_pass_descriptors: Final = ( @@ -4233,6 +4278,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, skip_tpm_check=self.tpm_reservation_enabled, parallel_slot_id=parallel_slot_id, + skip_parallel_check=reuses_request_slot, ) if response["overall_code"] == "OVER_LIMIT": @@ -4242,7 +4288,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, ) else: - stash.rate_limit_response = response + if not reuses_request_slot: + stash.rate_limit_response = response if parallel_slot_id is not None: stash.parallel_slot = ParallelSlotAcquisition( slot_id=parallel_slot_id, @@ -4340,7 +4387,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): raw_estimate=estimated_tokens, ) if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit: - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) raise PTUShareExceededError(stash.ptu_ceiling, ptu_estimated_tokens) tpm_response: Final = await self.reserve_tpm_tokens( descriptors=descriptors, @@ -4350,7 +4398,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if tpm_response["overall_code"] == "OVER_LIMIT": - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, @@ -4401,6 +4450,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_reservation_amount=tpm_reservation_amount, call_type=call_type, endpoint_type=endpoint_type, + reuses_request_slot=reuses_request_slot, ) def _create_pipeline_operations( @@ -4512,6 +4562,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) -> Usage | ResponseAPIUsage | Mapping[str, object] | None: if isinstance(response_obj, (Usage, ResponseAPIUsage)): return response_obj + if isinstance(response_obj, ExtractionResponse): + return get_usage_object(response_obj) if isinstance( response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse, BaseLiteLLMOpenAIResponseObject), @@ -5143,21 +5195,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # than parsed out of the body) carry their usage in # ``combined_usage_object`` instead, and would otherwise never charge # the TPM window. - _usage: Usage | dict[str, object] | None = None - if isinstance( - response_obj, - ( - ModelResponse, - EmbeddingResponse, - TextCompletionResponse, - BaseLiteLLMOpenAIResponseObject, - ), - ): - _usage = getattr(response_obj, "usage", None) - else: - _combined_usage: Final = kwargs.get("combined_usage_object") - if isinstance(_combined_usage, Usage): - _usage = _combined_usage + _combined_usage: Final = kwargs.get("combined_usage_object") + _usage: Final[Usage | dict[str, object] | None] = ( + getattr(response_obj, "usage", None) + if isinstance( + response_obj, + ( + ModelResponse, + EmbeddingResponse, + TextCompletionResponse, + BaseLiteLLMOpenAIResponseObject, + ), + ) + else get_usage_object(response_obj) + if isinstance(response_obj, ExtractionResponse) + else _combined_usage + if isinstance(_combined_usage, Usage) + else None + ) total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type) if total_tokens == 0: total_tokens = self._aggregate_only_total_tokens(usage=_usage) @@ -5310,8 +5365,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): try: verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") - stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True) + callback_kwargs: Final = _callback_kwargs_mapping(kwargs) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(callback_kwargs)) + await self._release_stashed_parallel_slot( + stash, + litellm_parent_otel_span, + in_logging_callback=True, + round_ended_at=end_time + if isinstance(end_time, datetime) and not callback_kwargs.get("stream") + else None, + ) pipeline_operations: Final = self._build_success_event_pipeline_operations( kwargs=kwargs, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1f8b352fed9..994310c6619 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -5,14 +5,14 @@ from contextlib import contextmanager from contextvars import ContextVar from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, cast +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final -from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED +from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, @@ -54,11 +54,13 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports _sanitize_error_information_for_spend_logs, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + deployment_params_provider, get_request_model_access_groups, sanitize_error_information_for_spend_logs, should_store_prompts_and_responses_in_spend_logs, ) from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend +from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, LiteLLMBatch, @@ -187,6 +189,14 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +DeploymentLookup: TypeAlias = Callable[[str], Deployment | None] + + +def _proxy_deployment(model_id: str) -> Deployment | None: + router: Final = get_llm_router() + return None if router is None else router.get_deployment(model_id=model_id) + + def _billed_failed_attempts(request_data: Mapping[str, object]) -> tuple[BilledFailedAttempt, ...]: """The earlier billed attempts ``_failure_fields_to_lift`` stashed on the request, if any.""" attempts: Final = request_data.get("billed_failed_attempts") @@ -197,7 +207,10 @@ def _billed_failed_attempts(request_data: Mapping[str, object]) -> tuple[BilledF def _failure_spend_rows( - request_data: Mapping[str, object], final_cost: float, late_guardrails_snapshot: object + request_data: Mapping[str, object], + final_cost: float, + late_guardrails_snapshot: object, + deployment_lookup: DeploymentLookup, ) -> tuple[tuple[float, Mapping[str, object]], ...]: """One ``(response_cost, kwargs)`` per spend row a failure writes: the final attempt under the request id, then every earlier attempt the provider also billed under its own row.""" @@ -209,25 +222,59 @@ def _failure_spend_rows( return ( (final_cost, final_row), *( - (attempt.response_cost, _attempt_spend_row_kwargs(request_data, attempt, position)) + ( + attempt.response_cost, + _attempt_spend_row_kwargs( + request_data, attempt, position, _attempt_deployment(attempt, deployment_lookup) + ), + ) for position, attempt in enumerate(_billed_failed_attempts(request_data)) ), ) +def _attempt_deployment(attempt: BilledFailedAttempt, deployment_lookup: DeploymentLookup) -> Deployment | None: + return None if attempt.deployment_id is None else deployment_lookup(attempt.deployment_id) + + +def _deployment_attribution(deployment: Deployment) -> Mapping[str, object]: + """The spend-log attribution fields of the deployment that served an attempt: the model string it + was called with, its provider, its configured ``api_base`` and its model group.""" + return MappingProxyType( + { + "model": deployment.litellm_params.model, + "custom_llm_provider": deployment_params_provider(deployment.litellm_params), + "api_base": deployment.litellm_params.api_base, + "model_group": deployment.model_name, + } + ) + + +_NO_ATTRIBUTION: Final[Mapping[str, object]] = MappingProxyType({}) + + def _attempt_spend_row_kwargs( - request_data: Mapping[str, object], attempt: BilledFailedAttempt, position: int + request_data: Mapping[str, object], attempt: BilledFailedAttempt, position: int, deployment: Deployment | None ) -> Mapping[str, object]: """The final attempt's row kwargs re-keyed for one earlier billed attempt: its own ``request_id`` (``-attempt-``, since ``request_id`` is the spend-log primary key and the retry count restarts after a fallback hop), its estimated usage and cost, its retry - count under ``attempted_retries`` and the deployment that billed it under ``model_info.id``.""" + count under ``attempted_retries``, the deployment that billed it under ``model_info.id`` and, + when the router still knows that deployment, its model, provider, ``api_base`` and model group. + The row is stamped ``billed_retry_attempt`` so the daily spend tables keep its spend, tokens + and model usage without counting it as a request the caller sent.""" suffix: Final = f"-attempt-{position}" + attribution: Final = _NO_ATTRIBUTION if deployment is None else _deployment_attribution(deployment) raw_params: Final = cast("Mapping[str, object]", request_data.get("litellm_params") or {}) # cast-ok: request dict - litellm_params: Final = dict(raw_params) + litellm_params: Final = { + **raw_params, + **({} if deployment is None else {"api_base": attribution["api_base"]}), + } metadata: Final = { **cast("Mapping[str, object]", litellm_params.get("metadata") or {}), # cast-ok: request dict "attempted_retries": attempt.attempted_retries, + BILLED_RETRY_ATTEMPT_METADATA_KEY: True, + **({} if deployment is None else {"model_group": attribution["model_group"]}), } metadata.pop("previous_models", None) if attempt.deployment_id is not None: @@ -238,6 +285,7 @@ def _attempt_spend_row_kwargs( litellm_params["metadata"] = metadata row: Final = { **request_data, + **({} if deployment is None else {k: attribution[k] for k in ("model", "custom_llm_provider")}), "litellm_params": litellm_params, "combined_usage_object": attempt.usage, "response_cost": attempt.response_cost, @@ -250,7 +298,11 @@ def _attempt_spend_row_kwargs( if isinstance(standard_logging_object, dict): payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict if isinstance(payload.get("id"), str): - row["standard_logging_object"] = {**payload, "id": f"{payload['id']}{suffix}"} + row["standard_logging_object"] = { + **payload, + "id": f"{payload['id']}{suffix}", + **({} if deployment is None else {**attribution, "model_id": attempt.deployment_id}), + } return row @@ -262,10 +314,12 @@ class _ProxyDBLogger(CustomLogger): turn_off_message_logging: bool = False, message_logging: bool = True, spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer, + deployment_lookup: DeploymentLookup = _proxy_deployment, ) -> None: super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging) self.spend_event_producer = spend_event_producer self._spend_writer: Final = spend_writer + self._deployment_lookup: Final = deployment_lookup async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime @@ -439,6 +493,7 @@ class _ProxyDBLogger(CustomLogger): cast("Mapping[str, object]", request_data), # cast-ok: request dict, copied per row below recovered_response_cost, failure_snapshot_with_late_guardrails, + self._deployment_lookup, ), ) diff --git a/litellm/proxy/hooks/search_free_tier.py b/litellm/proxy/hooks/search_free_tier.py index 4ef300dfbac..34f4d50550a 100644 --- a/litellm/proxy/hooks/search_free_tier.py +++ b/litellm/proxy/hooks/search_free_tier.py @@ -34,8 +34,9 @@ def freed_units(allowance: int, units: int, total_after: int) -> int: return max(0, min(units, allowance - (total_after - units))) -def search_free_tier_key(allowance: SearchFreeRequests, now: datetime) -> str: - return f"litellm:search_free_tier:{allowance['shared_group']}:{now.strftime(_PERIOD_STAMP[allowance['period']])}" +def search_free_tier_key(provider: str, allowance: SearchFreeRequests, now: datetime) -> str: + period_stamp: Final = now.strftime(_PERIOD_STAMP[allowance["period"]]) + return f"litellm:search_free_tier:{provider}:{allowance['shared_group']}:{period_stamp}" def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int: @@ -107,7 +108,8 @@ async def apply_search_free_tier( "apply_search_free_tier is on but the proxy has no Redis, grounding is billed at full price" ) return response_cost - key: Final = search_free_tier_key(allowance, now) + provider: Final = sl_object["custom_llm_provider"] or model_info["litellm_provider"] + key: Final = search_free_tier_key(provider, allowance, now) try: total_after: Final = await counter.async_increment(key, units, ttl=search_free_tier_ttl(allowance, now)) except Exception: # noqa: BLE001 # any counter failure means full price, never a lost spend row diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d69778136ee..d34787cd5e9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -29,6 +29,7 @@ from litellm._service_logger import ServiceLogging from litellm._uuid import uuid from litellm.constants import ( BAGGAGE_MAPPED_TAGS_METADATA_KEY, + BILLED_RETRY_ATTEMPT_METADATA_KEY, CALLER_SESSION_ID_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, @@ -267,7 +268,7 @@ def sanitize_for_log(value: object) -> str: from litellm.router import Router from litellm.secret_managers.main import get_secret_bool -from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES +from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES, STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.router import BaggageMappedTagsStamp, holds_api_key_secret_reference from litellm.types.services import ServiceTypes @@ -426,6 +427,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( CONSUMED_REQUEST_TAGS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, + BILLED_RETRY_ATTEMPT_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, "standard_logging_object", "litellm_roi_estimator", @@ -1218,6 +1220,32 @@ def get_dynamic_logging_metadata( return callback_settings_obj +def _callback_settings_fields(callback_settings_obj: TeamCallbackMetadata) -> Mapping[str, object]: + callback_vars: Final = callback_settings_obj.callback_vars + return { + "success_callback": callback_settings_obj.success_callback, + "failure_callback": callback_settings_obj.failure_callback, + **({**callback_vars, TRUSTED_CALLBACK_VARS_FIELD: callback_vars} if callback_vars is not None else {}), + } + + +def key_and_team_logging_fields(user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig) -> Mapping[str, object]: + callback_settings_obj: Final = get_dynamic_logging_metadata( + user_api_key_dict=user_api_key_dict, proxy_config=proxy_config + ) + disabled_callbacks: Final = (user_api_key_dict.metadata or {}).get("litellm_disabled_callbacks") + return MappingProxyType( + { + **(_callback_settings_fields(callback_settings_obj) if callback_settings_obj is not None else {}), + **( + {"litellm_disabled_callbacks": disabled_callbacks} + if disabled_callbacks and isinstance(disabled_callbacks, list) + else {} + ), + } + ) + + _get_dynamic_logging_metadata: Final = get_dynamic_logging_metadata @@ -2342,7 +2370,15 @@ def refresh_proxy_server_request_body_snapshot( if not isinstance(proxy_server_request, dict): return _body_snapshot_exclude: Final = ( - frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG}) + frozenset( + { + "secret_fields", + "proxy_server_request", + "litellm_logging_obj", + VECTOR_STORE_CALLER_METADATA_KWARG, + STREAM_USAGE_INJECTED_FIELD, + } + ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS | _CALLBACK_CREDENTIAL_KEYS ) @@ -2841,27 +2877,7 @@ async def add_litellm_data_to_request( request=request, request_data=data, metadata_key=_metadata_variable_name ) - # Team Callbacks controls - callback_settings_obj: Final = get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=proxy_config - ) - if callback_settings_obj is not None: - data["success_callback"] = callback_settings_obj.success_callback - data["failure_callback"] = callback_settings_obj.failure_callback - - if callback_settings_obj.callback_vars is not None: - # unpack callback_vars in data - for k, v in callback_settings_obj.callback_vars.items(): - data[k] = v - # Callbacks that must not honour request-supplied credentials read this - # proxy-owned field instead of the raw request kwargs. - data[TRUSTED_CALLBACK_VARS_FIELD] = callback_settings_obj.callback_vars - - # Add disabled callbacks from key metadata - if user_api_key_dict.metadata and "litellm_disabled_callbacks" in user_api_key_dict.metadata: - disabled_callbacks: Final = user_api_key_dict.metadata["litellm_disabled_callbacks"] - if disabled_callbacks and isinstance(disabled_callbacks, list): - data["litellm_disabled_callbacks"] = disabled_callbacks + data.update(key_and_team_logging_fields(user_api_key_dict=user_api_key_dict, proxy_config=proxy_config)) # Guardrails from key/team metadata and policy engine await move_guardrails_to_metadata( @@ -3321,12 +3337,12 @@ async def _resolve_user_provider_credentials_for_request( from litellm.proxy.proxy_server import prisma_client, user_api_key_cache from litellm.types.proxy.litellm_pre_call_utils import RedactedDict - from .credential_endpoints.user_provider_credentials import aget_user_provider_tokens + from .credential_endpoints.user_provider_credentials import aget_user_provider_credentials if prisma_client is None: return try: - tokens: Final = await aget_user_provider_tokens( + connections: Final = await aget_user_provider_credentials( prisma_client=prisma_client, cache=user_api_key_cache, user_id=user_id, @@ -3341,7 +3357,7 @@ async def _resolve_user_provider_credentials_for_request( secret_fields: Final = data.get("secret_fields") if not isinstance(secret_fields, dict): return - secret_fields["user_provider_credentials"] = RedactedDict(dict(tokens)) + secret_fields["user_provider_credentials"] = RedactedDict(dict(connections)) secret_fields["user_provider_credentials_user_id"] = user_id diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index ec905a9f187..668bd64227b 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -11,7 +11,7 @@ from dataclasses import dataclass from functools import partial, reduce from itertools import chain, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Final, NamedTuple, NoReturn, Protocol, TypeAlias, overload +from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeAlias, overload from fastapi import ( APIRouter, @@ -53,7 +53,12 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo delete_cache_team_object, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + broadcast_evictions, + evict_and_broadcast, + evict_local, + evict_shared, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export safe_get_request_headers, @@ -64,6 +69,7 @@ from litellm.proxy.management_endpoints.scim.scim_transformations import ( ScimTransformations, ) from litellm.proxy.management_endpoints.team_endpoints import ( + delete_validated_team_rows, new_team, team_member_add, team_member_delete, @@ -104,6 +110,7 @@ from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.management_endpoints.scim_v2 import * if TYPE_CHECKING: + from prisma import Prisma from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken @@ -317,11 +324,15 @@ async def _check_team_exists(team_id: str) -> LiteLLM_TeamTable: team: Final = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) if not team: - raise HTTPException(status_code=404, detail={"error": f"Group not found with ID: {team_id}"}) + raise _group_not_found(team_id) return team +def _group_not_found(team_id: str) -> HTTPException: + return HTTPException(status_code=404, detail={"error": f"Group not found with ID: {team_id}"}) + + def _extract_scim_user_data(user: SCIMUser) -> ScimUserData: """Extract common data from SCIMUser object.""" user_email = None @@ -482,10 +493,7 @@ async def _recompute_scim_member_roles(prisma_client: PrismaClient, user_ids: It Recompute and persist each user's global proxy role from their resulting team membership. No-op unless scim_admin_group is configured, so a SCIM group write that drops a member from the admin group demotes them just like the user - endpoints do, and the role is left untouched when the feature is off. The users' - rows are locked through ``USER_ROWS_LOCK_SQL`` before their teams are read, the - same lock a roster write takes, so a concurrent group push cannot change a - member's teams between the read and the role write. + endpoints do, and the role is left untouched when the feature is off. """ admin_group: Final = await _get_scim_admin_group() if admin_group is None: @@ -494,28 +502,46 @@ async def _recompute_scim_member_roles(prisma_client: PrismaClient, user_ids: It ids: Final = tuple(dict.fromkeys(user_ids)) if not ids: return - default_role: Final = _default_scim_user_role() async with prisma_client.tx(timeout=ROSTER_TX_TIMEOUT) as tx: - await tx.query_raw(USER_ROWS_LOCK_SQL, sorted(ids)) - rows: Final = await find_many_in(user_tx_db(tx), "user_id", ids) - team_ids: Final = tuple(dict.fromkeys(chain.from_iterable(row.teams for row in rows))) - teams: Final = await find_many_in(team_tx_db(tx), "team_id", team_ids) - alias_of: Final = MappingProxyType({team.team_id: team.team_alias for team in teams}) - resolved: Final = tuple( - ( - row.user_id, - _resolve_scim_user_role( - [SCIMUserGroup(value=team_id, display=alias_of.get(team_id)) for team_id in row.teams], - admin_group, - default_role, - ), - ) - for row in rows + await _write_scim_member_roles_tx(tx, ids, admin_group=admin_group) + + from litellm.proxy.proxy_server import user_api_key_cache + + evict_local(cache_keys=ids, user_api_key_cache=user_api_key_cache) + await evict_shared(cache_keys=ids, cache=user_api_key_cache) + await broadcast_evictions(cache_keys=ids) + + +async def _write_scim_member_roles_tx(tx: "Prisma", user_ids: Sequence[str], *, admin_group: str) -> None: + """ + Write each user's role resolved from the teams they are in, on ``tx``. The users' + rows are locked through ``USER_ROWS_LOCK_SQL`` before their teams are read, the + same lock a roster write takes, so a concurrent group push cannot change a + member's teams between the read and the role write. + """ + if not user_ids: + return + default_role: Final = _default_scim_user_role() + await tx.query_raw(USER_ROWS_LOCK_SQL, sorted(user_ids)) + rows: Final = await find_many_in(user_tx_db(tx), "user_id", user_ids) + team_ids: Final = tuple(dict.fromkeys(chain.from_iterable(row.teams for row in rows))) + teams: Final = await find_many_in(team_tx_db(tx), "team_id", team_ids) + alias_of: Final = MappingProxyType({team.team_id: team.team_alias for team in teams}) + resolved: Final = tuple( + ( + row.user_id, + _resolve_scim_user_role( + [SCIMUserGroup(value=team_id, display=alias_of.get(team_id)) for team_id in row.teams], + admin_group, + default_role, + ), + ) + for row in rows + ) + for role in dict.fromkeys(role for _, role in resolved): + await tx.execute_raw( + _SET_USER_ROLE_SQL, role, sorted(user_id for user_id, user_role in resolved if user_role == role) ) - for role in dict.fromkeys(role for _, role in resolved): - await tx.execute_raw( - _SET_USER_ROLE_SQL, role, sorted(user_id for user_id, user_role in resolved if user_role == role) - ) class _ResolvedUserMember(NamedTuple): @@ -3051,7 +3077,8 @@ async def update_group( dependencies=[Depends(user_api_key_auth)], ) async def delete_group( - group_id: str = Path(..., title="Group ID"), + group_id: Annotated[str, Path(title="Group ID")], + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Delete a group according to SCIM v2 protocol @@ -3064,23 +3091,22 @@ async def delete_group( existing_team: Final = await _check_team_exists(group_id) raise_if_teams_hold_ptu_shares((group_id,), llm_router) - member_ids: Final = await _get_team_member_user_ids_from_team(existing_team) - - # For each member, remove this team from their teams list - for member_id in member_ids: - user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": member_id}) - if user: - current_teams = user.teams or [] - if group_id in current_teams: - new_teams = [t for t in current_teams if t != group_id] - await _table(UserRepository(prisma_client)).update( - where={"user_id": member_id}, data={"teams": new_teams} - ) - - await _recompute_scim_member_roles(prisma_client, member_ids) - - # Delete team - await _table(TeamRepository(prisma_client)).delete(where={"team_id": group_id}) + admin_group: Final = await _get_scim_admin_group() + deleted: Final = await delete_validated_team_rows( + teams=(LiteLLM_TeamTable.model_validate(existing_team.model_dump()),), + prisma_client=prisma_client, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key=user_api_key_dict.api_key, + user_id=user_api_key_dict.user_id, + ), + litellm_changed_by=None, + before_commit=( + None if admin_group is None else partial(_write_scim_member_roles_tx, admin_group=admin_group) + ), + ) + if group_id not in deleted.team_ids: + raise _group_not_found(group_id) return Response(status_code=204) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 517a05773c8..d9763f72741 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -14,7 +14,7 @@ import copy import json import math import traceback -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -34,7 +34,7 @@ from typing import ( import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel, Field, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never import litellm @@ -102,6 +102,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo OrganizationNotFoundError, _cache_team_object, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export allowed_route_check_inside_route, + await_inflight_team_membership_load, cache_team_object, can_org_access_model, delete_cache_key_objects, @@ -112,6 +113,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo get_team_object, get_user_object, invalidate_team_member_spend_state, + jwt_key_mapping_cache_key, ) from litellm.proxy.auth.auth_utils import ( enforce_batch_limits_are_admin_only, @@ -119,12 +121,19 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + broadcast_evictions, evict_and_broadcast, + evict_local, + evict_shared, publish_auth_cache_invalidation, ) from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, +) from litellm.proxy.db.db_span import db_span from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import ( @@ -207,7 +216,7 @@ from litellm.proxy.management_helpers.utils import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy from litellm.repositories.budget_repository import BudgetRepository -from litellm.repositories.chunked_in import find_many_in +from litellm.repositories.chunked_in import delete_many_in, find_many_in from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( @@ -417,6 +426,23 @@ class _TeamDeleteTx(AccessGroupSyncTx, Protocol): @property def litellm_teammembership(self) -> "TableActions[prisma_models.LiteLLM_TeamMembership]": ... + @property + def litellm_deletedteamtable(self) -> "TableActions[prisma_models.LiteLLM_DeletedTeamTable]": ... + + @property + def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ... + + @property + def litellm_jwtkeymapping(self) -> "TableActions[prisma_models.LiteLLM_JWTKeyMapping]": ... + + +def _team_delete_tx(tx: "Prisma") -> _TeamDeleteTx: + return tx # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do + + +_LOCK_DELETED_TEAM_MEMBERS_SQL: Final = """ +SELECT user_id, teams FROM "LiteLLM_UserTable" WHERE teams && $1::text[] ORDER BY user_id COLLATE "C" FOR UPDATE +""" _STRIP_DELETED_TEAMS_FROM_USERS_SQL: Final = """ UPDATE "LiteLLM_UserTable" @@ -430,6 +456,22 @@ WHERE user_id IN ( ) """ + +class _DetachedMember(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + user_id: str + teams: tuple[str, ...] + + +_DetachedMembers: Final = TypeAdapter(tuple[_DetachedMember, ...]) + + +@dataclass(frozen=True, slots=True) +class DeletedTeams: + team_ids: tuple[str, ...] + + _INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True}) @@ -4526,18 +4568,7 @@ async def delete_team( }' ``` """ - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - is_audit_logging_enabled, - ) - from litellm.proxy.proxy_server import ( - create_audit_log_for_update, - litellm_proxy_admin_name, - llm_router, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.proxy_server import llm_router, prisma_client if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -4567,47 +4598,32 @@ async def delete_team( raise_if_teams_hold_ptu_shares(data.team_ids, llm_router) - await _persist_deleted_team_records( + await delete_validated_teams( teams=team_rows, prisma_client=prisma_client, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, ) + deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids} + return deleted_teams - # we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes - if is_audit_logging_enabled(): - # make an audit log for each team deleted - for team_id in data.team_ids: - team_row: LiteLLM_TeamTable | None = await prisma_client.get_data( - team_id=team_id, table_name="team", query_type="find_unique" - ) - if team_row is None: - continue +async def delete_validated_teams( + teams: Sequence[LiteLLM_TeamTable], + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: str | None, +) -> DeletedTeams: + """ + What ``/team/delete`` does once the teams exist and the caller may delete them: every key + and BYOK model of the teams, then ``delete_validated_team_rows``. + """ + from litellm.proxy.proxy_server import ( + proxy_logging_obj, + user_api_key_cache, + ) - _team_row = team_row.json(exclude_none=True) - - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.TEAM_TABLE_NAME, - object_id=team_id, - action="deleted", - updated_values="{}", - before_value=_team_row, - ) - ) - ) - - # End of Audit logging + team_ids: Final = [team.team_id for team in teams] ## DELETE ASSOCIATED KEYS # Fetch keys before deletion to persist them @@ -4615,7 +4631,7 @@ async def delete_team( persist_deleted_verification_tokens, ) - keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": data.team_ids}}) + keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": team_ids}}) jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_tokens( hashed_tokens=tuple(key.token for key in keys_to_delete), prisma_client=prisma_client, @@ -4629,7 +4645,7 @@ async def delete_team( litellm_changed_by=litellm_changed_by, ) - await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key") + await prisma_client.delete_data(team_id_list=team_ids, table_name="key") if keys_to_delete: KeyManagementEventHooks.create_key_deleted_audit_logs( @@ -4654,18 +4670,41 @@ async def delete_team( from litellm.proxy.proxy_server import llm_router await delete_team_models( - team_ids=data.team_ids, + team_ids=team_ids, prisma_client=prisma_client, llm_router=llm_router, ) - await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) - - member_ids_per_team: Final = await _resolve_deleted_team_member_user_ids( - teams=team_rows, + return await delete_validated_team_rows( + teams=teams, prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, ) + +async def delete_validated_team_rows( + teams: Sequence[LiteLLM_TeamTable], + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: str | None, + before_commit: Callable[["Prisma", tuple[str, ...]], Awaitable[None]] | None = None, +) -> DeletedTeams: + """ + Under every team's advisory lock, delete the teams still there with the keys their rostered + members hold on them, their membership rows and user references, and write their deleted-team + records; then evict caches and write audit rows. Keys of users off the roster and the teams' + models stay. ``team_ids`` names only the teams this call deleted. ``before_commit`` runs on the + delete's transaction with the members of the rosters read under the lock, so its writes commit + or roll back with the delete. + """ + from litellm.proxy.proxy_server import ( + proxy_logging_obj, + user_api_key_cache, + ) + + team_ids: Final = [team.team_id for team in teams] + ## DELETE TEAMS # Both the delete and the reconcile sweep run under every team's advisory lock # (TEAM_ADVISORY_LOCK_SQL, the same one /team/member_add takes before its own writes), @@ -4674,14 +4713,28 @@ async def delete_team( # the lock before this transaction starts, in which case this sweep reaches what it wrote, # or is still waiting on the lock, in which case its own re-read happens after this commits # and sees the row gone before it writes anything. - delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}} - async with prisma_client.tx(timeout=timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx: - for team_id in sorted(data.team_ids): + tx_budget: Final = timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS) + async with prisma_client.tx(max_wait=tx_budget, timeout=tx_budget) as tx: + for team_id in sorted(team_ids): await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) - await tx.litellm_teamtable.delete_many(where=delete_filter) - await _sweep_deleted_team_references_tx(team_ids=data.team_ids, tx=tx) + team_tx: Final = _team_delete_tx(tx) + locked: Final = await _locked_teams_tx(team_ids=team_ids, tx=team_tx) + member_keys, jwt_mapping_cache_keys = await _delete_roster_member_keys_tx( + locked=locked, + tx=tx, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + await _delete_locked_teams_tx( + locked=locked, tx=team_tx, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by + ) + detached_per_team: Final = await _sweep_deleted_team_references_tx(team_ids=team_ids, tx=team_tx) + if before_commit is not None: + await before_commit(tx, _locked_roster_user_ids(locked)) - deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids} + rostered_per_team: Final = await _resolve_deleted_team_member_user_ids(teams=locked, prisma_client=prisma_client) + member_ids_per_team: Final = _members_per_deleted_team(rostered_per_team, detached_per_team) # Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and # `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a @@ -4689,7 +4742,7 @@ async def delete_team( # invalidated anything. Nothing fallible runs between the delete and this, or a failure there # would strand the deleted team in cache. await _invalidate_deleted_team_cache( - teams=team_rows, + teams=locked, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -4697,44 +4750,167 @@ async def delete_team( member_ids_per_team=member_ids_per_team, user_api_key_cache=user_api_key_cache, ) + await _invalidate_deleted_key_cache( + keys=member_keys, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) - for deleted_team in team_rows: + _schedule_deleted_team_audit_rows( + teams=locked, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by + ) + if member_keys: + KeyManagementEventHooks.create_key_deleted_audit_logs( + keys_being_deleted=member_keys, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + + for deleted_team in locked: _emit_team_members_metric(deleted_team.model_copy(update={"members_with_roles": ()})) await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id) - return deleted_teams + return DeletedTeams(team_ids=tuple(team.team_id for team in locked)) -async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client: PrismaClient) -> None: +def _locked_roster_user_ids(locked: Sequence[LiteLLM_TeamTable]) -> tuple[str, ...]: + return tuple(dict.fromkeys(chain.from_iterable(_roster_user_ids(team) for team in locked))) + + +def _members_per_deleted_team( + rostered: Sequence[tuple[str, Sequence[str]]], detached: Sequence[tuple[str, Sequence[str]]] +) -> tuple[tuple[str, tuple[str, ...]], ...]: + detached_by_team: Final = MappingProxyType({team_id: user_ids for team_id, user_ids in detached}) + return tuple( + (team_id, tuple(dict.fromkeys((*user_ids, *detached_by_team[team_id])))) for team_id, user_ids in rostered + ) + + +def _roster_user_ids(team: LiteLLM_TeamTable) -> tuple[str, ...]: + return tuple(member.user_id for member in team.members_with_roles if member.user_id is not None) + + +def _is_roster_member_key( + key: "prisma_models.LiteLLM_VerificationToken", roster_by_team: Mapping[str, frozenset[str]] +) -> bool: + return key.team_id is not None and key.user_id in roster_by_team.get(key.team_id, frozenset()) + + +async def _locked_teams_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> tuple[LiteLLM_TeamTable, ...]: + rows: Final = await find_many_in(tx.litellm_teamtable, "team_id", team_ids) + return tuple(LiteLLM_TeamTable.model_validate(row.model_dump()) for row in rows) + + +async def _delete_roster_member_keys_tx( + locked: Sequence[LiteLLM_TeamTable], + tx: "Prisma", + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: str | None, +) -> tuple[tuple["prisma_models.LiteLLM_VerificationToken", ...], tuple[str, ...]]: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + persist_deleted_verification_tokens, + ) + + team_tx: Final = _team_delete_tx(tx) + roster_by_team: Final = MappingProxyType({team.team_id: frozenset(_roster_user_ids(team)) for team in locked}) + team_keys: Final = await find_many_in(team_tx.litellm_verificationtoken, "team_id", tuple(roster_by_team)) + member_keys: Final = tuple(key for key in team_keys if _is_roster_member_key(key, roster_by_team)) + if not member_keys: + return (), () + tokens: Final = tuple(key.token for key in member_keys) + jwt_mappings: Final = await find_many_in(team_tx.litellm_jwtkeymapping, "token", tokens) + await persist_deleted_verification_tokens( + keys=member_keys, # pyright: ignore[reportArgumentType] # generated row model carries the same columns as LiteLLM_VerificationToken + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + tx=tx, + ) + await delete_many_in(team_tx.litellm_verificationtoken, "token", tokens, atomicity="caller_transaction") + return member_keys, tuple( + jwt_key_mapping_cache_key(mapping.jwt_claim_name, mapping.jwt_claim_value, mapping.jwt_issuer) + for mapping in jwt_mappings + ) + + +async def _delete_locked_teams_tx( + locked: Sequence[LiteLLM_TeamTable], + tx: _TeamDeleteTx, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: str | None, +) -> None: + await _persist_deleted_team_records( + teams=locked, + deleted_teams=tx.litellm_deletedteamtable, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + await delete_many_in( + tx.litellm_teamtable, + "team_id", + tuple(team.team_id for team in locked), + atomicity="caller_transaction", + ) + + +def _schedule_deleted_team_audit_rows( + teams: Sequence[LiteLLM_TeamTable], + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: str | None, +) -> None: + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + is_audit_logging_enabled, + ) + from litellm.proxy.proxy_server import ( + create_audit_log_for_update, + litellm_proxy_admin_name, + ) + + if not is_audit_logging_enabled(): + return + for team in teams: + asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id=team.team_id, + action="deleted", + updated_values="{}", + before_value=team.model_dump_json(exclude_none=True), + ) + ) + ) + + +async def _sweep_deleted_team_references_tx( + team_ids: Sequence[str], tx: _TeamDeleteTx +) -> tuple[tuple[str, tuple[str, ...]], ...]: """ - Strip the deleted team ids from every user row and team-membership row that still references them. - - The per-member `team_member_delete` pass above only reaches users listed in the team's - `members_with_roles`, so a user row that outlived its roster entry is invisible to it and keeps - surfacing the team on `/user/info` after the team is gone. - - #36839 closed the route that created that drift, by resolving member removal off the roster - entry's `user_id` rather than the identifier the caller happened to pass. It does not backfill - rows that already drifted, which is the state this was reported against, so the sweep still has - to run on delete. - - One statement filters the array in place rather than read-filter-write: rewriting the whole - array from a snapshot read outside a transaction drops any team a concurrent `/team/member_add` - appended in between. It locks the user rows in the order `USER_ROWS_LOCK_SQL` takes them, so it - never waits on a SCIM group push in a cycle. + Strip the deleted team ids from every user and team-membership row still referencing them, on + the transaction holding the teams' advisory locks. User rows are locked first in + ``USER_ROWS_LOCK_SQL`` order so the sweep never deadlocks with a SCIM group push, and one + statement filters each array in place so a team a concurrent ``/team/member_add`` appended is + kept. Returns, per team, the users read under the lock as referencing it. """ - _ = await prisma_client.db.execute_raw(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, list(team_ids)) - - _ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)})) - - -async def _sweep_deleted_team_references_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> None: - """Same sweep as `_sweep_deleted_team_references`, run on the transaction that holds - every id's advisory lock and deletes the team rows, so it commits or rolls back with them.""" + locked: Final = _DetachedMembers.validate_python(await tx.query_raw(_LOCK_DELETED_TEAM_MEMBERS_SQL, list(team_ids))) _ = await tx.execute_raw(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, list(team_ids)) membership_filter: Final[_TeamIdInFilter] = {"team_id": {"in": tuple(team_ids)}} _ = await tx.litellm_teammembership.delete_many(where=membership_filter) + return tuple( + (team_id, tuple(member.user_id for member in locked if team_id in member.teams)) for team_id in team_ids + ) async def _invalidate_deleted_key_cache( @@ -4791,17 +4967,20 @@ async def _evict_deleted_team_member_cache( member_user_ids: Sequence[str], user_api_key_cache: UserApiKeyCache, ) -> None: - await evict_and_broadcast(cache_keys=tuple(member_user_ids), user_api_key_cache=user_api_key_cache) - await asyncio.gather( - *( - invalidate_team_member_spend_state( - user_id=user_id, - team_id=team_id, - user_api_key_cache=user_api_key_cache, - ) - for user_id in member_user_ids - ) + _ = await asyncio.gather( + *(await_inflight_team_membership_load(user_id=user_id, team_id=team_id) for user_id in member_user_ids) ) + cache_keys: Final = tuple(member_cache_keys(team_id, member_user_ids)) + evict_local(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache) + await evict_shared(cache_keys=cache_keys, cache=user_api_key_cache) + await broadcast_evictions(cache_keys=cache_keys) + + +def member_cache_keys(team_id: str, user_ids: Sequence[str]) -> Iterator[str]: + for user_id in user_ids: + yield user_id + yield team_membership_auth_cache_key(team_id=team_id, user_id=user_id) + yield team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) async def _resolve_deleted_team_member_user_ids( @@ -4832,7 +5011,7 @@ async def _deleted_team_member_user_ids(team: LiteLLM_TeamTable, prisma_client: def _transform_teams_to_deleted_records( - teams: list[LiteLLM_TeamTable], + teams: Sequence[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None = None, ) -> list[dict[str, object]]: @@ -4881,17 +5060,17 @@ def _transform_teams_to_deleted_records( async def _save_deleted_team_records( records: list[dict[str, object]], - prisma_client: PrismaClient, + deleted_teams: "TableActions[prisma_models.LiteLLM_DeletedTeamTable]", ) -> None: """Save deleted team records to the database.""" if not records: return - await _deleted_team_db(prisma_client).create_many(data=records) + await deleted_teams.create_many(data=records) async def _persist_deleted_team_records( - teams: list[LiteLLM_TeamTable], - prisma_client: PrismaClient, + teams: Sequence[LiteLLM_TeamTable], + deleted_teams: "TableActions[prisma_models.LiteLLM_DeletedTeamTable]", user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None = None, ) -> None: @@ -4903,7 +5082,7 @@ async def _persist_deleted_team_records( ) await _save_deleted_team_records( records=records, - prisma_client=prisma_client, + deleted_teams=deleted_teams, ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 4cc849762aa..19bca3d2bf8 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -55,6 +55,7 @@ from litellm.constants import ( CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, CLI_SSO_SESSION_CACHE_KEY_PREFIX, CLI_SSO_SESSION_TTL_SECONDS, + CP_RETURN_TO_COOKIE_PREFIX, LITELLM_CLI_SOURCE_IDENTIFIER, LITELLM_UI_SESSION_DURATION, MAX_SPENDLOG_ROWS_TO_QUERY, @@ -2285,9 +2286,7 @@ async def auth_callback(request: Request, state: str | None = None): sso_assertion=sso_assertion, ) - # Control-plane cross-origin: read return_to from cookie. - # Starlette's cookie_parser already handles RFC 2109 unquoting. - cp_return_to: Final[str | None] = request.cookies.get("litellm_cp_return_to") + cp_return_to: Final[str | None] = read_return_to_cookie(request) return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=result, @@ -3090,10 +3089,11 @@ def _persist_return_to_cookie(response: Response, return_to: str | None, request safe: Final = _is_same_origin_return_path(return_to) or SSOAuthenticationHandler.validate_return_to(return_to) except HTTPException: return # a non-matching absolute return_to is ignored, never blocks sign-in - if safe: + sealed: Final = _seal_return_to(return_to) if safe else None + if sealed is not None: response.set_cookie( key="litellm_cp_return_to", - value=return_to, + value=sealed, max_age=600, httponly=True, samesite="lax", @@ -3101,6 +3101,45 @@ def _persist_return_to_cookie(response: Response, return_to: str | None, request ) +def _seal_return_to(return_to: str) -> str | None: + """``return_to`` sealed under the salt key as an opaque cookie value. The plaintext URL carries the + client's loopback ``redirect_uri`` (``http://127.0.0.1:/...`` for Codex or Claude Code), which + edge firewalls read as an SSRF probe and block before the proxy sees the request. None when there + is no salt or master key yet, so the login entrypoint still renders.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token + + try: + return encrypt_bearer_token(return_to, prefix=CP_RETURN_TO_COOKIE_PREFIX) + except ValueError: + return None + + +def _clear_rejected_return_to_cookie(response: Response, request: Request) -> None: + """Expire a ``litellm_cp_return_to`` that ``read_return_to_cookie`` rejected (plaintext, stale, or sealed + under another key) so the browser stops replaying it until it expires on its own.""" + if "litellm_cp_return_to" in request.cookies: + response.delete_cookie("litellm_cp_return_to") + + +def _dashboard_redirect_response(request: Request, jwt_token: str, litellm_dashboard_ui: str) -> RedirectResponse: + """The post-SSO dashboard landing, used when there is no honored ``return_to``. Sets the session + cookie and expires a ``litellm_cp_return_to`` the reader rejected.""" + verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui) + redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) + set_session_token_cookie(redirect_response, request, jwt_token) + _clear_rejected_return_to_cookie(redirect_response, request) + return redirect_response + + +def read_return_to_cookie(request: Request) -> str | None: + """The ``return_to`` this proxy sealed into ``litellm_cp_return_to``, or None when the cookie is + missing, was not sealed here, or was tampered with.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_bearer_token + + sealed: Final = request.cookies.get("litellm_cp_return_to") + return None if sealed is None else decrypt_bearer_token(sealed, prefix=CP_RETURN_TO_COOKIE_PREFIX) + + def _is_blank_sso_user_id(user_id: object) -> bool: return isinstance(user_id, str) and not user_id.strip() @@ -3976,10 +4015,7 @@ class SSOAuthenticationHandler: if user_id is not None and isinstance(user_id, str): litellm_dashboard_ui += "?login=success" - verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui) - redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - set_session_token_cookie(redirect_response, request, jwt_token) - return redirect_response + return _dashboard_redirect_response(request, jwt_token, litellm_dashboard_ui) @staticmethod @with_service_target(SSO_SESSIONS_TARGET) diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index ef95e8e2c0f..aa31f44dd13 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -4,6 +4,7 @@ Functions to create audit logs for LiteLLM Proxy import asyncio import json +from collections.abc import Awaitable, Sequence from datetime import datetime, timezone from typing import Final, TypeVar @@ -258,12 +259,19 @@ async def create_object_audit_log( ) -async def _insert_audit_log_row(repository: AuditLogRepository, request_data: LiteLLM_AuditLogs) -> None: +async def _insert_audit_logs(insert: Awaitable[object], description: str) -> None: try: - await repository.table.create(data={**request_data.model_dump(exclude_none=True)}) + await insert except Exception as e: # [Non-Blocking Exception. Do not allow blocking LLM API call] - verbose_proxy_logger.error("Failed Creating audit log %s", e) + verbose_proxy_logger.error("Failed Creating %s %s", description, e) + + +def _json_encoded(request_data: LiteLLM_AuditLogs) -> LiteLLM_AuditLogs: + values: Final = (("updated_values", request_data.updated_values), ("before_value", request_data.before_value)) + return request_data.model_copy( + update={field: json.dumps(value) for field, value in values if isinstance(value, dict)} + ) async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): @@ -279,15 +287,10 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): return verbose_proxy_logger.debug("creating audit log for %s", request_data) - - if isinstance(request_data.updated_values, dict): - request_data.updated_values = json.dumps(request_data.updated_values) - - if isinstance(request_data.before_value, dict): - request_data.before_value = json.dumps(request_data.before_value) + encoded: Final = _json_encoded(request_data) # Dispatch to external audit log callbacks regardless of DB availability - await _dispatch_audit_log_to_callbacks(request_data) + await _dispatch_audit_log_to_callbacks(encoded) if prisma_client is None: verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB") @@ -295,7 +298,41 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): await track_audit_task( asyncio.create_task( - _insert_audit_log_row(AuditLogRepository(prisma_client), request_data), - name=_audit_task_name(request_data, "the database"), + _insert_audit_logs( + AuditLogRepository(prisma_client).table.create(data={**encoded.model_dump(exclude_none=True)}), + "audit log", + ), + name=_audit_task_name(encoded, "the database"), + ) + ) + + +async def create_audit_logs_for_updates(rows: Sequence[LiteLLM_AuditLogs]) -> None: + """Create audit logs for many objects with one insert.""" + if not rows or not is_audit_logging_enabled(): + return + + from litellm.proxy.proxy_server import premium_user, prisma_client + + if premium_user is not True: + return + + encoded: Final = tuple(_json_encoded(row) for row in rows) + for row in encoded: + await _dispatch_audit_log_to_callbacks(row) + + if prisma_client is None: + verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB") + return + + await track_audit_task( + asyncio.create_task( + _insert_audit_logs( + AuditLogRepository(prisma_client).table.create_many( + data=tuple(row.model_dump(exclude_none=True) for row in encoded) + ), + f"{len(encoded)} audit logs", + ), + name=f"{len(encoded)} audit logs to the database", ) ) diff --git a/litellm/proxy/management_helpers/model_group_custom_fields.py b/litellm/proxy/management_helpers/model_group_custom_fields.py new file mode 100644 index 00000000000..2e94d47a952 --- /dev/null +++ b/litellm/proxy/management_helpers/model_group_custom_fields.py @@ -0,0 +1,68 @@ +import math +from collections.abc import Iterator, Mapping, Sequence +from types import MappingProxyType +from typing import Final, TypeAlias + +from pydantic import TypeAdapter + +CustomMetadataValue: TypeAlias = str | int | float | bool + +_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[object, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) + + +def _string_tags(tags: object) -> Iterator[str]: + if isinstance(tags, list): + yield from (tag for tag in _OBJECT_LIST_ADAPTER.validate_python(tags) if isinstance(tag, str)) + + +def _is_non_finite(value: object) -> bool: + return isinstance(value, float) and not math.isfinite(value) + + +def _scalar_metadata(metadata: object) -> Iterator[tuple[str, CustomMetadataValue]]: + if isinstance(metadata, dict): + yield from ( + (key, value) + for key, value in _OBJECT_MAPPING_ADAPTER.validate_python(metadata).items() + if isinstance(key, str) and isinstance(value, (str, int, float, bool)) and not _is_non_finite(value) + ) + + +def _custom_tag_values(model_infos: Sequence[Mapping[object, object]]) -> Iterator[str]: + for model_info in model_infos: + yield from _string_tags(model_info.get("custom_tags")) + + +def _custom_metadata_entries( + model_infos: Sequence[Mapping[object, object]], +) -> Iterator[tuple[str, CustomMetadataValue]]: + for model_info in model_infos: + yield from _scalar_metadata(model_info.get("custom_metadata")) + + +def _merged_custom_tags(model_infos: Sequence[Mapping[object, object]]) -> tuple[str, ...] | None: + return tuple(dict.fromkeys(_custom_tag_values(model_infos))) or None + + +def _merged_custom_metadata( + model_infos: Sequence[Mapping[object, object]], +) -> Mapping[str, CustomMetadataValue] | None: + entries: Final = tuple(_custom_metadata_entries(model_infos)) + first_seen: Final = MappingProxyType(dict(reversed(entries))) + return MappingProxyType({key: first_seen[key] for key in dict.fromkeys(key for key, _ in entries)}) or None + + +def _model_info(deployment: Mapping[str, object]) -> Mapping[object, object]: + model_info: Final = deployment.get("model_info") + return _OBJECT_MAPPING_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else MappingProxyType({}) + + +def model_group_custom_fields(deployments: Sequence[Mapping[str, object]]) -> Mapping[str, object]: + model_infos: Final = tuple(_model_info(deployment) for deployment in deployments) + return MappingProxyType( + { + "custom_tags": _merged_custom_tags(model_infos), + "custom_metadata": _merged_custom_metadata(model_infos), + } + ) diff --git a/litellm/proxy/management_helpers/team_roster_sync.py b/litellm/proxy/management_helpers/team_roster_sync.py index 4b29d4676f8..dc0ba2e47d4 100644 --- a/litellm/proxy/management_helpers/team_roster_sync.py +++ b/litellm/proxy/management_helpers/team_roster_sync.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import timedelta from types import MappingProxyType @@ -23,8 +23,6 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( ) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, - team_membership_auth_cache_key, - team_membership_reservation_cache_key, ) from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -33,6 +31,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.team_endpoints import ( _emit_team_members_metric, # pyright: ignore[reportPrivateUsage] # same gauge /team/member_add moves _schedule_team_membership_audit_log, # pyright: ignore[reportPrivateUsage] # same audit row /team/member_add writes + member_cache_keys, ) from litellm.proxy.management_helpers.access_group_team_sync import ( DETACH_TEAM_FROM_USERS_SQL, @@ -277,7 +276,7 @@ async def _settle_caches( evicted: Final = ( *deleted_tokens, *removal.jwt_mapping_cache_keys, - *_member_cache_keys(team.team_id, changed_user_ids), + *member_cache_keys(team.team_id, changed_user_ids), ) evict_local(cache_keys=evicted, user_api_key_cache=user_api_key_cache) await evict_shared(cache_keys=evicted, cache=user_api_key_cache) @@ -292,13 +291,6 @@ async def _settle_caches( await broadcast_evictions(cache_keys=evicted) -def _member_cache_keys(team_id: str, user_ids: Sequence[str]) -> Iterator[str]: - for user_id in user_ids: - yield user_id - yield team_membership_auth_cache_key(team_id=team_id, user_id=user_id) - yield team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) - - async def sync_team_roster( prisma_client: PrismaClient, team_id: str, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fe7ac30e434..e1e8864678e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -880,8 +880,8 @@ def _build_passthrough_failure_request_payload( @dataclass(frozen=True, slots=True) class _TeamCallbackWiring: - success_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg - failure_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg + success_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg + failure_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg logging_kwargs: dict[str, str | dict[str, str]] | None = None # mutable-ok: Logging.__init__ arg diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c6d3e12a378..b48b5ab1c1e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -166,6 +166,7 @@ from litellm.proxy.common_utils.token_count_settings import ( resolve_token_count_settings, ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog +from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields from litellm.router_utils.access_windows import access_windows_config_error from litellm.router_utils.add_retry_fallback_headers import ( get_fallback_errors_from_headers, @@ -975,6 +976,7 @@ from litellm.secret_managers.main import ( ) from litellm.tracing.config import is_lens_tracing_enabled from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.anthropic import ( AnthropicMessagesRequest, AnthropicResponse, @@ -17747,10 +17749,7 @@ def get_model_group_info( ) -> list[ModelGroupInfoProxy]: model_groups: Final[list[ModelGroupInfoProxy]] = [] - unique_models: Final = [] - for model in all_models_str: - if model not in unique_models: - unique_models.append(model) + unique_models: Final = tuple(dict.fromkeys(all_models_str)) for model in unique_models: if model_group is not None and model_group != model: @@ -17759,7 +17758,18 @@ def get_model_group_info( _model_group_info = llm_router.get_model_group_info(model_group=model, team_id=team_id) if _model_group_info is not None: - model_groups.append(ModelGroupInfoProxy.model_validate(_model_group_info.model_dump())) + model_groups.append( + ModelGroupInfoProxy.model_validate( + { + **_model_group_info.model_dump(), + **model_group_custom_fields( + llm_router.get_model_list_of_routed_group( + resolve_model_group_alias(llm_router.model_group_alias, model) or model + ) + ), + } + ) + ) else: model_group_info = ModelGroupInfoProxy( model_group=model, @@ -18210,6 +18220,7 @@ async def async_queue_request( try: data = await request.json() data.pop("_litellm_strip_stream_usage", None) + data.pop(STREAM_USAGE_INJECTED_FIELD, None) # Include original request and headers in the data data["proxy_server_request"] = { @@ -18400,6 +18411,7 @@ async def login(request: Request): # _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the # one-shot cookie is cleared after use. from litellm.proxy.management_endpoints.ui_sso import ( + read_return_to_cookie, set_session_token_cookie, sso_return_to_redirect, ) @@ -18411,7 +18423,7 @@ async def login(request: Request): # return_to is honored for EVERY sign-in branch, so the write and read sets cannot diverge: it # sets the token cookie on the same-origin arm and hands off via a one-time login code on the # cross-origin arm, and clears the one-shot cookie in both. - cp_return_to: Final = request.cookies.get("litellm_cp_return_to") + cp_return_to: Final = read_return_to_cookie(request) if cp_return_to: try: resumed = await sso_return_to_redirect( # rebind-ok: pre-existing rebinding on a rename-only line @@ -18435,7 +18447,7 @@ async def login(request: Request): # Create redirect response with cookie redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) set_session_token_cookie(redirect_response, request, jwt_token) - if cp_return_to: + if "litellm_cp_return_to" in request.cookies: redirect_response.delete_cookie(key="litellm_cp_return_to") return redirect_response diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index a794854ce19..4bec42a7868 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1731,6 +1731,16 @@ "field_type": "password", "options": null, "default_value": null + }, + { + "key": "github_copilot_host", + "label": "GitHub host", + "placeholder": "github.com", + "tooltip": "Leave empty for github.com. For GHE.com data residency enter your SUBDOMAIN.ghe.com host", + "required": false, + "field_type": "text", + "options": null, + "default_value": null } ], "default_model_placeholder": "gpt-3.5-turbo" diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 8903b22c120..6ced7827846 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy._types import ( CommonProxyErrors, ) +from litellm.proxy.db.health_check_latest import latest_by_model_name from litellm.proxy.utils import get_custom_url from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.router_strategy.complexity_router.fuse_presets import FusePresetCatalog, get_fuse_presets @@ -295,13 +296,8 @@ async def public_model_hub(): if prisma_client is not None: try: latest_checks: Final = await prisma_client.get_all_latest_health_checks() - for check in latest_checks: - key = check.model_id if check.model_id else check.model_name - if key: - health_check_dict = convert_health_check_to_dict(check) - health_checks_map[key] = health_check_dict - if check.model_name: - health_checks_map[check.model_name] = health_check_dict + for model_name, check in latest_by_model_name(latest_checks).items(): + health_checks_map[model_name] = convert_health_check_to_dict(check) except Exception: pass diff --git a/litellm/proxy/public_endpoints/public_v1/model_hub.py b/litellm/proxy/public_endpoints/public_v1/model_hub.py index 2f5eeeb3d4b..cd634e490b6 100644 --- a/litellm/proxy/public_endpoints/public_v1/model_hub.py +++ b/litellm/proxy/public_endpoints/public_v1/model_hub.py @@ -12,6 +12,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.health_check_latest import latest_by_model_name from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem from litellm.proxy.list_api.in_memory import Cells, InMemoryListExecutor from litellm.proxy.list_api.list_framework import ( @@ -59,12 +60,12 @@ class PrismaHealthSnapshotLookup: checks: Final = await self.prisma_client.get_latest_health_checks_for_models(model_groups) return MappingProxyType( { - check.model_name: HealthSnapshot( + model_group: HealthSnapshot( status=check.status, response_time_ms=check.response_time_ms, checked_at=check.checked_at.isoformat() if check.checked_at else None, ) - for check in checks + for model_group, check in latest_by_model_name(checks).items() } ) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index e16e14c8d54..cb226929167 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,8 +1,13 @@ import asyncio -from collections.abc import Mapping +import hashlib +import json +from collections import OrderedDict +from collections.abc import Callable, Mapping from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, Literal +from uuid import uuid4 import httpx from fastapi import HTTPException, status @@ -10,6 +15,7 @@ from pydantic import BeforeValidator, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import is_proxy_admin_request @@ -54,7 +60,34 @@ _ROUTER_ONLY_SETTINGS: Final = frozenset( ) -def _route_user_config_request(data: dict, route_type: str): +class UserConfigRouters: + def __init__(self, max_size: int) -> None: + self._max_size: Final = max_size + self._routers: Final[OrderedDict[str, LitellmRouter]] = OrderedDict() # mutable-ok: LRU evicts in place + + def get(self, config: Mapping[str, object], build: Callable[[], LitellmRouter]) -> LitellmRouter: + key: Final = hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest() + cached: Final = self._routers.get(key) + if cached is not None: + self._routers.move_to_end(key) + return cached + router: Final = build() + router.serve_through_request_logging() + self._routers[key] = router + if len(self._routers) > self._max_size: + self._routers.popitem(last=False)[1].discard() + return router + + def clear(self) -> None: + for router in self._routers.values(): + router.discard() + self._routers.clear() + + +user_config_routers: Final = UserConfigRouters(max_size=32) + + +async def _route_user_config_request(data: dict, route_type: str): """Route a request using the user-provided router config.""" router_config: Final = data.pop("user_config") @@ -63,10 +96,24 @@ def _route_user_config_request(data: dict, route_type: str): valid_args: Final = litellm.Router.get_valid_args() filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args} - user_router: Final = litellm.Router(**filtered_config) - ret_val: Final = getattr(user_router, f"{route_type}")(**data) - user_router.discard() - return ret_val + logging_obj: Final = _request_logging_obj(JSON_OBJECT_ADAPTER.validate_python(data), route_type) + cached_router: Final = user_config_routers.get(filtered_config, lambda: litellm.Router(**filtered_config)) + cached_router.bind_to_request_logging(logging_obj) + return await getattr(cached_router, f"{route_type}")(**{**data, "litellm_logging_obj": logging_obj}) + + +def _request_logging_obj(data: Mapping[str, object], route_type: str) -> LiteLLMLoggingObj: + arrived_with: Final = data.get("litellm_logging_obj") + if isinstance(arrived_with, LiteLLMLoggingObj): + return arrived_with + logging_obj, _ = litellm.utils.function_setup( + original_function=route_type, + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), # noqa: DTZ005 # Logging subtracts it from naive end times + is_async_call=True, + **{"litellm_call_id": str(uuid4()), **data}, + ) + return logging_obj def _is_a2a_agent_model(model_name: object) -> bool: @@ -111,6 +158,7 @@ ROUTE_ENDPOINT_MAPPING: Final = { "asearch": "/search", "adecisions": "/decisions", "asystemone": "/systemone", + "aextraction": "/extraction", "avideo_generation": "/videos", "avideo_list": "/videos", "avideo_status": "/videos/{video_id}", @@ -573,6 +621,7 @@ RouteType = Literal[ "asearch", "adecisions", "asystemone", + "aextraction", "avideo_generation", "avideo_list", "avideo_status", diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 091fb5d1305..62aee788ed0 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -15,6 +15,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + BILLED_RETRY_ATTEMPT_METADATA_KEY, CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, EMPTY_MAPPING, @@ -187,6 +188,7 @@ _STAMPED_METADATA_KEYS: Final = frozenset( "used_client_oauth_token", KUBERNETES_POD_ROUTING_KEY, "litellm_roi_estimator", + BILLED_RETRY_ATTEMPT_METADATA_KEY, "user_api_key_user_email", ) ) @@ -249,6 +251,7 @@ def _get_spend_logs_metadata( guardrail_information=None, internal_call_origin=None, litellm_roi_estimator=False, + billed_retry_attempt=False, eval_information=None, cold_storage_object_key=cold_storage_object_key, litellm_overhead_time_ms=None, @@ -270,6 +273,8 @@ def _get_spend_logs_metadata( ) verbose_proxy_logger.debug("getting payload for SpendLogs, available keys in metadata: %s", [*metadata]) + retry_key: Final = BILLED_RETRY_ATTEMPT_METADATA_KEY + billed_retry_attempt: Final = retry_key in metadata and metadata[retry_key] is True # Filter the metadata dictionary to include only the specified keys clean_metadata: Final = SpendLogsMetadata( **MappingProxyType( @@ -283,6 +288,7 @@ def _get_spend_logs_metadata( used_client_oauth_token=used_client_oauth_token, kubernetes_pod_routing=kubernetes_pod_routing, litellm_roi_estimator=metadata.get("litellm_roi_estimator") is True, + billed_retry_attempt=billed_retry_attempt, user_api_key_user_email=_stored_user_email(metadata), ) _raw_key: Final = clean_metadata.get("user_api_key") @@ -453,7 +459,12 @@ def _sl_attribution_fallback( def _deployment_provider(deployment: DeploymentTypedDict) -> str | None: - litellm_params: Final = LiteLLM_Params.model_validate(deployment["litellm_params"]) + litellm_params: Final[LiteLLM_Params] = LiteLLM_Params.model_validate(deployment["litellm_params"]) + return deployment_params_provider(litellm_params) + + +def deployment_params_provider(litellm_params: LiteLLM_Params) -> str | None: + """The provider a deployment's ``litellm_params`` authenticate against, None when its model cannot be resolved.""" if litellm.LiteLLMProxyChatConfig.should_use_litellm_proxy_by_default(litellm_params=litellm_params): return LlmProviders.LITELLM_PROXY.value declared: Final = declared_authenticating_provider(litellm_params.model, litellm_params.custom_llm_provider) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 20a0d4d5053..6c9b23860da 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,7 +91,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse -from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, CostPerToken, ModelInfo, Usage +from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage try: from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( @@ -158,7 +158,8 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int -from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token 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.served_output_texts import ( @@ -268,6 +269,7 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.router_utils.common_utils import resolve_model_group_alias +from litellm.router_utils.deployment_pricing import pricing_models from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES from litellm.types.llms.openai import ResponsesAPIResponse @@ -1159,6 +1161,43 @@ def _timed_out_earlier_records( return tuple(record for identity, record in by_identity.items() if identity != final_identity) +def _record_pricing( + model_call_details: Mapping[str, object], record: Mapping[str, object], base_model: str | None +) -> tuple[Mapping[str, object], str | None]: + """The ``(model_call_details, base_model)`` to price one earlier attempt with: the record's own + deployment model and ``base_model`` when it served a different model than the final attempt + (the provider is then inferred from that model), else the final attempt's.""" + record_model: Final = record.get("model") + if not isinstance(record_model, str) or record_model == model_call_details.get("model"): + return model_call_details, base_model + record_base_model: Final = record.get("base_model") + return ( + {**model_call_details, "model": record_model, "custom_llm_provider": None}, + record_base_model if isinstance(record_base_model, str) else None, + ) + + +def _billed_earlier_attempt( + model_call_details: Mapping[str, object], + record: Mapping[str, object], + index: int, + estimated_usage: Usage, + base_model: str | None, +) -> BilledFailedAttempt: + pricing_details, pricing_base_model = _record_pricing(model_call_details, record, base_model) + return BilledFailedAttempt( + attempted_retries=_attempt_number(record, index), + deployment_id=_record_deployment_id(record), + usage=estimated_usage, + response_cost=_estimated_input_cost( + pricing_details, + estimated_usage.prompt_tokens, + deployment_id=_record_deployment_id(record), + base_model=pricing_base_model, + ), + ) + + def _billed_earlier_attempts( model_call_details: Mapping[str, object], litellm_params: Mapping[str, object] | None, @@ -1172,54 +1211,45 @@ def _billed_earlier_attempts( request's own deployment and retry count (``attempted_retries`` restarts at 0 after a fallback hop, so the count alone cannot tell the groups apart). An earlier attempt that ended in a ``Timeout`` had the same prompt accepted by its deployment, so it is billed like the final - attempt, at that deployment's input rate (else the cost map's rate for ``base_model`` or the - model). Only the exception class survives on the record, so an earlier connect timeout is - billed too; any other failure class is not.""" + attempt, at that deployment's input rate (else the cost map's rate for the record's own + ``base_model`` or model, which a fallback hop can make a different model than the final one). + Only the exception class survives on the record, so an earlier connect timeout is billed too; + any other failure class is not.""" final_attempt: Final = _retry_metadata_bucket(litellm_params).get("attempted_retries") final_identity: Final = (deployment_id, final_attempt if isinstance(final_attempt, int) else None) return tuple( - BilledFailedAttempt( - attempted_retries=_attempt_number(record, index), - deployment_id=_record_deployment_id(record), - usage=estimated_usage, - response_cost=_estimated_input_cost( - model_call_details, - estimated_usage.prompt_tokens, - deployment_id=_record_deployment_id(record), - base_model=base_model, - ), - ) + _billed_earlier_attempt(model_call_details, record, index, estimated_usage, base_model) for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity)) ) -def _deployment_cost_per_token(deployment_id: str | None, model: str) -> CostPerToken | None: - """The pricing the router registered for this deployment under its ``model_info.id``, read - through ``deployment_pricing_model_info`` exactly as the success path does; None means the - public cost map. The id is ``Logging.get_router_model_id()`` (the logging object's own - ``litellm_params``), never the flattened ``model_call_details`` copy a request body key can - overwrite.""" - model_info: Final = deployment_pricing_model_info(model_id=deployment_id, deployment_model=model) - if model_info is None: - return None - return CostPerToken( - input_cost_per_token=model_info["input_cost_per_token"] or 0.0, - output_cost_per_token=model_info["output_cost_per_token"] or 0.0, - ) +def _deployment_priced_input_cost( + model: str, custom_llm_provider: str | None, prompt_tokens: int, deployment_pricing: ModelInfo +) -> float: + """Input price for the estimated prompt at the deployment's registered rates, read through + ``deployment_pricing_model_info`` and priced by ``generic_cost_per_token`` the way the success + path prices them, so a long-context tier the deployment declares (``input_cost_per_token_above_200k_tokens``) + is honoured instead of flattened to the standard rate. An entry it cannot price keeps the row at $0.""" + usage: Final = Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens) + try: + prompt_cost, _ = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider or str(deployment_pricing.get("litellm_provider") or ""), + model_info=deployment_pricing, + ) + except Exception as exc: # noqa: BLE001 # a malformed registered entry must not drop the failure row + verbose_proxy_logger.debug("failure spend estimate: could not price deployment of %s: %s", model, exc) + return 0.0 + return prompt_cost -def _priced_input_cost( - pricing_model: str, - custom_llm_provider: str | None, - prompt_tokens: int, - custom_cost_per_token: CostPerToken | None, -) -> float | None: +def _priced_input_cost(pricing_model: str, custom_llm_provider: str | None, prompt_tokens: int) -> float | None: try: estimated_input_cost, _ = litellm.cost_per_token( model=pricing_model, custom_llm_provider=custom_llm_provider, prompt_tokens=prompt_tokens, - custom_cost_per_token=custom_cost_per_token, ) except Exception as exc: # noqa: BLE001 # cost_per_token raises BadRequestError for an unmapped model and the row must still be written verbose_proxy_logger.debug("failure spend estimate: could not price %s: %s", pricing_model, exc) @@ -1233,18 +1263,20 @@ def _estimated_input_cost( deployment_id: str | None, base_model: str | None, ) -> float: - """Input price for the estimated prompt: the deployment's registered rate when it has one, else - the cost map's rate for ``base_model``, else the deployment model, the same candidate order - ``completion_cost`` walks (an Azure deployment name is not in the cost map, its ``base_model`` - is; a ``base_model`` the cost map does not know falls through to the deployment model). A model - no candidate can price keeps the row at $0 rather than dropping the failure row.""" + """Input price for the estimated prompt: the deployment's registered rates when it has any (the id + is ``Logging.get_router_model_id()``, never the flattened ``model_call_details`` copy a request body + key can overwrite), else the cost map's rate for ``base_model``, else the deployment model, the same + candidate order ``completion_cost`` walks (an Azure deployment name is not in the cost map, its + ``base_model`` is; a ``base_model`` the cost map does not know falls through to the deployment + model). A model no candidate can price keeps the row at $0 rather than dropping the failure row.""" model: Final = str(model_call_details.get("model") or "") custom_llm_provider: Final = str(model_call_details.get("custom_llm_provider") or "") or None - custom_cost_per_token: Final = _deployment_cost_per_token(deployment_id, model) - pricing_models: Final = (base_model, model) if base_model is not None else (model,) + deployment_pricing: Final = deployment_pricing_model_info(model_id=deployment_id, deployment_model=model) + if deployment_pricing is not None: + return _deployment_priced_input_cost(model, custom_llm_provider, prompt_tokens, deployment_pricing) + pricing_models_to_try: Final = pricing_models(base_model, model) candidate_costs: Final = ( - _priced_input_cost(pricing_model, custom_llm_provider, prompt_tokens, custom_cost_per_token) - for pricing_model in pricing_models + _priced_input_cost(pricing_model, custom_llm_provider, prompt_tokens) for pricing_model in pricing_models_to_try ) return next((cost for cost in candidate_costs if cost is not None), 0.0) @@ -1383,8 +1415,9 @@ def _timed_out_reading_from_provider(exception: Exception) -> bool: def _client_disconnected(exception: Exception) -> bool: - """The 499 the proxy raises when the caller hung up while the provider was still working on the request.""" - return getattr(exception, "status_code", None) == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED + """The 499 the proxy itself raises when the caller hung up while the provider was still working on the + request. A provider's own 499 arrives as a litellm exception and is a plain upstream failure.""" + return isinstance(exception, HTTPException) and exception.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED def _provider_billed_the_failure(exception: Exception) -> bool: @@ -1406,17 +1439,19 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: _litellm_params: Final = getattr(_logging_obj, "litellm_params", None) _base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details) _deployment_id: Final = _logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None + _success_already_logged: Final = isinstance(_logging_obj, Logging) and _logging_obj.async_success_log_enqueued + _billed: Final = provider_billed and not _success_already_logged _usage_to_lift: Final = _failure_usage_to_lift( model_call_details=_model_call_details, request_body=request_data, dispatched=_first_handoff is not None, - provider_billed=provider_billed, + provider_billed=_billed, deployment_id=_deployment_id, base_model=_base_model, ) _prompt_estimate: Final = ( _estimated_prompt_usage(_model_call_details, request_data) - if provider_billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") + if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") else None ) _earlier_attempts: Final = ( @@ -3859,6 +3894,7 @@ class ProxyLogging: original_exception, traceback_str, ), + kwargs={"for_async_request": True}, daemon=True, ).start() @@ -4696,19 +4732,25 @@ class ProxyLogging: @staticmethod def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool: - """Drop the parked success dispatch for an assembled chat stream that ends in an error - ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. + """Drop the parked success dispatch for an assembled chat stream or a stream that billed several rounds when it + ends in an error ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. Returns False when the parked dispatch should still be flushed by the caller.""" logging_obj: Final = request_data.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): return False _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None) - assembled: Final = _args[0] if _args else None - if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse): + if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS): return False + match _args: + case (ModelResponse() as assembled, *_): + logging_obj.record_assembled_response_for_failure(assembled) + case (BilledRoundsDispatch() as billed,): + billed.dispatch.close() + logging_obj.record_billed_rounds_for_failure(billed.spend) + case _: + return False logging_obj._on_deferred_stream_complete = None logging_obj._deferred_stream_complete_args = None - logging_obj.record_assembled_response_for_failure(assembled) return True async def arelease_max_parallel_requests_on_disconnect( @@ -4796,6 +4838,7 @@ def jsonify_object(data: dict) -> dict: _deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(max_size=1000) _DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60 _PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5) +_PRISMA_DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2) async def _lookup_deprecated_key( @@ -5202,13 +5245,17 @@ class PrismaClient: return self.db.read_target return self.db - def tx(self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT) -> "TransactionManager": + def tx( + self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT, max_wait: timedelta = _PRISMA_DEFAULT_TX_MAX_WAIT + ) -> "TransactionManager": """Open an interactive transaction on the writer. Callers go through this instead of reaching into ``self.db`` so writer selection and read-replica routing stay encapsulated in the wrapper. """ - return cast("TransactionManager", self.db.tx(timeout=timeout)) # cast-ok: untyped __getattr__ delegate + return cast( # cast-ok: untyped __getattr__ delegate + "TransactionManager", self.db.tx(max_wait=max_wait, timeout=timeout) + ) def get_request_status(self, payload: dict | SpendLogsPayload) -> Literal["success", "failure"]: """ diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py index ad642fb427c..0dda3f93660 100644 --- a/litellm/responses/dispatch.py +++ b/litellm/responses/dispatch.py @@ -1,12 +1,17 @@ import inspect from collections.abc import Awaitable, Callable, Coroutine, Mapping +from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.responses import main from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator -from litellm.rust_bridge.catalog import Route, RouteContext +from litellm.rust_bridge import detectors +from litellm.rust_bridge.catalog import Feature, Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch from litellm.rust_bridge.public_call import ( + AsyncFlag, NativeCall, bind, native_call, @@ -14,6 +19,7 @@ from litellm.rust_bridge.public_call import ( optional_str, signature, ) +from litellm.rust_bridge.responses import detectors as route_detectors from litellm.rust_bridge.responses.entrypoints import ( NATIVE_ARESPONSES, NATIVE_RESPONSES, @@ -43,6 +49,7 @@ def _python_aresponses() -> PythonAresponses: _PYTHON_RESPONSES: Final = _python_responses() _RESPONSES: Final = signature(_PYTHON_RESPONSES) +_FLAG: Final = AsyncFlag("aresponses", sync=_RESPONSES) _PYTHON_ARESPONSES: Final = _python_aresponses() _ARESPONSES: Final = signature(_PYTHON_ARESPONSES) @@ -59,25 +66,64 @@ def _public_request( return native_call(legacy, args, kwargs) -def _context(request: NativeCall) -> RouteContext: +def _provider(request: NativeCall) -> str | None: + provider: Final = optional_str(request.resolved.get("custom_llm_provider")) + try: + return get_llm_provider(str(request.resolved["model"]), provider)[1] + except BadRequestError: + return provider + + +_FEATURES: Final[detectors.Detectors] = MappingProxyType( + { + **detectors.COMMON, + Feature.DROP_PARAMS: detectors.drop_params, + Feature.MODIFY_PARAMS: detectors.modify_params, + Feature.STREAMING: detectors.streaming, + Feature.PROMPT_MANAGEMENT: detectors.prompt_management, + Feature.TEXT_FORMAT: route_detectors.text_format, + Feature.CHAT_COMPLETIONS_BRIDGE: route_detectors.chat_completions_bridge, + Feature.RETRIES_FALLBACKS: detectors.retries_fallbacks, + Feature.DEPLOYMENT_RESPONSE_ID: route_detectors.deployment_response_id, + Feature.ENCODED_ITEM_IDS: route_detectors.encoded_item_ids, + Feature.INPUT_REWRITE: route_detectors.input_rewrite, + Feature.EXTRA_BODY: detectors.extra_body, + Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, + Feature.RESPONSE_CACHE: detectors.response_cache, + Feature.MANAGED_RESOURCE_IDS: route_detectors.managed_resource_ids, + Feature.CACHE_CONTROL_MARKS: route_detectors.cache_control_marks, + Feature.ANTHROPIC_TOOL_SEARCH: route_detectors.anthropic_tool_search, + Feature.TOOL_SCHEMA_REPAIR: route_detectors.tool_schema_repair, + Feature.MODEL_PARAM_GATING: route_detectors.model_param_gating, + Feature.WORKLOAD_IDENTITY: route_detectors.openai_workload_identity, + Feature.CLIENT_HEADERS: route_detectors.client_headers, + } +) + +_SYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType( + {**_FEATURES, Feature.DEPLOYMENT_HOOK: route_detectors.deployment_hook} +) + + +def _context(request: NativeCall, features: detectors.Detectors) -> RouteContext: return RouteContext( Route.RESPONSES, - provider=optional_str(request.resolved.get("custom_llm_provider")), + provider=_provider(request), model=str(request.resolved["model"]), + features=detectors.detected(features, request), ) _DISPATCH: Final = PublicDispatch( route=Route.RESPONSES, request=lambda args, kwargs: _public_request(_RESPONSES, args, kwargs), - context=_context, - bypass=lambda request: request.kwargs.get("aresponses") is True, + context=lambda request: _context(request, _SYNC_FEATURES), ) _ADISPATCH: Final = PublicDispatch( route=Route.RESPONSES, request=lambda args, kwargs: _public_request(_ARESPONSES, args, kwargs), - context=_context, + context=lambda request: _context(request, _FEATURES), ) @@ -86,9 +132,11 @@ def responses( **kwargs: object, # kwargs-ok: preserve the public Responses call shape ) -> ResponsesResult | Coroutine[object, object, ResponsesResult]: python: Final = _PYTHON_RESPONSES + if _FLAG.requested(kwargs): + return aresponses(**_FLAG.keywords(args, kwargs)) return _DISPATCH.run( args, - kwargs, + _FLAG.dropped(kwargs), python=python, binding=NATIVE_RESPONSES, native=native_call_hook, @@ -99,7 +147,7 @@ async def aresponses(*args: object, **kwargs: object) -> ResponsesResult: # kwa python: Final = _PYTHON_ARESPONSES return await _ADISPATCH.arun( args, - kwargs, + _FLAG.dropped(kwargs), python=python, binding=NATIVE_ARESPONSES, native=native_call_hook, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 84d4648f75f..4ef68aa072a 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -5,7 +5,8 @@ from itertools import filterfalse from typing import Any, Final, cast import litellm -from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params +from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.main import stream_chunk_builder from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, @@ -1363,6 +1364,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return False return not any(getattr(part, "text", None) for part in getattr(item, "content", None) or ()) + def finished_round_event(self) -> ResponseCompletedEvent | ResponseIncompleteEvent | None: + if not isinstance(self.litellm_model_response, ModelResponse): + return None + finished: Final = self._emit_terminal_response_event(self.litellm_model_response) + if finished is not None and served_from_cache(self.litellm_custom_stream_wrapper): + finished.response.hidden_params = {**finished.response.hidden_params, "cache_hit": True} + return finished + + def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None: + self.litellm_custom_stream_wrapper.bill_failure_after_earlier_rounds(earlier) + def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: """ The response.incomplete event to bill when the client went away mid-stream. diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 73cfba8dca4..030911238ec 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -3,7 +3,9 @@ from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, @@ -352,6 +354,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._round_max_output_index = -1 self._composed_output: list[object] = [] # mutable-ok: grows as each round finishes self._pending_mcp_call_items: list[dict[str, object]] = [] # mutable-ok: grows per executed tool + self._iteration_started = False + self._round_billing: StreamedRoundBilling | None = None + self._bridged_round_billed: LiteLLMCompletionStreamingIterator | None = None def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" @@ -436,7 +441,15 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: - chunk: Final = await self._anext_impl() + if not self._iteration_started: + self._iteration_started = True + self._round_billing = self._bill_rounds_on_one_row() + try: + chunk: Final = await self._anext_impl() + except Exception: + if self._round_billing is not None: + self._round_billing.finish() + raise sequence_number: Final = getattr(chunk, "sequence_number", None) if isinstance(sequence_number, int): if sequence_number <= self._last_sequence_number and self._last_sequence_number > 0: @@ -446,6 +459,49 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._last_sequence_number = max(self._last_sequence_number, sequence_number) return chunk + def __del__(self) -> None: + round_billing: Final[object] = getattr(self, "_round_billing", None) + if isinstance(round_billing, StreamedRoundBilling): + round_billing.release_abandoned() + + def _bill_rounds_on_one_row(self) -> StreamedRoundBilling | None: + if not self.should_auto_execute: + return None + if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator): + return StreamedRoundBilling(self.base_iterator.litellm_logging_obj) + if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator): + return StreamedRoundBilling(self.base_iterator.logging_obj) + return None + + def _restore_round_headers(self, logged_event: ResponsesAPIStreamingResponse) -> None: + if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator): + return + if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator): + self.base_iterator._restore_provider_response_headers(logged_event) + + def _round_billed_early(self) -> bool: + return self._bridged_round_billed is not None and self._bridged_round_billed is self.base_iterator + + def _bill_finished_bridged_round(self) -> None: + base: Final = self.base_iterator + if self._round_billing is None or not isinstance(base, LiteLLMCompletionStreamingIterator): + return + if self._round_billed_early(): + return + finished: Final = base.finished_round_event() + if finished is None: + return + self._bridged_round_billed = base + round_output: Final = _output_items(finished.response) + self._round_billing.complete_round( + finished, + finished.response, + round_output + if finished.type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE + else (*self._composed_output, *round_output), + self._restore_round_headers, + ) + async def _anext_impl(self) -> ResponsesAPIStreamingResponse: """ Phase-based streaming: @@ -617,6 +673,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): follow-up rounds past the items already emitted, and lists every round's items on the final response.completed. """ + self._bill_finished_bridged_round() chunk_type: Final[object] = getattr(chunk, "type", None) if self._round_index > 0 and chunk_type in ( ResponsesAPIStreamEvents.RESPONSE_CREATED, @@ -630,12 +687,21 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self._output_index_offset: _set_event_field(chunk, "output_index", output_index + self._output_index_offset) + if self.should_auto_execute and chunk_type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE: + return self._charge_incomplete_round(chunk) if not (self.should_auto_execute and self._is_response_completed(chunk)): return chunk response_obj: Final[object] = getattr(chunk, "response", None) if isinstance(response_obj, ResponsesAPIResponse): self.collected_response = response_obj + if self._round_billing is not None and not self._round_billed_early(): + self._round_billing.complete_round( + chunk, + response_obj, + (*self._composed_output, *_output_items(response_obj)), + self._restore_round_headers, + ) # Move to tool execution phase after this chunk self.phase = "tool_execution" await self._generate_tool_execution_events() @@ -645,13 +711,33 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self._follow_up_pending(): self._absorb_round(response_obj) return None - if self._composed_output: - merged_output: Final[list[object]] = [ # mutable-ok: the response model declares output as a list - *self._composed_output, - *_output_items(response_obj), - ] - merged_response: Final = response_obj.model_copy(update={"output": merged_output}) - _set_event_field(chunk, "response", merged_response) + merged_response: Final = ( + response_obj.model_copy(update={"output": [*self._composed_output, *_output_items(response_obj)]}) + if self._composed_output + else response_obj + ) + client_response: Final = ( + merged_response + if self._round_billing is None + else self._round_billing.charge_final_round(chunk, merged_response, self._restore_round_headers) + ) + if client_response is not response_obj: + _set_event_field(chunk, "response", client_response) + return chunk + + def _charge_incomplete_round(self, chunk: ResponsesAPIStreamingResponse) -> ResponsesAPIStreamingResponse: + response_obj: Final[object] = getattr(chunk, "response", None) + if self._round_billing is None or not isinstance(response_obj, ResponsesAPIResponse): + return chunk + if not self._round_billed_early(): + self._round_billing.complete_round( + chunk, response_obj, _output_items(response_obj), self._restore_round_headers + ) + client_response: Final = self._round_billing.charge_final_round( + chunk, response_obj, self._restore_round_headers + ) + if client_response is not response_obj: + _set_event_field(chunk, "response", client_response) return chunk async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: @@ -949,6 +1035,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.collected_response = None self._cached_response_id = None self._round_index += 1 + if self._round_billing is not None: + self._round_billing.record_billed_rounds(follow_up_response) except Exception as e: verbose_logger.error("Error creating follow-up iterator: %s", e) diff --git a/litellm/responses/mcp/round_billing.py b/litellm/responses/mcp/round_billing.py index ae34ab31633..93769fbe82a 100644 --- a/litellm/responses/mcp/round_billing.py +++ b/litellm/responses/mcp/round_billing.py @@ -4,7 +4,7 @@ from contextlib import contextmanager from types import MappingProxyType from typing import Final, TypeVar -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from pydantic_core import to_jsonable_python from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params @@ -46,15 +46,18 @@ def _summed_mapping(first: Mapping[str, object], final: Mapping[str, object]) -> return MappingProxyType({key: _summed_entry(key, first, final) for key in keys}) +def _summed_usage(first: BaseModel, final: BaseModel) -> object: + return to_jsonable_python(_summed_mapping(first.model_dump(), final.model_dump()), fallback=dict) + + def merged_round_usage(first: ResponseAPIUsage | None, final: ResponseAPIUsage | None) -> ResponseAPIUsage | None: if first is None or final is None: return final if first is None else first - return ResponseAPIUsage.model_validate(_summed_mapping(first.model_dump(), final.model_dump())) + return ResponseAPIUsage.model_validate(_summed_usage(first, final)) -def _merged_chat_usage(first: Usage, final: Usage) -> Usage: - summed: Final = _summed_mapping(first.model_dump(), final.model_dump()) - return Usage.model_validate(to_jsonable_python(summed, fallback=dict)) +def merged_chat_usage(first: Usage, final: Usage) -> Usage: + return Usage.model_validate(_summed_usage(first, final)) def _chat_usage(round_response: ResponsesAPIResponse | ModelResponse) -> Usage | None: @@ -82,7 +85,7 @@ def _usage_of_both_rounds( final_usage: Final = _chat_usage(final) if first_usage is None or final_usage is None: return final_usage if first_usage is None else first_usage - return _merged_chat_usage(first_usage, final_usage) + return merged_chat_usage(first_usage, final_usage) def _repeated_fixed_margin( @@ -105,6 +108,14 @@ def _less_repeated_fixed_margin(summed: Mapping[str, object], repeated: float) - return MappingProxyType({**summed, **charged_once}) +def breakdown_less_repeated_fixed_margin( + breakdown: CostBreakdown | None, repeated_fixed_margin: float +) -> CostBreakdown | None: + if breakdown is None: + return None + return _COST_BREAKDOWN.validate_python(_less_repeated_fixed_margin(breakdown, repeated_fixed_margin)) + + def summed_cost_breakdown( first: ResponsesAPIResponse | ModelResponse, first_breakdown: CostBreakdown | None, diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py new file mode 100644 index 00000000000..443147230bf --- /dev/null +++ b/litellm/responses/mcp/stream_round_billing.py @@ -0,0 +1,376 @@ +import asyncio +from collections.abc import Callable, Coroutine, Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Final, Protocol, TypeAlias + +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse +from litellm.types.utils import CostBreakdown, ModelResponse, Usage + +Parked: TypeAlias = tuple[object, ...] +RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None] + +_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class _DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: + return logger + + +def _event_loop_running() -> bool: + try: + asyncio.get_running_loop() + except RuntimeError: + return False + return True + + +def _priced(response: ResponsesAPIResponse, price: UncachedPrice) -> ResponsesAPIResponse: + priced: Final = response.model_copy() + priced.hidden_params = {**response.hidden_params, "response_cost": price.cost} + return priced + + +def _assembled(parked: Parked | None) -> ModelResponse | None: + match parked: + case (ModelResponse() as assembled, bool()): + return assembled + case _: + return None + + +def _bridged_round_usage(parked: Parked | None) -> Usage | None: + usage: Final[object] = getattr(_assembled(parked), "usage", None) + return usage if isinstance(usage, Usage) else None + + +def _usage_of_bridged_rounds(billed: Usage | None, round_usage: Usage | None, cost: object) -> Usage | None: + if billed is None or round_usage is None: + return None + return merged_chat_usage(billed, round_usage).model_copy(update={"cost": cost}) + + +def _logged_cache_hit(response: ResponsesAPIResponse, chat_usage: Usage | None) -> bool | None: + if chat_usage is not None: + return served_from_cache(response) + return served_from_cache(response) or None + + +@dataclass(frozen=True, slots=True) +class _RoundRow: + logged: ResponsesAPIStreamingResponse | ModelResponse + cache_hit: bool | None + end_time: datetime + + +@dataclass(frozen=True, slots=True) +class _SuccessLog: + logger: LiteLLMLoggingObj + row: _RoundRow + chat_usage: Usage | None + + def carry_bridged_usage(self) -> None: + if self.chat_usage is not None: + self.logger.model_call_details["combined_usage_object"] = self.chat_usage + + +def _round_row( + completed: ResponsesAPIStreamingResponse, + response: ResponsesAPIResponse, + chat_usage: Usage | None, + restore_headers: RestoresHeaders, +) -> _RoundRow: + logged_response: Final = response.model_copy() + logged_response.hidden_params = {**response.hidden_params} + logged: Final = completed.model_copy(update={"response": logged_response}) + restore_headers(logged) + return _RoundRow( + logged, + _logged_cache_hit(response, chat_usage), + datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time + ) + + +@dataclass(frozen=True, slots=True) +class _BilledRounds: + response: ResponsesAPIResponse + breakdown: CostBreakdown | None + chat_usage: Usage | None + rounds: int = 1 + + @property + def stamped_cost(self) -> float | None: + cost: Final = self.response.hidden_params.get("response_cost") + return float(cost) if isinstance(cost, (int, float)) else None + + @property + def cost(self) -> float: + return self.stamped_cost or 0.0 + + @property + def spend(self) -> BilledRoundsSpend: + usage: Final = ( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(self.response.usage) + if self.chat_usage is None + else self.chat_usage + ) + return BilledRoundsSpend(usage, self.cost, self.breakdown) + + def with_round( + self, priced: ResponsesAPIResponse, breakdown: CostBreakdown | None, chat_usage: Usage | None + ) -> "_BilledRounds": + merged: Final = billed_for_every_round(self.response, self.breakdown, priced, breakdown) + total: Final = merged.hidden_params.get("response_cost") + charged: Final = merged.model_copy( + update={"usage": None if merged.usage is None else merged.usage.model_copy(update={"cost": total})} + ) + charged.hidden_params = { + **merged.hidden_params, + "cache_hit": served_from_cache(self.response) and served_from_cache(priced), + } + return _BilledRounds( + charged, + summed_cost_breakdown(self.response, self.breakdown, priced, breakdown), + _usage_of_bridged_rounds(self.chat_usage, chat_usage, total), + self.rounds + 1, + ) + + def row( + self, + completed: ResponsesAPIStreamingResponse, + client_output: Sequence[object], + restore_headers: RestoresHeaders, + ) -> _RoundRow: + return _round_row( + completed, + self.response.model_copy(update={"output": list(client_output)}), + self.chat_usage, + restore_headers, + ) + + def bridged_row(self, assembled: ModelResponse | None, logged_id: str) -> _RoundRow | None: + if assembled is None or self.chat_usage is None: + return None + logged: Final = assembled.model_copy(update={"id": logged_id, "usage": self.chat_usage}) + logged.hidden_params = { + **(get_hidden_params(assembled) or {}), + "response_cost": self.stamped_cost, + } + return _RoundRow( + logged, + served_from_cache(self.response), + datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time + ) + + +class StreamedRoundBilling: + """One success row per streamed auto-execute request, billed for every round it ran. + + Each round's iterator parks its success dispatch on the shared logging object rather than sending it, either as a + ready dispatch or, for a stream bridged through chat completions, as the assembled response and its cache hit. + When a round completes, before its tools run, it is folded into the rounds before it and a dispatch for every + round so far stays parked, so a client disconnect or a later round that fails still bills them. Once a later round + starts, that dispatch logs from a copy of the logging state taken then, so the later round's failure, logged on the + request's own state where the proxy reads its failure row, never reaches their success row. Only the latest + dispatch is sent, through the proxy's deferred release when the proxy armed one.""" + + def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: + self._logging_obj: Final = logging_obj + self._billed: _BilledRounds | None = None + self._banked: Parked | None = None + self._banked_row: _RoundRow | None = None + if getattr(logging_obj, _DISPATCH_RELEASE, None) is None: + setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked) + + def complete_round( + self, + completed: ResponsesAPIStreamingResponse, + response: ResponsesAPIResponse, + client_output: Sequence[object], + restore_headers: RestoresHeaders, + ) -> None: + round_parked: Final = self._take_parked() + price: Final = self._price(response, round_parked) + priced: Final = _priced(response, price) + chat_usage: Final = _bridged_round_usage(round_parked) + if self._billed is None: + first_round: Final = _BilledRounds(priced, price.cost_breakdown, chat_usage) + self._billed = first_round + self._bank( + round_parked, + first_round.bridged_row(_assembled(round_parked), priced.id) + or first_round.row(completed, client_output, restore_headers), + ) + return + billed: Final = self._billed.with_round(priced, price.cost_breakdown, chat_usage) + self._close(round_parked) + self._close(self._banked) + self._logging_obj.model_call_details.pop("combined_usage_object", None) + self._logging_obj.model_call_details.pop("response_cost", None) + self._billed = billed + self._logging_obj.cost_breakdown = billed.breakdown + row: Final = billed.bridged_row(_assembled(round_parked), billed.response.id) or billed.row( + completed, client_output, restore_headers + ) + self._bank(None if round_parked is None else (self._dispatch(row, billed, self._logging_obj),), row) + + def record_billed_rounds(self, follow_up: object) -> None: + if self._billed is None or not isinstance(follow_up, BaseResponsesAPIStreamingIterator): + return + follow_up.bill_failure_after_earlier_rounds(self._billed.spend) + self._bank_apart_from_later_rounds(self._billed) + + def charge_final_round( + self, completed: ResponsesAPIStreamingResponse, response: ResponsesAPIResponse, restore_headers: RestoresHeaders + ) -> ResponsesAPIResponse: + parked: Final = self._take_parked() + self._banked = None + if self._billed is None or self._billed.rounds == 1: + if parked is not None: + self._release(parked) + return response + self._close(parked) + charged: Final = response.model_copy(update={"usage": self._billed.response.usage}) + charged.hidden_params = {**self._billed.response.hidden_params} + if parked is not None: + final_row: Final = self._billed.bridged_row(self._latest_assembled(), charged.id) or _round_row( + completed, charged, self._billed.chat_usage, restore_headers + ) + self._release((self._dispatch(final_row, self._billed, self._logging_obj),)) + return charged + + def finish(self) -> None: + parked: Final = self._take_parked() + self._banked = None + if parked is not None: + self._release(parked) + return + self._disarm() + + def release_abandoned(self) -> None: + if not self._owns_release(): + return + if _event_loop_running(): + self.finish() + return + self.abandon() + + def abandon(self) -> None: + self._close(self._take_parked()) + self._close(self._banked) + self._banked = None + self._disarm() + + def _latest_assembled(self) -> ModelResponse | None: + row: Final = self._banked_row + return row.logged if row is not None and isinstance(row.logged, ModelResponse) else None + + def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: + if served_from_cache(response): + return _SERVED_FROM_CACHE + match round_parked: + case (ModelResponse() as assembled, bool()): + return self._logging_obj.uncached_price(assembled) + case _: + return self._logging_obj.uncached_price(response) + + def _bank(self, parked: Parked | None, row: _RoundRow) -> None: + self._banked = parked + self._banked_row = row + if parked is not None: + setattr(self._logging_obj, _PARKED_DISPATCH, parked) + + def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: + row: Final = self._banked_row + if row is None: + return + parked: Final = self._take_parked() + if parked is None: + return + self._close(parked) + self._bank((self._dispatch(row, billed, self._logging_obj.attempt_scoped_copy()),), row) + + def _dispatch( + self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj + ) -> BilledRoundsDispatch | Coroutine[object, object, None]: + dispatch: Final = self._log_success(_SuccessLog(logger, row, billed.chat_usage)) + if billed.chat_usage is None: + return dispatch + return BilledRoundsDispatch(dispatch, billed.spend) + + @staticmethod + async def _log_success(success: _SuccessLog) -> None: + success.carry_bridged_usage() + await _success_dispatcher(success.logger).dispatch_success_handlers( + success.row.logged, + start_time=success.logger.start_time, + end_time=success.row.end_time, + cache_hit=success.row.cache_hit, + prefer_async_handlers=True, + ) + + def _take_parked(self) -> Parked | None: + parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + if not isinstance(parked, tuple): + return None + setattr(self._logging_obj, _PARKED_DISPATCH, None) + return parked + + async def _run_parked(self, *parked: object) -> None: + match parked: + case (BilledRoundsDispatch(dispatch=dispatch),): + await dispatch + case (dispatch,) if asyncio.iscoroutine(dispatch): + await dispatch + case (assembled, cache_hit): + await _success_dispatcher(self._logging_obj).dispatch_success_handlers( + assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + case _: + return + + def _owns_release(self) -> bool: + release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + return release == self._run_parked + + def _release(self, parked: Parked) -> None: + if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, _PARKED_DISPATCH, parked) + return + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + spawn_logging_task( + asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging" + ) + + def _disarm(self) -> None: + if self._owns_release(): + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + + @staticmethod + def _close(parked: Parked | None) -> None: + match parked: + case (BilledRoundsDispatch(dispatch=dispatch),): + dispatch.close() + case (dispatch,) if asyncio.iscoroutine(dispatch): + dispatch.close() + case _: + return diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 3ecfc31f632..636888fd266 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -46,6 +46,7 @@ from litellm.litellm_core_utils.internal_call_metadata import ( 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 from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -367,6 +368,7 @@ class BaseResponsesAPIStreamingIterator: """ _pending_logging_tasks: tuple[asyncio.Task[object], ...] = () + _earlier_rounds_spend: BilledRoundsSpend | None = None def __init__( self, @@ -710,18 +712,22 @@ class BaseResponsesAPIStreamingIterator: if usage_obj is None: return try: - self.logging_obj.model_call_details["combined_usage_object"] = ( - ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj) - ) + failed_usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj) except (TypeError, ValueError) as usage_error: verbose_logger.debug( "could not record usage for failed responses stream: %s", usage_error, ) return - self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj.response_cost_calculator(result=response_obj) or 0.0 - ) + failed_cost: Final = self.logging_obj.response_cost_calculator(result=response_obj) or 0.0 + if self._earlier_rounds_spend is None: + self.logging_obj.model_call_details["combined_usage_object"] = failed_usage + self.logging_obj.model_call_details["response_cost"] = failed_cost + return + self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, failed_usage, failed_cost) + + def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None: + self._earlier_rounds_spend = earlier def _map_error_event_exception(self, error_obj: object) -> Exception: return map_stream_error_to_exception( @@ -996,10 +1002,14 @@ class BaseResponsesAPIStreamingIterator: estimate ) served_tier: Final[object] = getattr(self._created_response, "service_tier", None) - self.logging_obj.record_assembled_response_for_failure( - ModelResponse( - model=self.model, usage=usage, service_tier=served_tier if isinstance(served_tier, str) else 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): diff --git a/litellm/router.py b/litellm/router.py index 1de2330fb1a..15001b5a935 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -38,6 +38,7 @@ from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast @@ -242,6 +243,7 @@ from litellm.router_utils.cooldown_handlers import ( is_video_lookup_not_found, set_cooldown_deployments, ) +from litellm.router_utils.deployment_pricing import deployment_pricing_models from litellm.router_utils.embedding_batch_size_check import ( async_log_rejection_before_provider_call, clear_rejected_before_provider_call, @@ -360,10 +362,11 @@ from litellm.router_utils.routing_groups import ( validate_routing_strategy, ) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch +from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage from litellm.scheduler import FlowItem, Scheduler from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.litellm_params import RoutingStrategyName +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD, RoutingStrategyName from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolParam, @@ -551,6 +554,7 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _CallbackT = TypeVar("_CallbackT") _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) + _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( {"model", "api_base", "api_key", "api_version", "responses_api_version"} ) @@ -1045,6 +1049,8 @@ class RoutingArgs(enum.Enum): # as the per-request one built from a caller-supplied user_config, drops out on its # own rather than leaving entries behind that nothing can withdraw. _live_routers: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() +_routers_holding_global_callbacks: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() +_routers_holding_global_callbacks_lock: Final = threading.Lock() def _replay_live_router_model_cost() -> None: @@ -1577,22 +1583,18 @@ class Router: self._init_routing_groups(self._routing_groups_input) self._override_selectors: dict[str, RouterStrategySelector | None] = {} self._override_selectors_lock = threading.Lock() + self._request_scoped_selectors: tuple[CustomLogger, ...] = () + self._request_scoped_policies: tuple[CustomLogger, ...] = () self.access_groups = None + with _routers_holding_global_callbacks_lock: + _routers_holding_global_callbacks.add(self) ## USAGE TRACKING ## - if isinstance(litellm._async_success_callback, list): - litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) - else: - litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) + litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) if isinstance(litellm.success_callback, list): litellm.logging_callback_manager.add_litellm_success_callback(self.sync_deployment_callback_on_success) else: litellm.success_callback = [self.sync_deployment_callback_on_success] - if isinstance(litellm._async_failure_callback, list): - litellm.logging_callback_manager.add_litellm_async_failure_callback( - self.async_deployment_callback_on_failure - ) - else: - litellm._async_failure_callback = [self.async_deployment_callback_on_failure] + litellm.logging_callback_manager.add_litellm_async_failure_callback(self.async_deployment_callback_on_failure) ## COOLDOWNS ## if isinstance(litellm.failure_callback, list): litellm.logging_callback_manager.add_litellm_failure_callback(self.deployment_callback_on_failure) @@ -1688,20 +1690,116 @@ class Router: # Stop contributing to cost-map rebuilds straight away rather than waiting # for this router to be collected. _live_routers.discard(self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_success_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.success_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_failure_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.failure_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.input_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.service_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, self) + self.retire_routing_state() + self._release_global_callbacks() - # Remove ForwardClientSideHeadersByModelGroup if it exists - if self.optional_callbacks is not None: - for callback in self.optional_callbacks: - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.callbacks, callback, require_self=False - ) + def _release_global_callbacks(self) -> None: + with _routers_holding_global_callbacks_lock: + _routers_holding_global_callbacks.discard(self) + slack_alerting_loggers: Final = self._slack_alerting_loggers() + adaptive_router_hooks: Final = self._adaptive_router_hooks() + vacated_keys: Final = litellm.logging_callback_manager.drop_callbacks( + ( + *self._callback_selectors(), + *self._override_selectors.values(), + *(self.optional_callbacks or ()), + *slack_alerting_loggers, + *adaptive_router_hooks, + ), + method_owners=(self, *slack_alerting_loggers), + ) + with _routers_holding_global_callbacks_lock: + survivors: Final = tuple(_routers_holding_global_callbacks) + candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors) + for candidate in litellm.logging_callback_manager.with_custom_logger_keys(candidates, vacated_keys): + self._register_router_selector(candidate) + adaptive_keys: Final = frozenset( + map(litellm.logging_callback_manager.get_custom_logger_key, adaptive_router_hooks) + ) + adaptive_survivor: Final = next((router for router in survivors if router.adaptive_routers), None) + if adaptive_survivor is not None and vacated_keys & adaptive_keys: + adaptive_survivor._sync_adaptive_router_hooks() + + def serve_through_request_logging(self) -> None: + self._request_scoped_selectors = (*self._callback_selectors(), *self._request_adaptive_router_hooks()) + self._request_scoped_policies = (*self._optional_loggers(), *self._slack_alerting_loggers()) + self._release_global_callbacks() + + def bind_to_request_logging(self, logging_obj: LiteLLMLogging) -> None: + logging_obj.add_dynamic_function_callbacks( + async_success=self.deployment_callback_on_success, + failure=self.deployment_callback_on_failure, + async_failure=self.async_deployment_callback_on_failure, + ) + for selector in self._request_scoped_selectors: + logging_obj.replace_global_callback(selector) + for policy in self._request_policies_not_set_globally(): + logging_obj.add_dynamic_callback(policy) + + def _request_policies_not_set_globally(self) -> tuple[CustomLogger, ...]: + global_keys: Final = frozenset( + map( + litellm.logging_callback_manager.get_custom_logger_key, + (callback for callback in litellm.callbacks if isinstance(callback, CustomLogger)), + ) + ) + return tuple( + policy + for policy in self._request_scoped_policies + if litellm.logging_callback_manager.get_custom_logger_key(policy) not in global_keys + ) + + def _request_adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + return tuple( + AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) + for tagged in chain.from_iterable(self.adaptive_routers.values()) + ) + + def _routing_callbacks(self) -> Iterable[object]: + if not (self._request_scoped_selectors or self._request_scoped_policies): + return litellm.callbacks + replaced_keys: Final = frozenset( + map(litellm.logging_callback_manager.get_custom_logger_key, self._request_scoped_selectors) + ) + return chain( + litellm.logging_callback_manager.without_custom_logger_keys(litellm.callbacks, replaced_keys), + self._request_scoped_selectors, + self._request_policies_not_set_globally(), + ) + + def retire_routing_state(self) -> None: + for strategy in (*self._global_callback_candidates(), *self._override_selectors.values()): + if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting): + strategy.retire() + + def _slack_alerting_loggers(self) -> tuple[CustomLogger, ...]: + slack_alerting_logger: Final = getattr(self, "slack_alerting_logger", None) + return (slack_alerting_logger,) if isinstance(slack_alerting_logger, CustomLogger) else () + + def _adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + own_strategy_ids: Final = frozenset( + id(tagged.strategy) for tagged in chain.from_iterable(self.adaptive_routers.values()) + ) + hooks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + return tuple( + hook + for hook in hooks + if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own_strategy_ids + ) + + def _optional_loggers(self) -> tuple[CustomLogger, ...]: + return tuple(callback for callback in self.optional_callbacks or () if isinstance(callback, CustomLogger)) + + def _global_callback_candidates(self) -> tuple[CustomLogger, ...]: + return (*self._callback_selectors(), *self._optional_loggers(), *self._slack_alerting_loggers()) @staticmethod def _create_redis_cache( @@ -1807,7 +1905,7 @@ class Router: return selector @staticmethod - def _register_router_selector(selector: RouterStrategySelector) -> None: + def _register_router_selector(selector: CustomLogger) -> None: if isinstance(selector, LeastBusyLoggingHandler): if isinstance(litellm.input_callback, list): litellm.logging_callback_manager.add_litellm_input_callback(selector) @@ -1816,23 +1914,28 @@ class Router: if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(selector) + def _callback_selectors(self) -> tuple[RouterStrategySelector, ...]: + return tuple( + selector + for selector in ( + *(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()), + *chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()), + ) + if selector is not None + ) + def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: """ Drop router-owned strategy selectors from litellm's global callback - lists by identity. Used before re-init (`routing_strategy_init` / - `_init_routing_groups`) so repeated `update_settings` calls don't - accumulate dead selectors that keep receiving callback events. + lists. Used before re-init (`routing_strategy_init` / + `_init_routing_groups`) and by `discard` so dead selectors stop + receiving callback events. `function_setup` copies `litellm.callbacks` + into the per-event lists on every request, so those are purged too. """ for selector in selectors: if isinstance(selector, BaseRoutingStrategy): selector.retire() - selector_ids: Final = {id(s) for s in selectors if s is not None} - if not selector_ids: - return - if isinstance(litellm.callbacks, list): - litellm.callbacks = [c for c in litellm.callbacks if id(c) not in selector_ids] - if isinstance(litellm.input_callback, list): - litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids] + litellm.logging_callback_manager.drop_callbacks(selectors) def _apply_updated_routing_strategy_args(self) -> None: """ @@ -2524,6 +2627,11 @@ class Router: self.asystemone = self.factory_function(asystemone, call_type="asystemone") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable self.systemone = self.factory_function(systemone, call_type="systemone") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable + from litellm.extraction import aextraction, extraction + + self.aextraction = self.factory_function(aextraction, call_type="aextraction") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable + self.extraction = self.factory_function(extraction, call_type="extraction") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable + def _initialize_video_endpoints(self): """Initialize video endpoints.""" from litellm.videos import ( @@ -4667,7 +4775,8 @@ class Router: deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) ): kwargs.pop(key, None) - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=self.default_litellm_params) + defaults: Final = _LITELLM_PARAMS_ADAPTER.validate_python(self.default_litellm_params) + self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=defaults) model_info = deployment.get("model_info", {}).copy() deployment_litellm_model_name = deployment["litellm_params"]["model"] @@ -4811,6 +4920,19 @@ class Router: data=_DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]), ) + deployment_litellm_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) + if router_config_requests_stream_usage( + _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + ( + _LITELLM_PARAMS_ADAPTER.validate_python( + deployment_litellm_params + if isinstance(deployment_litellm_params, Mapping) + else deployment_litellm_params.model_dump(exclude_none=True) + ), + defaults, + ), + ): + kwargs[STREAM_USAGE_INJECTED_FIELD] = True self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) if _is_fallback_hop(hop_kwargs): @@ -7707,6 +7829,8 @@ class Router: "decisions", "asystemone", "systemone", + "aextraction", + "extraction", "aadapter_generate_content", "avideo_generation", "video_generation", @@ -7782,6 +7906,7 @@ class Router: "search", "decisions", "systemone", + "extraction", "video_generation", "video_list", "video_status", @@ -7951,6 +8076,7 @@ class Router: "ocr", "adecisions", "asystemone", + "aextraction", "avideo_generation", "avideo_list", "avideo_status", @@ -9500,17 +9626,24 @@ class Router: When a retry or fallback happens, record which model group, deployment and attempt just failed and why, and count it toward the request-wide num_retries_per_request cap """ + from litellm.litellm_core_utils.get_litellm_params import get_base_model_from_litellm_call_metadata from litellm.types.router import RetryAttemptRecord _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var] + request_metadata: Final = cast("Mapping[str, object]", kwargs[_metadata_var]) # cast-ok: untyped kwargs model_group: Final = kwargs.get("model") model_info: Final = request_metadata.get("model_info") deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None + base_model: Final = get_base_model_from_litellm_call_metadata( + metadata=request_metadata + ) or self._deployment_params_base_model(deployment_id if isinstance(deployment_id, str) else None) + deployment_model: Final = request_metadata.get("deployment") attempted_retries: Final = request_metadata.get("attempted_retries") attempt_record: Final[RetryAttemptRecord] = { "model_group": model_group if isinstance(model_group, str) else None, "deployment_id": deployment_id if isinstance(deployment_id, str) else None, + "model": deployment_model if isinstance(deployment_model, str) else None, + "base_model": base_model if isinstance(base_model, str) else None, "exception_type": type(e).__name__, "exception_string": str(e), "attempted_retries": attempted_retries if type(attempted_retries) is int else None, @@ -9541,6 +9674,18 @@ class Router: kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap return kwargs + def _deployment_params_base_model(self, deployment_id: str | None) -> str | None: + """The ``base_model`` a deployment declares under ``litellm_params`` (the Azure named-deployment + shape), the same key the success path's ``get_base_model_from_metadata`` reads.""" + deployment: Final = self.get_deployment(model_id=deployment_id) if deployment_id is not None else None + if deployment is None: + return None + extra_params: Final = cast( # cast-ok: pydantic extras are untyped, base_model is isinstance-checked below + "Mapping[str, object]", deployment.litellm_params.model_extra or {} + ) + base_model: Final = extra_params.get("base_model") + return base_model if isinstance(base_model, str) else None + def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int: """ Update deployment rpm for that minute @@ -9763,7 +9908,7 @@ class Router: Raises: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): _callback.pre_call_check(deployment) @@ -9784,7 +9929,7 @@ class Router: Raises: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): try: await _callback.async_pre_call_check(deployment, parent_otel_span) @@ -9921,7 +10066,7 @@ class Router: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ returned_healthy_deployments = healthy_deployments - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): try: returned_healthy_deployments = await _callback.async_filter_deployments( @@ -10576,6 +10721,7 @@ class Router: changes, otherwise a released router keeps recording turns through its hook.""" from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, ) for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): @@ -10585,6 +10731,8 @@ class Router: litellm.logging_callback_manager.add_litellm_callback( AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) ) + if self.adaptive_routers: + litellm.logging_callback_manager.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) def init_adaptive_router_deployment(self, deployment: Deployment) -> None: """ @@ -11690,19 +11838,9 @@ class Router: ) for deployment in deployments ) - params: Final = tuple(deployment.get("litellm_params") or MappingProxyType({}) for deployment in deployments) - # base_model resolution mirrors get_router_model_info: unset or blank means the - # deployment's own model name is the cost-map key. - cost_map_keys: Final = tuple( - dict.fromkeys( # deduplicates while preserving config order - key - for key in ( - model_info.get("base_model") or litellm_params.get("base_model") or litellm_params.get("model") - for model_info, litellm_params in zip(model_infos, params) - ) - if isinstance(key, str) and key - ) - ) + raw_deployments: Final[tuple[object, ...]] = deployments + pricing_keys: Final = (deployment_pricing_models(deployment) for deployment in raw_deployments) + cost_map_keys: Final = tuple(dict.fromkeys(keys[0] for keys in pricing_keys if keys)) return DeploymentModelListingInfo( cost_map_keys=cost_map_keys, max_input_tokens=self._widest_configured_limit(model_infos, "max_input_tokens"), diff --git a/litellm/router_strategy/adaptive_router/__init__.py b/litellm/router_strategy/adaptive_router/__init__.py index d7f55ebced9..1f7e6b93169 100644 --- a/litellm/router_strategy/adaptive_router/__init__.py +++ b/litellm/router_strategy/adaptive_router/__init__.py @@ -1,6 +1,9 @@ """Adaptive router strategy. See README.md for design overview.""" from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter -from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, +) -__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"] +__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook", "AdaptiveRouterResponseHeaderHook"] diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 92207e07c52..3427b81be68 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -170,12 +170,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequen return content, tool_calls -class AdaptiveRouterPostCallHook(CustomLogger): - """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" - - def __init__(self, adaptive_router: AdaptiveRouter) -> None: - self.adaptive_router = adaptive_router - +class AdaptiveRouterResponseHeaderHook(CustomLogger): async def async_post_call_response_headers_hook( self, data: Mapping[str, object], @@ -200,6 +195,13 @@ class AdaptiveRouterPostCallHook(CustomLogger): return None return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} + +class AdaptiveRouterPostCallHook(CustomLogger): + """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" + + def __init__(self, adaptive_router: AdaptiveRouter) -> None: + self.adaptive_router = adaptive_router + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): await self._record(kwargs, response_obj, response_status=200) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 51400e6da24..2cea7389808 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -24,6 +24,8 @@ class BaseRoutingStrategy(ABC): self.dual_cache = dual_cache self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] self._sync_task: asyncio.Task[None] | None = None + self._retired = False + self._flush_after_retire: asyncio.Task[None] | None = None if should_batch_redis_writes: self.setup_sync_task(default_sync_interval) @@ -47,13 +49,25 @@ class BaseRoutingStrategy(ABC): def retire(self) -> None: self.cancel_sync_task() - if not self.redis_increment_operation_queue: + self._retired = True + self._schedule_flush_after_retire() + + def _schedule_flush_after_retire(self) -> None: + if self._flush_after_retire is not None and not self._flush_after_retire.done(): + return + if not self.redis_increment_operation_queue or self.dual_cache.redis_cache is None: return try: loop: Final = asyncio.get_running_loop() except RuntimeError: return - loop.create_task(self._push_in_memory_increments_to_redis()) + self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) + + async def _drain_increments_after_retire(self) -> None: + if not self.redis_increment_operation_queue: + return + await self._push_in_memory_increments_to_redis() + await self._drain_increments_after_retire() async def cleanup(self): """Cleanup method to be called when shutting down""" @@ -98,6 +112,8 @@ class BaseRoutingStrategy(ABC): self.redis_increment_operation_queue.append(increment_op) self.add_to_in_memory_keys_to_update(key=key) + if self._retired: + self._schedule_flush_after_retire() return result @with_service_target("router_usage") diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 147ce60a99b..eab232eae18 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -128,7 +128,9 @@ class RouterBudgetLimiting(CustomLogger): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None - asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._sync_task: asyncio.Task[None] = asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._retired = False + self._flush_after_retire: asyncio.Task[None] | None = None self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config self.deployment_budget_config: GenericBudgetConfigType | None = None self.tag_budget_config: GenericBudgetConfigType | None = None @@ -140,6 +142,30 @@ class RouterBudgetLimiting(CustomLogger): if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(self) + def retire(self) -> None: + self._sync_task.cancel() + self._retired = True + self._schedule_flush_after_retire() + + def _schedule_flush_after_retire(self) -> None: + if self._flush_after_retire is not None and not self._flush_after_retire.done(): + return + if not self._has_unwritten_increments(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) + + def _has_unwritten_increments(self) -> bool: + return bool(self.redis_increment_operation_queue) or self._detached_increment_operations is not None + + async def _drain_increments_after_retire(self) -> None: + while self._has_unwritten_increments(): + if not await self._push_in_memory_increments_to_redis(): + await asyncio.sleep(DEFAULT_REDIS_SYNC_INTERVAL) + @with_service_target("router_budgets") async def async_filter_deployments( self, @@ -427,6 +453,8 @@ class RouterBudgetLimiting(CustomLogger): if self.dual_cache.redis_cache is None: return self.redis_increment_operation_queue.append(increment_op) + if self._retired: + self._schedule_flush_after_retire() def _get_redis_increment_queue_lock(self) -> asyncio.Lock: return self._redis_increment_queue_lock diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index d567b6acccc..5894da9cb4f 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -1,5 +1,6 @@ #### What this does #### # picks based on response time (for streaming, this is time to first token) +from collections.abc import Mapping from datetime import datetime from typing import Final @@ -10,6 +11,30 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed + + +def _deployment_cost_map_rates(deployment: object) -> tuple[float | None, float | None]: + """Input and output per-token prices from the first cost-map row among the deployment's + pricing models that carries one; (None, None) when no row does, so the caller falls back to + its unmapped default.""" + rows: Final = (str_keyed(_cost_map().get(model)) for model in deployment_pricing_models(deployment)) + row: Final = next((row for row in rows if _declares_token_price(row)), None) + if row is None: + return None, None + return _as_rate(row.get("input_cost_per_token")), _as_rate(row.get("output_cost_per_token")) + + +def _cost_map() -> Mapping[str, object]: + return litellm.model_cost + + +def _declares_token_price(row: Mapping[str, object]) -> bool: + return row.get("input_cost_per_token") is not None or row.get("output_cost_per_token") is not None + + +def _as_rate(value: object) -> float | None: + return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None class LowestCostLoggingHandler(CustomLogger): @@ -243,8 +268,8 @@ class LowestCostLoggingHandler(CustomLogger): or _deployment.get("model_info", {}).get("rpm", None) or float("inf") ) - item_litellm_model_name = _deployment.get("litellm_params", {}).get("model") - item_litellm_model_cost_map = litellm.model_cost.get(item_litellm_model_name, {}) + deployment_view: object = _deployment + map_input_cost, map_output_cost = _deployment_cost_map_rates(deployment_view) # check if user provided input_cost_per_token and output_cost_per_token in litellm_params item_input_cost = None @@ -256,10 +281,10 @@ class LowestCostLoggingHandler(CustomLogger): item_output_cost = _deployment.get("litellm_params", {}).get("output_cost_per_token") if item_input_cost is None: - item_input_cost = item_litellm_model_cost_map.get("input_cost_per_token", 5.0) + item_input_cost = map_input_cost if map_input_cost is not None else 5.0 if item_output_cost is None: - item_output_cost = item_litellm_model_cost_map.get("output_cost_per_token", 5.0) + item_output_cost = map_output_cost if map_output_cost is not None else 5.0 # if litellm["model"] is not in model_cost map -> use item_cost = $10 diff --git a/litellm/router_utils/deployment_pricing.py b/litellm/router_utils/deployment_pricing.py new file mode 100644 index 00000000000..622cc91ab5e --- /dev/null +++ b/litellm/router_utils/deployment_pricing.py @@ -0,0 +1,37 @@ +from collections.abc import Iterator, Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +_EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) +_STR_KEYED: Final = TypeAdapter(dict[str, object]) + + +def pricing_models(base_model: object, model: object) -> tuple[str, ...]: + """Cost-map keys to price a deployment under, ``model_info.base_model`` first, then the + deployment's own model, blanks and duplicates dropped. Every spend and routing surface that + prices a deployment walks this order, so a ``base_model`` the cost map does not know falls + through to the deployment model instead of billing or ranking at the unmapped default.""" + return tuple(dict.fromkeys(_named_models(base_model, model))) + + +def deployment_pricing_models(deployment: object) -> tuple[str, ...]: + """``pricing_models`` read off a raw deployment dict (``model_info.base_model``, then + ``litellm_params.base_model``, then ``litellm_params.model``); malformed blocks count as unset.""" + view: Final = str_keyed(deployment) + params: Final = str_keyed(view.get("litellm_params")) + info: Final = str_keyed(view.get("model_info")) + return pricing_models(info.get("base_model") or params.get("base_model"), params.get("model")) + + +def _named_models(*candidates: object) -> Iterator[str]: + return (candidate for candidate in candidates if isinstance(candidate, str) and candidate) + + +def str_keyed(value: object) -> Mapping[str, object]: + """``value`` as a read-only str-keyed mapping, empty when it is not a dict.""" + try: + return MappingProxyType(_STR_KEYED.validate_python(value)) + except ValidationError: + return _EMPTY diff --git a/litellm/router_utils/stream_usage_injection.py b/litellm/router_utils/stream_usage_injection.py new file mode 100644 index 00000000000..946543d9b83 --- /dev/null +++ b/litellm/router_utils/stream_usage_injection.py @@ -0,0 +1,17 @@ +from collections.abc import Mapping, Sequence + + +def router_config_requests_stream_usage( + request_kwargs: Mapping[str, object], configured_params: Sequence[Mapping[str, object]] +) -> bool: + if request_kwargs.get("stream") is not True or "stream_options" in request_kwargs: + return False + return any(_includes_usage(params.get("stream_options")) for params in configured_params) + + +def _includes_usage(stream_options: object) -> bool: + match stream_options: + case {"include_usage": True}: + return True + case _: + return False diff --git a/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index 607bb257a0e..c197b21bad0 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -1,15 +1,53 @@ -# Python boundary +# Ownership -This package owns native rollout and fallback selection, Python public API compatibility, settings projection and the Python bindings supplied to the Rust bridge. Rust core owns provider execution. `litellm-host-python` owns CPython runtime mechanics. `callbacks-legacy-python` owns legacy callback sharing and dispatch policy +- This package owns + - choosing Rust or Python for each call, and falling back to Python when Rust is unavailable + - Python public API compatibility at the native boundary + - projecting Python settings into the values Rust reads + - the Python bindings handed to the Rust bridge +- It does not own + - provider execution, which is Rust core + - CPython runtime mechanics, which is `litellm-host-python` + - legacy callback sharing and dispatch policy, which is `callbacks-legacy-python` +- `host/` holds the modules Rust imports by name to read Python-owned state. `host/AGENTS.md` states that contract -`host/` holds the Python modules Rust imports by name to read Python-owned state (provider resolution, model capabilities, settings, caches, diagnostics, legacy callbacks, the secret manager, transport and timeout conversions). Its AGENTS.md states that contract +# Routing -`execution.py` owns generic inline execution and stream iteration, plus the step classes Rust constructs. Native entrypoints return an unstarted `Execution`, and `execution.synchronous` or `execution.asynchronous` hands it to the public caller through `run` or `drive` with the one stream pair, so Rust never names a Python driver or stream class. Keep one driver implementing generator-style `send` (`send(None)` starts), `throw` and idempotent `close`. Do not create a second implementation +- `catalog.py` decides, the detectors describe the request, `dispatch.py` applies the decision +- Shared detectors live in `detectors.py`; route-only detectors live in `/detectors.py` +- The first matching rule wins. A context no rule matches stays on Python +- Every Python-only exclusion for a route (a `targets` or `features` rule) sits above any Rust rule for that route +- A route whose Python side is `NO_PYTHON` resolves to `RUST_REQUIRED` for every context it can produce + - `tests/test_litellm_rust/test_public_entries.py` marks such a route `python=None`, so dropping a route's Python changes its entries, its rules and that field together +- A detector reads the same inputs the legacy Python path reads, so the catalog and the fallback agree on which requests need Python + - A route's context, its target as well as its features, reads `NativeCall.resolved`, so an argument counts the same whether it came by position, by keyword or as a signature default + - A route's sync and async dispatches read different detector maps when the legacy handlers differ: Messages runs pre-request hooks, web search interception and interceptors only on async calls, and the sync Responses handler runs the deployment hook itself +- A detector that raises keeps the call on Python. `PublicDispatch` logs it and never fails the call over routing +- `catalog.RULES` is rollout policy with no tests of its own. Tests pass the rules they need, so check these invariants by review whenever `RULES` changes -Generic execution steps and inline suspension handling must not depend on LiteLLM response metadata. Public stream construction, `_hidden_params` and header compatibility remain product responsibilities. Keep generic stream heads opaque and preserve public header behavior in the product wrappers +# Public entries -Await each selected suspension in the caller's task so hooks retain thread, loop and context identity. A final awaitable value is returned as a value, not awaited implicitly. Closing, cancellation and `GeneratorExit` release the execution without replaying provider work or starting further callbacks +- Every route serves `a(...)`, `(..., a=True)` and `a(..., a=True)` the same way, handing native one identical `NativeCall` with no async flag on the async binding +- `tests/test_litellm_rust/test_public_entries.py` holds this for every route in its table, so a new public route gets a row there -Native admission may select legacy before execution starts. A failure after execution starts is terminal and never authorizes fallback. Creating an unstarted native coroutine must not acquire clients or credentials or capture execution context +# Execution -Verify cross-language behavior using a fresh installed extension, including unstarted-call release, setup ordering, exception identity, sync and async streams, cancellation and GC. Keep `_native.pyi` accurate about Future-returning and coroutine-returning APIs +- `execution.py` owns generic inline execution, stream iteration and the step classes Rust constructs +- Native entrypoints return an unstarted `Execution`. `execution.synchronous` or `execution.asynchronous` hands it to the caller through `run` or `drive`, so Rust never names a Python driver or stream class +- Keep one driver, implementing generator-style `send` (`send(None)` starts), `throw` and idempotent `close` +- Generic steps and suspension handling never depend on LiteLLM response metadata. Public stream construction, `_hidden_params` and header compatibility stay in the product wrappers, and generic stream heads stay opaque +- Await each selected suspension in the caller's task, so hooks keep their thread, loop and context +- Return a final awaitable as a value instead of awaiting it implicitly +- Closing, cancellation and `GeneratorExit` release the execution without replaying provider work or starting more callbacks + +# Fallback + +- Native admission may pick Python before execution starts +- 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 + +# 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 +- Keep `_native.pyi` accurate about which APIs return a Future and which return a coroutine +- Unit tests of this package live in `tests/unit/rust_bridge/`, end-to-end Rust path tests in `tests/test_litellm_rust/` diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index ca271e201b0..ebb221e98be 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -25,11 +25,94 @@ class Route(str, Enum): TOKENIZER = "tokenizer" +class Target(str, Enum): + CHAT_COMPLETIONS = "chat_completions" + MESSAGES = "messages" + RESPONSES = "responses" + + +class Feature(str, Enum): + SAFEGUARDS_CLASSIFIER = "safeguards_classifier" + """A non-Claude model with dangerous_tool_use safeguards, which the proxy classifier answers.""" + MOCK_RESPONSE = "mock_response" + """mock_response short-circuits before any provider call.""" + PROXY_MCP_TOOLS = "proxy_mcp_tools" + """litellm_proxy MCP tools that Python expands and may auto-execute.""" + CALLER_CLIENT = "caller_client" + """A caller-supplied HTTP client or shared session that Python sends through.""" + TIMEOUT_VARIANTS = "timeout_variants" + """request_timeout, stream_timeout or httpx.Timeout phases beyond a plain timeout.""" + DROP_PARAMS = "drop_params" + """drop_params or additional_drop_params dropping fields per provider.""" + MODIFY_PARAMS = "modify_params" + """litellm.modify_params rewriting the request, like clamping max_tokens.""" + STREAMING = "streaming" + """stream=True.""" + PROMPT_MANAGEMENT = "prompt_management" + """prompt_id resolved through a prompt manager.""" + TEXT_FORMAT = "text_format" + """A text_format model that Python turns into a JSON schema and parses back.""" + CHAT_COMPLETIONS_BRIDGE = "chat_completions_bridge" + """use_chat_completions_api or the openai/chat_completions/ prefix, served through Chat Completions.""" + RETRIES_FALLBACKS = "retries_fallbacks" + """num_retries, max_retries or fallbacks run by the SDK.""" + DEPLOYMENT_RESPONSE_ID = "deployment_response_id" + """A deployment id that Python encodes into the returned response id.""" + ENCODED_ITEM_IDS = "encoded_item_ids" + """Input items carrying litellm-encoded ids that Python restores before sending.""" + INPUT_REWRITE = "input_rewrite" + """Responses input items that Python normalizes or removes before replay.""" + EXTRA_BODY = "extra_body" + """Extra body fields that Python transforms and merges into the provider request.""" + PROVIDER_AFFINITY = "provider_affinity" + """A provider affinity header that Python validates and populates from the session id.""" + PRE_REQUEST_HOOK = "pre_request_hook" + """A logger overriding async_pre_request_hook, which may rewrite the request.""" + WEBSEARCH_INTERCEPTION = "websearch_interception" + """A WebSearchInterceptionLogger that may answer without calling the provider.""" + MESSAGES_INTERCEPTOR = "messages_interceptor" + """A registered MessagesInterceptor that takes over the call.""" + AGENTIC_LOOP_HOOK = "agentic_loop_hook" + """A logger overriding async_should_run_agentic_loop, which may run follow-up calls.""" + RESPONSE_CACHE = "response_cache" + """litellm.cache, whose hit and store order differs from native.""" + TLS_POLICY = "tls_policy" + """An OpenSSL security level or ECDH curve that native TLS cannot apply.""" + MANAGED_RESOURCE_IDS = "managed_resource_ids" + """Managed file or vector store ids that Python maps to provider ids.""" + CACHE_CONTROL_MARKS = "cache_control_marks" + """Anthropic cache_control marks that Python strips or turns into prompt cache breakpoints.""" + ANTHROPIC_TOOL_SEARCH = "anthropic_tool_search" + """Anthropic tool search tools that Python maps to OpenAI tool search.""" + TOOL_SCHEMA_REPAIR = "tool_schema_repair" + """Tool schemas that Python decodes, flattens or strips of non-Python regex patterns.""" + MODEL_PARAM_GATING = "model_param_gating" + """Params that Python clamps, drops or rejects from the model map, like gpt-5 temperature.""" + WORKLOAD_IDENTITY = "workload_identity" + """No static key, so Python mints a workload identity federation token.""" + CACHE_CONTROL_INJECTION = "cache_control_injection" + """Prompt caching injection that Python applies to messages and system.""" + STREAM_COST_INJECTION = "stream_cost_injection" + """litellm.include_cost_in_streaming_usage, where Python adds cost to streamed usage.""" + THINKING_SIGNATURE_RETRY = "thinking_signature_retry" + """Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature.""" + CLIENT_HEADERS = "client_headers" + """Proxy-forwarded headers that Python merges into extra_headers.""" + DEPLOYMENT_HOOK = "deployment_hook" + """A logger overriding async_pre_call_deployment_hook, which the sync Python handler runs itself.""" + MODEL_LIST = "model_list" + """A model_list that Python fans out across its deployments.""" + MODEL_ALIAS_MAP = "model_alias_map" + """A model that litellm.model_alias_map rewrites before the call.""" + + @dataclass(frozen=True, slots=True) class RouteContext: route: Route provider: str | None = None model: str | None = None + target: Target | None = None + features: frozenset[Feature] = frozenset() @dataclass(frozen=True, slots=True) @@ -38,23 +121,96 @@ class RouteRule: rollout: Rollout providers: frozenset[str] | None = None models: frozenset[str] | None = None + targets: frozenset[Target] | None = None + features: frozenset[Feature] | None = None + + def unconstrained(self) -> bool: + return self.providers is None and self.models is None and self.targets is None and self.features is None def matches(self, context: RouteContext) -> bool: return ( context.route is self.route and (self.providers is None or context.provider in self.providers) and (self.models is None or context.model in self.models) + and (self.targets is None or context.target in self.targets) + and (self.features is None or not self.features.isdisjoint(context.features)) ) Rules: TypeAlias = tuple[RouteRule, ...] +COMMON_PYTHON_FEATURES: Final = frozenset( + { + Feature.MOCK_RESPONSE, + Feature.PROXY_MCP_TOOLS, + Feature.CALLER_CLIENT, + Feature.TIMEOUT_VARIANTS, + Feature.AGENTIC_LOOP_HOOK, + Feature.TLS_POLICY, + } +) + +CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { + Feature.RESPONSE_CACHE, + Feature.STREAMING, + Feature.DROP_PARAMS, + Feature.MODIFY_PARAMS, + Feature.EXTRA_BODY, + Feature.RETRIES_FALLBACKS, + Feature.PROMPT_MANAGEMENT, + Feature.CACHE_CONTROL_INJECTION, + Feature.MODEL_PARAM_GATING, + Feature.MODEL_LIST, + Feature.MODEL_ALIAS_MAP, +} + +MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { + Feature.SAFEGUARDS_CLASSIFIER, + Feature.PRE_REQUEST_HOOK, + Feature.WEBSEARCH_INTERCEPTION, + Feature.MESSAGES_INTERCEPTOR, + Feature.RESPONSE_CACHE, + Feature.MODIFY_PARAMS, + Feature.WORKLOAD_IDENTITY, + Feature.CACHE_CONTROL_INJECTION, + Feature.STREAM_COST_INJECTION, + Feature.THINKING_SIGNATURE_RETRY, +} + +RESPONSES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { + Feature.DROP_PARAMS, + Feature.MODIFY_PARAMS, + Feature.STREAMING, + Feature.PROMPT_MANAGEMENT, + Feature.TEXT_FORMAT, + Feature.CHAT_COMPLETIONS_BRIDGE, + Feature.RETRIES_FALLBACKS, + Feature.DEPLOYMENT_RESPONSE_ID, + Feature.ENCODED_ITEM_IDS, + Feature.INPUT_REWRITE, + Feature.EXTRA_BODY, + Feature.PROVIDER_AFFINITY, + Feature.RESPONSE_CACHE, + Feature.MANAGED_RESOURCE_IDS, + Feature.CACHE_CONTROL_MARKS, + Feature.ANTHROPIC_TOOL_SEARCH, + Feature.TOOL_SCHEMA_REPAIR, + Feature.MODEL_PARAM_GATING, + Feature.WORKLOAD_IDENTITY, + Feature.CLIENT_HEADERS, + Feature.DEPLOYMENT_HOOK, +} + RULES: Final[Rules] = ( + RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY, features=CHAT_PYTHON_FEATURES), RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY), RouteRule(Route.OCR, Rollout.RUST_REQUIRED), - RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "deepseek", "vertex_ai"})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, targets=frozenset({Target.CHAT_COMPLETIONS, Target.RESPONSES})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=MESSAGES_PYTHON_FEATURES), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), + RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY, features=RESPONSES_PYTHON_FEATURES), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY), @@ -69,3 +225,9 @@ def rollout(context: RouteContext, rules: Rules | None = None) -> Rollout: def decision(context: RouteContext, rules: Rules | None = None) -> Decision: return _decision(rollout(context, rules)) + + +def rust_might_handle(route: Route, rules: Rules) -> bool: + own: Final = tuple(rule for rule in rules if rule.route is route) + last: Final = next((i for i, rule in enumerate(own) if rule.unconstrained()), len(own)) + return any(_decision(rule.rollout) is not Decision.PYTHON for rule in own[: last + 1]) diff --git a/litellm/rust_bridge/chat_completions/detectors.py b/litellm/rust_bridge/chat_completions/detectors.py new file mode 100644 index 00000000000..5aa65a945b1 --- /dev/null +++ b/litellm/rust_bridge/chat_completions/detectors.py @@ -0,0 +1,137 @@ +"""Detectors for Python-only Chat Completions behavior.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +import litellm +from litellm.rust_bridge.detectors import ( + agentic_loop_hook, + api_base, + call_callbacks, + mock_response, + overridden, + present, + retries_fallbacks, +) +from litellm.rust_bridge.public_call import NativeCall, optional_sequence, optional_str + +_PARAMS: Final = TypeAdapter(dict[str, object]) + + +def chat_mock_response(request: NativeCall) -> bool: + """Chat mocks can supply text, tool calls or a timeout without provider IO.""" + return ( + mock_response(request) + or bool(request.resolved.get("mock_tool_calls")) + or bool(request.resolved.get("mock_timeout")) + ) + + +def model_list(request: NativeCall) -> bool: + """Python fans the request out across the supplied model deployments.""" + return present(request, "model_list") + + +def model_alias_map(request: NativeCall) -> bool: + """Python rewrites this model through the configured alias map.""" + return bool(litellm.model_alias_map) and str(request.resolved["model"]) in litellm.model_alias_map + + +def chat_retries_fallbacks(request: NativeCall) -> bool: + """Chat also reads global model fallbacks and context-window fallbacks.""" + return ( + retries_fallbacks(request) + or bool(optional_sequence(litellm.model_fallbacks)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm annotates it with a bare list + or present(request, "context_window_fallback_dict") + ) + + +_CHAT_UNGATED: Final = frozenset( + { + "model", + "messages", + "custom_llm_provider", + "api_version", + "drop_params", + "allowed_openai_params", + "additional_drop_params", + } +) + + +_NAMES: Final = TypeAdapter(list[str]) + + +def _chat_non_default_params(request: NativeCall) -> Mapping[str, object]: + from litellm.constants import DEFAULT_CHAT_COMPLETION_PARAM_VALUES + + defaults: Final = _PARAMS.validate_python(DEFAULT_CHAT_COMPLETION_PARAM_VALUES) + return { + name: value + for name, value in request.resolved.items() + if name in defaults and name not in _CHAT_UNGATED and value != defaults[name] + } + + +def _chat_provider(request: NativeCall) -> tuple[str, str | None]: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model: Final = str(request.resolved["model"]) + provider: Final = optional_str(request.resolved.get("custom_llm_provider")) + try: + resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider, api_base=api_base(request)) + except BadRequestError: + return model, provider + return resolved_model, resolved_provider + + +def _unvalidated(params: Mapping[str, object]) -> frozenset[str]: + from litellm.utils import PROVIDER_UNVALIDATED_PARAMS + + return PROVIDER_UNVALIDATED_PARAMS | (frozenset({"n"}) if params.get("n") == 1 else frozenset[str]()) + + +def chat_model_param_gating(request: NativeCall) -> bool: + """Python rejects unsupported chat params or maps them differently with drop_params.""" + from litellm.exceptions import UnsupportedParamsError + from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, # pyright: ignore[reportUnknownVariableType] # returns a bare list + ) + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + params: Final = _chat_non_default_params(request) + model, provider = _chat_provider(request) + if not params or provider is None: + return False + supported: Final = frozenset(_NAMES.validate_python(get_supported_openai_params(model, provider) or [])) + if any(name not in supported for name in params.keys() - _unvalidated(params)): + return True + try: + llm_provider: Final = LlmProviders(provider) + except ValueError: + return False + config: Final = ProviderConfigManager.get_provider_chat_config(model, llm_provider) + if config is None: + return False + try: + strict: Final = _PARAMS.validate_python( + config.map_openai_params(dict(params), {}, model, drop_params=False) # pyright: ignore[reportUnknownMemberType] # returns a bare dict + ) + except UnsupportedParamsError: + return True + return strict != _PARAMS.validate_python( + config.map_openai_params(dict(params), {}, model, drop_params=True) # pyright: ignore[reportUnknownMemberType] # returns a bare dict + ) + + +def chat_agentic_loop_hook(request: NativeCall) -> bool: + """Chat callbacks can also override the chat-completion-specific agentic loop.""" + return agentic_loop_hook(request) or overridden( + "async_should_run_chat_completion_agentic_loop", call_callbacks(request) + ) diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py new file mode 100644 index 00000000000..0ad1b787e38 --- /dev/null +++ b/litellm/rust_bridge/detectors.py @@ -0,0 +1,183 @@ +"""Shared feature detectors and helpers for native-route selection.""" + +from __future__ import annotations + +import os +from collections.abc import Callable, Iterable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias, cast # noqa: TID251 # tool lists stay caller-owned + +import httpx + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.rust_bridge.catalog import Feature +from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str + +if TYPE_CHECKING: + from litellm.responses.mcp.litellm_proxy_mcp_handler import ToolParam + +Detector: TypeAlias = Callable[[NativeCall], bool] +Detectors: TypeAlias = Mapping[Feature, Detector] + + +def detected(detectors: Detectors, request: NativeCall) -> frozenset[Feature]: + return frozenset(feature for feature, detector in detectors.items() if detector(request)) + + +def present(request: NativeCall, name: str) -> bool: + return request.resolved.get(name) is not None + + +def mock_response(request: NativeCall) -> bool: + """Python returns a supplied mock instead of calling the provider.""" + return present(request, "mock_response") + + +def caller_client(request: NativeCall) -> bool: + """Python sends through the caller's HTTP client or shared session.""" + return present(request, "client") or present(request, "shared_session") + + +def extra_body(request: NativeCall) -> bool: + """Python transforms and merges additional fields into the provider body.""" + return bool(request.resolved.get("extra_body")) + + +def streaming(request: NativeCall) -> bool: + """The caller requested streaming, which some native routes do not support yet.""" + return request.resolved.get("stream") is True + + +def timeout_variants(request: NativeCall) -> bool: + """Python handles request and stream overrides and httpx phase timeouts.""" + return ( + present(request, "request_timeout") + or (streaming(request) and present(request, "stream_timeout")) + or isinstance(request.resolved.get("timeout"), httpx.Timeout) + ) + + +def proxy_mcp_tools(request: NativeCall) -> bool: + """Python expands proxy MCP tools and may execute their tool calls.""" + from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + + tools: Final = optional_sequence(request.resolved.get("tools")) + return tools is not None and LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway( + cast("Iterable[ToolParam]", tools) # cast-ok: the predicate skips entries that are not MCP tool dicts + ) + + +def prompt_management(request: NativeCall) -> bool: + """Python resolves prompt-manager inputs or strips the litellm_agent prefix.""" + from litellm.integrations.litellm_agent.litellm_agent_model_resolver import LITELLM_AGENT_PREFIX + from litellm.types.utils import DynamicPromptManagementParamLiteral + + return ( + present(request, "prompt_id") + or str(request.resolved["model"]).startswith(LITELLM_AGENT_PREFIX) + or any(request.resolved.get(param) for param in DynamicPromptManagementParamLiteral.list_all_params()) + ) + + +def retries_fallbacks(request: NativeCall) -> bool: + """Python runs the requested retry policy or fallback chain.""" + return bool(litellm.num_retries) or any( + present(request, name) for name in ("num_retries", "max_retries", "fallbacks", "retry_policy") + ) + + +def api_base(request: NativeCall) -> str | None: + return optional_str(request.resolved.get("api_base")) or optional_str(request.resolved.get("base_url")) + + +def cache_control_injection(request: NativeCall) -> bool: + """Python injects Anthropic prompt-cache breakpoints into messages and system content.""" + from litellm.integrations.anthropic_cache_control_hook import configured_injection_points + + return ( + litellm.enable_anthropic_prompt_caching is True + or request.resolved.get("enable_prompt_caching") is True + or bool(configured_injection_points(request.resolved.get("cache_control_injection_points"))) + ) + + +def drop_params(request: NativeCall) -> bool: + """Python drops unsupported or explicitly excluded request parameters.""" + return ( + bool(litellm.drop_params) + or normalize_drop_params(request.resolved.get("drop_params")) is True + or bool(optional_sequence(request.resolved.get("additional_drop_params"))) + ) + + +def modify_params(_request: NativeCall) -> bool: + """Python rewrites unsupported request values when modify_params is enabled.""" + return bool(litellm.modify_params) + + +def _custom_loggers(callbacks: Iterable[object]) -> tuple[CustomLogger, ...]: + from litellm.litellm_core_utils.litellm_logging import get_custom_logger_compatible_class + + resolved: Final = ( + get_custom_logger_compatible_class(callback) if isinstance(callback, str) else callback + for callback in callbacks + ) + return tuple(callback for callback in resolved if isinstance(callback, CustomLogger)) + + +def overridden(method: str, callbacks: Iterable[object]) -> bool: + base: Final = getattr(CustomLogger, method) + return any(getattr(type(callback), method, base) is not base for callback in _custom_loggers(callbacks)) + + +def _request_success_callbacks(request: NativeCall) -> tuple[object, ...]: + from litellm.litellm_core_utils.litellm_logging import Logging + + supplied: Final = optional_sequence(request.resolved.get("success_callback")) or () + logging_obj: Final = request.resolved.get("litellm_logging_obj") + dynamic: Final = ( + optional_sequence(logging_obj.dynamic_success_callbacks) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging annotates the list with a bare Callable + if isinstance(logging_obj, Logging) + else None + ) + return (*supplied, *(dynamic or ())) + + +def call_callbacks(request: NativeCall) -> tuple[object, ...]: + return (*litellm.callbacks, *_request_success_callbacks(request)) + + +def agentic_loop_hook(request: NativeCall) -> bool: + """A callback can request follow-up tool calls through Python's agentic loop.""" + return overridden("async_should_run_agentic_loop", call_callbacks(request)) + + +def response_cache(request: NativeCall) -> bool: + """Python owns cache lookup and storage unless this call disables cache reads.""" + cache: Final = optional_mapping(request.resolved.get("cache")) or {} + return litellm.cache is not None and cache.get("no-cache") is not True + + +_NATIVE_ECDH_CURVES: Final = frozenset({"x25519", "prime256v1", "secp256r1", "p-256", "secp384r1", "p-384"}) + + +def tls_policy(_request: NativeCall) -> bool: + """Python applies TLS settings or ECDH curves the native client cannot configure.""" + curve: Final = os.getenv("SSL_ECDH_CURVE", litellm.ssl_ecdh_curve) + return bool(os.getenv("SSL_SECURITY_LEVEL", litellm.ssl_security_level)) or ( + bool(curve) and str(curve).strip().lower() not in _NATIVE_ECDH_CURVES + ) + + +COMMON: Final[Detectors] = MappingProxyType( + { + Feature.MOCK_RESPONSE: mock_response, + Feature.PROXY_MCP_TOOLS: proxy_mcp_tools, + Feature.CALLER_CLIENT: caller_client, + Feature.TIMEOUT_VARIANTS: timeout_variants, + Feature.AGENTIC_LOOP_HOOK: agentic_loop_hook, + Feature.TLS_POLICY: tls_policy, + } +) diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index da9b9bed54f..ca98fe2b5b2 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -4,11 +4,10 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from typing import Final, Generic, TypeAlias, TypeVar +from litellm._logging import verbose_logger from litellm.rust_bridge import catalog, runtime from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import Route, RouteContext, Rules -from litellm.rust_bridge.configuration import Decision -from litellm.rust_bridge.configuration import decision as rollout_decision RequestT: Final = TypeVar("RequestT") NativeT: Final = TypeVar("NativeT") @@ -31,18 +30,6 @@ class PublicDispatch(Generic[RequestT]): route: Route request: Callable[[tuple[object, ...], Mapping[str, object]], RequestT | None] context: Callable[[RequestT], RouteContext] - bypass: Callable[[RequestT], bool] | None = None - - def _requires_projection(self, rules: Rules) -> bool: - for rule in rules: - if rule.route is not self.route: - continue - if rule.providers is not None or rule.models is not None: - if rollout_decision(rule.rollout) is not Decision.PYTHON: - return True - continue - return rollout_decision(rule.rollout) is not Decision.PYTHON - return False def _native_request(self, args: tuple[object, ...], kwargs: Mapping[str, object]) -> RequestT: request: Final = self.request(args, kwargs) @@ -50,13 +37,17 @@ class PublicDispatch(Generic[RequestT]): raise runtime.NoPythonImplementationError( f"{self.route.value} has no Python implementation, so every call must project to a native request" ) - if self.bypass is not None and self.bypass(request): - raise runtime.NoPythonImplementationError( - f"{self.route.value} has no Python implementation, so a call its bypass predicate matches cannot " - "be served" - ) return request + def _guarded_context(self, request: RequestT) -> RouteContext | None: + try: + return self.context(request) + except Exception: # noqa: BLE001 # detectors call legacy handlers, so any error must fall back to Python + verbose_logger.warning( + "%s route detection failed, serving the call on Python", self.route.value, exc_info=True + ) + return None + def run( self, args: tuple[object, ...], @@ -77,13 +68,16 @@ class PublicDispatch(Generic[RequestT]): python=python, rules=selected_rules, ) - if not self._requires_projection(selected_rules): + if not catalog.rust_might_handle(self.route, selected_rules): return python(*args, **kwargs) request: Final = self.request(args, kwargs) - if request is None or (self.bypass is not None and self.bypass(request)): + if request is None: + return python(*args, **kwargs) + context: Final = self._guarded_context(request) + if context is None: return python(*args, **kwargs) return runtime.run( - self.context(request), + context, binding=binding, native=lambda hook: native(hook, request, args, kwargs), python=lambda: python(*args, **kwargs), @@ -110,13 +104,16 @@ class PublicDispatch(Generic[RequestT]): python=python, rules=selected_rules, ) - if not self._requires_projection(selected_rules): + if not catalog.rust_might_handle(self.route, selected_rules): return await python(*args, **kwargs) request: Final = self.request(args, kwargs) - if request is None or (self.bypass is not None and self.bypass(request)): + if request is None: + return await python(*args, **kwargs) + context: Final = self._guarded_context(request) + if context is None: return await python(*args, **kwargs) return await runtime.arun( - self.context(request), + context, binding=binding, native=lambda hook: native(hook, request, args, kwargs), python=lambda: python(*args, **kwargs), diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index a82e734b47d..ae913882968 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -11,6 +11,7 @@ import contextvars import datetime import traceback import uuid +import weakref from collections.abc import Awaitable, Coroutine, Mapping from dataclasses import dataclass from typing import ( @@ -61,9 +62,21 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return _claim_budget_reservation(CallSetup(supplied, arguments), asynchronous) + return _begin_call(CallSetup(supplied, arguments), asynchronous) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return _claim_budget_reservation(CallSetup(logger, prepared), asynchronous) + return _begin_call(CallSetup(logger, prepared), asynchronous) + + +# Loggers whose current native call is async: their sync failure step is followed by the async one +_ASYNC_CALLS: Final[weakref.WeakSet[Logging]] = weakref.WeakSet() + + +def _begin_call(call_setup: CallSetup, asynchronous: bool) -> CallSetup: + if asynchronous: + _ASYNC_CALLS.add(call_setup.logger) + else: + _ASYNC_CALLS.discard(call_setup.logger) + return _claim_budget_reservation(call_setup, asynchronous) def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> CallSetup: @@ -135,6 +148,8 @@ class LoggingSurface(Protocol): traceback_exception: str, start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, + *, + for_async_request: bool = False, ) -> None: ... def async_failure_handler( @@ -261,7 +276,7 @@ def failure_handler( if not is_internal_call(): unbind_budget_reservation_from_callbacks(logger.litellm_params) return logger.async_failure_handler(error, trace, start, end) - logger.failure_handler(error, trace, start, end) + logger.failure_handler(error, trace, start, end, for_async_request=logger in _ASYNC_CALLS) return None diff --git a/litellm/rust_bridge/host/routes/responses.py b/litellm/rust_bridge/host/routes/responses.py index 470a89a0e05..22e8cbbf3b0 100644 --- a/litellm/rust_bridge/host/routes/responses.py +++ b/litellm/rust_bridge/host/routes/responses.py @@ -4,9 +4,8 @@ from collections.abc import Mapping from typing import Final import litellm -from litellm import get_llm_provider from litellm.rust_bridge import failures -from litellm.rust_bridge.public_call import inference_decline_reason, known_params, optional_str +from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams, ResponsesAPIResponse PARAMETERS: Final = tuple(ResponsesAPIOptionalRequestParams.__annotations__) @@ -34,14 +33,3 @@ def map_failure(error: Exception, request: Mapping[str, object]) -> Exception: request, optional_str(request.get("api_base")) or optional_str(request.get("base_url")), ) - - -def decline_reason(request: Mapping[str, object]) -> str | None: - if optional_str(request.get("custom_llm_provider")) is None and "/" not in str(request["model"]): - try: - _, provider, _, _ = get_llm_provider(model=str(request["model"])) - except litellm.exceptions.BadRequestError: - return "native Responses could not resolve the provider" - if provider != "openai": - return "native HTTP responses provider" - return inference_decline_reason(PARAMETERS, request) diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py new file mode 100644 index 00000000000..691c04f2daa --- /dev/null +++ b/litellm/rust_bridge/messages/detectors.py @@ -0,0 +1,85 @@ +"""Detectors for Python-only Anthropic Messages behavior.""" + +from __future__ import annotations + +from typing import Final + +import litellm +from litellm.rust_bridge.detectors import overridden, present, streaming +from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str + + +def anthropic_workload_identity(request: NativeCall) -> bool: + """Python mints an Anthropic federation token when no static credential is configured.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + from litellm.secret_managers.main import get_secret_str + + def configured(param: str, env: str) -> bool: + return present(request, param) or get_secret_str(env) is not None + + return ( + AnthropicModelInfo.get_api_key(optional_str(request.resolved.get("api_key"))) is None + and AnthropicModelInfo.get_auth_token() is None + and configured("anthropic_federation_rule_id", "ANTHROPIC_FEDERATION_RULE_ID") + and configured("anthropic_organization_id", "ANTHROPIC_ORGANIZATION_ID") + ) + + +_THINKING_BLOCKS: Final = frozenset({"thinking", "redacted_thinking"}) + + +def _replays_thinking(message: object) -> bool: + content: Final = optional_sequence((optional_mapping(message) or {}).get("content")) or () + return any((optional_mapping(block) or {}).get("type") in _THINKING_BLOCKS for block in content) + + +def thinking_signature_retry(request: NativeCall) -> bool: + """Python may strip replayed thinking blocks and retry after a rejected signature.""" + return any(_replays_thinking(message) for message in optional_sequence(request.resolved.get("messages")) or ()) + + +def stream_cost_injection(request: NativeCall) -> bool: + """Python adds response cost to streaming usage when the global setting is truthy.""" + return streaming(request) and bool(litellm.include_cost_in_streaming_usage) + + +def pre_request_hook(_request: NativeCall) -> bool: + """The async Messages handler runs callbacks that can rewrite the request.""" + return overridden("async_pre_request_hook", litellm.callbacks) + + +def websearch_interception(_request: NativeCall) -> bool: + """The async Messages handler can answer a web-search request without an LLM call.""" + from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger + + return any(isinstance(callback, WebSearchInterceptionLogger) for callback in litellm.callbacks) + + +def _resolved_provider(request: NativeCall) -> str | None: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + provider: Final = optional_str(request.resolved.get("custom_llm_provider")) + try: + return get_llm_provider(str(request.resolved["model"]), provider)[1] + except BadRequestError: + return provider + + +def intercepted(request: NativeCall) -> bool: + """The async Messages handler delegates matching tools to a registered interceptor.""" + from litellm.llms.anthropic.pass_through.messages.interceptors import get_messages_interceptors + + tools: Final = optional_sequence(request.resolved.get("tools")) + if not tools: + return False + mappings: Final = tuple(mapping for mapping in map(optional_mapping, tools) if mapping is not None) + provider: Final = _resolved_provider(request) + return any(interceptor.can_handle(mappings, provider) for interceptor in get_messages_interceptors()) + + +def safeguards_classifier(request: NativeCall) -> bool: + """Python's proxy classifier evaluates dangerous tool use for non-Claude models.""" + from litellm.llms.anthropic.pass_through.safeguards import native_route_classifies + + return native_route_classifies(str(request.resolved["model"]), request.resolved.get("safeguards")) diff --git a/litellm/rust_bridge/public_call.py b/litellm/rust_bridge/public_call.py index 332421c68cc..8625ca71bb2 100644 --- a/litellm/rust_bridge/public_call.py +++ b/litellm/rust_bridge/public_call.py @@ -8,33 +8,6 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # narrows caller-owned containers without copying them -import litellm - -_INFERENCE_CONTEXT: Final = frozenset( - { - "model", - "messages", - "input", - "api_key", - "api_base", - "base_url", - "custom_llm_provider", - "extra_headers", - "timeout", - "request_timeout", - "callbacks", - "success_callback", - "failure_callback", - "metadata", - "litellm_metadata", - "litellm_call_id", - "litellm_trace_id", - "litellm_logging_obj", - "litellm_credential_name", - "proxy_server_request", - } -) - def signature(legacy: Callable[..., object]) -> inspect.Signature: return inspect.signature(legacy) @@ -75,19 +48,6 @@ def known_params(names: Collection[str], request: Mapping[str, object]) -> Mappi return {name: value for name, value in request.items() if name in names and value is not None} -def inference_decline_reason(parameters: tuple[str, ...], kwargs: Mapping[str, object]) -> str | None: - if litellm.drop_params or litellm.modify_params: - return "native inference does not implement the configured parameter rewrites" - for name, value in kwargs.items(): - if value is None: - continue - if name in {"cache", "caching"}: - continue - if name not in parameters and name not in _INFERENCE_CONTEXT: - return f"native inference does not implement {name}" - return None - - _BAGS: Final = frozenset({inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}) @@ -113,6 +73,32 @@ def _without_bags(legacy: inspect.Signature, named: Mapping[str, object]) -> Map return MappingProxyType({name: value for name, value in named.items() if legacy.parameters[name].kind not in _BAGS}) +@dataclass(frozen=True, slots=True) +class AsyncFlag: + """The keyword that makes a sync entry hand off to its async entry, as in ``completion(..., acompletion=True)``. + + ``is_set`` is the route's own reading of the flag. ``@client`` hands off on ``is True``, the Messages handler on truthiness. + + Completion uses ``AsyncFlag("acompletion", sync=signature(_PYTHON_COMPLETION))``. + Messages uses ``AsyncFlag("is_async", sync=signature(_PYTHON_MESSAGES), is_set=bool)``. + """ + + name: str + sync: inspect.Signature + is_set: Callable[[object], bool] = lambda value: value is True + + def requested(self, kwargs: Mapping[str, object]) -> bool: + return self.is_set(kwargs.get(self.name)) + + def dropped(self, kwargs: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({name: value for name, value in kwargs.items() if name != self.name}) + + def keywords(self, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]: + """Name positionals by the sync signature, since the async entry may order its parameters differently.""" + named: Final = _without_bags(self.sync, self.sync.bind_partial(*args, **kwargs).arguments) + return self.dropped({**named, **kwargs}) + + def native_call(legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NativeCall: positional: Final = legacy.bind_partial(*args) positional.apply_defaults() diff --git a/litellm/rust_bridge/responses/detectors.py b/litellm/rust_bridge/responses/detectors.py new file mode 100644 index 00000000000..392c11f48a8 --- /dev/null +++ b/litellm/rust_bridge/responses/detectors.py @@ -0,0 +1,210 @@ +"""Detectors for Python-only Responses behavior.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from itertools import accumulate, chain, repeat, takewhile +from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # tool lists stay caller-owned + +from pydantic import TypeAdapter + +import litellm +from litellm.rust_bridge.detectors import api_base, call_callbacks, overridden, present +from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str + +if TYPE_CHECKING: + from openai.types.responses import ResponseInputParam + + from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponsesAPIOptionalRequestParams + + +def client_headers(request: NativeCall) -> bool: + """Python merges proxy-forwarded headers into extra_headers before sending.""" + return bool(optional_mapping(request.resolved.get("headers"))) + + +def provider_affinity(request: NativeCall) -> bool: + """Python validates and fills the affinity header from the session ID.""" + return present(request, "provider_affinity_header") + + +def text_format(request: NativeCall) -> bool: + """Python builds the output schema and parses the response into the supplied type.""" + return present(request, "text_format") + + +def chat_completions_bridge(request: NativeCall) -> bool: + """Python serves this Responses request through Chat Completions.""" + from litellm.responses.main import ( + _normalize_openai_chat_completions_responses_model, # pyright: ignore[reportPrivateUsage] # the handler's own bridge prefix check + ) + + return ( + bool(request.resolved.get("use_chat_completions_api")) + or _normalize_openai_chat_completions_responses_model(str(request.resolved["model"]))[1] + ) + + +def _openai_model(request: NativeCall) -> str: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model: Final = str(request.resolved["model"]) + try: + return get_llm_provider(model, optional_str(request.resolved.get("custom_llm_provider")))[0] + except BadRequestError: + return model + + +def _model_info(request: NativeCall) -> Mapping[str, object] | None: + return optional_mapping(request.resolved.get("model_info")) + + +def _tools(request: NativeCall) -> Sequence[object]: + return optional_sequence(request.resolved.get("tools")) or () + + +def managed_resource_ids(request: NativeCall) -> bool: + """Python replaces managed file IDs in Responses input and tools with provider IDs.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + update_responses_input_with_model_file_ids, + update_responses_tools_with_model_file_ids, + ) + + model_id: Final = optional_str((_model_info(request) or {}).get("id")) + mapping: Final = request.resolved.get("model_file_id_mapping") + file_ids: Final = cast("dict[str, dict[str, str]] | None", mapping) # cast-ok: handler passes it unvalidated + input_value: Final = request.resolved.get("input") + tools: Final = [dict(tool) for tool in map(optional_mapping, _tools(request)) if tool is not None] + rewritten_tools: Final = update_responses_tools_with_model_file_ids(tools, model_id, file_ids) if tools else tools + return update_responses_input_with_model_file_ids(input_value, model_id, file_ids) != input_value or ( + rewritten_tools != tools + ) + + +def input_rewrite(request: NativeCall) -> bool: + """Python normalizes replayed Responses items and removes unsupported IDs or namespaces.""" + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + input_value: Final = request.resolved.get("input") + if not isinstance(input_value, list): + return False + config: Final = OpenAIResponsesAPIConfig() + replayable: Final = config._drop_bridge_minted_reasoning_items( # pyright: ignore[reportPrivateUsage] # use the handler's replay policy + cast("ResponseInputParam", input_value), # cast-ok: the handler accepts unvalidated input items + ) + normalized: Final = config._validate_input_param(replayable) # pyright: ignore[reportPrivateUsage] # use the handler's normalization + cleaned: Final = config._drop_foreign_tool_call_item_ids(normalized) # pyright: ignore[reportPrivateUsage] # use the handler's id cleanup + return config.strip_custom_tool_call_namespace_from_responses_input(cleaned) != input_value + + +def _children(value: object) -> tuple[object, ...]: + if isinstance(value, Mapping): + return tuple(cast("Mapping[object, object]", value).values()) # cast-ok: only values are walked + if isinstance(value, list | tuple): + return tuple(cast("Sequence[object]", value)) # cast-ok: only items are walked + return () + + +def _next_level(level: tuple[object, ...], _: None) -> tuple[object, ...]: + return tuple(chain.from_iterable(map(_children, level))) + + +def _mentions(value: object, key: str) -> bool: + levels: Final = takewhile(bool, accumulate(repeat(None), _next_level, initial=(value,))) + return any( + key in cast("Mapping[object, object]", node) # cast-ok: node passed the Mapping guard + for node in chain.from_iterable(levels) + if isinstance(node, Mapping) + ) + + +def cache_control_marks(request: NativeCall) -> bool: + """Python strips or translates cache_control marks in Responses input and tools.""" + return _mentions(request.resolved.get("input"), "cache_control") or _mentions(_tools(request), "cache_control") + + +def anthropic_tool_search(request: NativeCall) -> bool: + """Python translates Anthropic tool-search definitions for the Responses API.""" + from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool + + return any(is_anthropic_tool_search_chat_tool(tool) for tool in _tools(request)) + + +def _undecoded_schema(tool: object) -> bool: + mapping: Final = optional_mapping(tool) + parameters: Final = None if mapping is None else mapping.get("parameters") + return parameters is not None and not isinstance(parameters, Mapping) + + +def tool_schema_repair(request: NativeCall) -> bool: + """Python decodes or repairs tool schemas before sending them to OpenAI.""" + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.types.router import GenericLiteLLMParams + + tools: Final = _tools(request) + if any(_undecoded_schema(tool) for tool in tools): + return True + sanitized: Final = OpenAIResponsesAPIConfig()._sanitized_tool_schemas_for_openai( # pyright: ignore[reportPrivateUsage] # the handler's own schema rewrite + model=_openai_model(request), + tools=cast("Sequence[ALL_RESPONSES_API_TOOL_PARAMS]", tools), # cast-ok: sanitizer skips non-dict tools + litellm_params=GenericLiteLLMParams(model_info=_model_info(request)), + ) + return bool(tools) and sanitized is not tools + + +_PARAMS: Final = TypeAdapter(dict[str, object]) + + +def model_param_gating(request: NativeCall) -> bool: + """Python clamps, drops or remaps Responses params according to the model's capabilities.""" + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.rust_bridge.host.routes.responses import PARAMETERS + + params: Final = {name: request.resolved[name] for name in PARAMETERS if request.resolved.get(name) is not None} + mapped: Final = _PARAMS.validate_python( + OpenAIResponsesAPIConfig().map_openai_params( # pyright: ignore[reportUnknownMemberType] # returns a bare dict + cast("ResponsesAPIOptionalRequestParams", params), # cast-ok: the handler passes these kwargs + _openai_model(request), + drop_params=True, + ) + ) + return mapped != params + + +def openai_workload_identity(request: NativeCall) -> bool: + """Python mints an OpenAI federation token when no static API key is configured.""" + from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config + + api_key: Final = optional_str(request.resolved.get("api_key")) or litellm.api_key or litellm.openai_key + return resolve_openai_workload_identity_config(api_key, api_base(request), request.resolved) is not None + + +def deployment_response_id(request: NativeCall) -> bool: + """Python encodes the deployment ID into the returned Responses ID.""" + metadata: Final = optional_mapping(request.resolved.get("litellm_metadata")) or {} + model_info: Final = optional_mapping(metadata.get("model_info")) or {} + return model_info.get("id") is not None + + +def _encoded_item(item: Mapping[str, object]) -> bool: + from litellm.responses.utils import ResponsesAPIRequestUtils + + item_id: Final = optional_str(item.get("id")) + encrypted_content: Final = optional_str(item.get("encrypted_content")) + return (item_id is not None and ResponsesAPIRequestUtils.decode_encrypted_item_id(item_id) is not None) or ( + encrypted_content is not None + and ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(encrypted_content)[1] != encrypted_content + ) + + +def encoded_item_ids(request: NativeCall) -> bool: + """Python restores provider IDs or encrypted content from previously wrapped items.""" + items: Final = optional_sequence(request.resolved.get("input")) or () + mappings: Final = (optional_mapping(item) for item in items) + return any(_encoded_item(item) for item in mappings if item is not None) + + +def deployment_hook(request: NativeCall) -> bool: + """The sync Responses handler runs deployment hooks that native currently skips.""" + return overridden("async_pre_call_deployment_hook", call_callbacks(request)) diff --git a/litellm/types/extraction.py b/litellm/types/extraction.py new file mode 100644 index 00000000000..4311e549b9a --- /dev/null +++ b/litellm/types/extraction.py @@ -0,0 +1,162 @@ +from collections.abc import Mapping, Sequence +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import ( + ConfigDict, + Field, + JsonValue, + PrivateAttr, + StrictBool, + StrictFloat, + StrictInt, + StrictStr, + model_validator, + with_config, +) +from typing_extensions import ReadOnly, Required, Self, TypedDict + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + +ExtractionFieldType: TypeAlias = Literal[ + "text", "number", "date", "time", "datetime", "duration", "list", "estimate", "forecast" +] +ExtractionFieldName: TypeAlias = Annotated[str, Field(pattern=r"^[A-Za-z_][A-Za-z0-9_.\-]{0,63}$")] +JsonNumber: TypeAlias = StrictInt | StrictFloat +ExtractionBound: TypeAlias = JsonNumber | StrictStr +EXTRACTION_STATE_REWRITTEN_METADATA_KEY: Final = "extraction_state_rewritten_by_guardrail" + + +def _is_none(value: object) -> bool: + return value is None + + +_OMITTED_WHEN_NONE: Final = Field(exclude_if=_is_none) +OptionalJson: TypeAlias = Annotated[JsonValue, _OMITTED_WHEN_NONE] +OptionalFloat: TypeAlias = Annotated[float | None, _OMITTED_WHEN_NONE] +OptionalInt: TypeAlias = Annotated[int | None, _OMITTED_WHEN_NONE] + + +class _ClairicObject(LiteLLMPydanticObjectBase): + model_config = ConfigDict(extra="allow", frozen=True) + + +class ExtractionField(_ClairicObject): + type: Annotated[ExtractionFieldType, Field(description="The answer type for this field")] + instructions: Annotated[str | None, Field(description="What the field should contain")] = None + unit: Annotated[str | None, Field(description="Unit label. For durations, seconds through years")] = None + integer: Annotated[StrictBool | None, Field(description="Numbers and durations only")] = None + decimals: Annotated[JsonNumber | None, Field(description="Numbers, durations and number lists only, 0 to 6")] = None + items: Annotated[ + Literal["text", "number", "date"] | None, Field(description="Item type, required for list fields") + ] = None + min: Annotated[ExtractionBound | None, Field(description="Estimate and forecast lower bound")] = None + max: Annotated[ExtractionBound | None, Field(description="Estimate and forecast upper bound")] = None + step: Annotated[ExtractionBound | None, Field(description="Estimate grid step, or month or year")] = None + scale: Literal["linear", "log"] | None = None + null_threshold: Annotated[ + JsonNumber | None, Field(description="Probability of null above which the value is null") + ] = None + horizon: Annotated[JsonNumber | None, Field(description="Forecast only, number of steps")] = None + interval: Literal["step", "minute", "hour", "day", "week", "month", "quarter", "year"] | None = None + + @model_validator(mode="after") + def _list_fields_name_their_item_type(self) -> Self: + if self.type == "list" and self.items is None: + raise ValueError("list fields need items set to text, number or date") + return self + + +class ExtractionBase64Image(_ClairicObject): + content_type: Annotated[str, Field(description="The image media type, such as image/png")] + base64: Annotated[str, Field(description="Base64-encoded image bytes")] + + +ExtractionImage: TypeAlias = Annotated[str, Field(pattern=r"^[Dd][Aa][Tt][Aa]:")] | ExtractionBase64Image + + +class ExtractionOptions(_ClairicObject): + include_quantiles: Annotated[StrictBool | None, Field(description="Add 99 percentiles to estimates")] = None + return_state_text: Annotated[StrictBool | None, Field(description="Return the text that spans index into")] = None + + +class ExtractionRequestBody(_ClairicObject): + state: Annotated[ + str, + Field(description="The text to read. Must be non-empty unless images are sent. Text answers are spans of it"), + ] = "" + fields: Annotated[ + Mapping[ExtractionFieldName, ExtractionField], + Field(min_length=1, description="Field name to typed field spec. Each name becomes a key in the answers"), + ] + images: Annotated[ + tuple[ExtractionImage, ...] | None, + Field(max_length=4, description="Up to 4 images as data URLs or content_type plus base64, no remote URLs"), + ] = None + reference_date: Annotated[str | None, Field(description="ISO date that anchors dates written without a year")] = ( + None + ) + options: ExtractionOptions | None = None + + +class ExtractionRequest(ExtractionRequestBody): + model: Annotated[str, Field(description="An extraction model from the proxy model_list")] + + +class ExtractionAnswer(_ClairicObject): + type: str + value: Annotated[JsonValue, Field(description="The answer, null when the input does not determine it")] + guess: OptionalJson = None + p_null: OptionalFloat = None + confidence: OptionalFloat = None + span: Annotated[tuple[int, ...] | None, _OMITTED_WHEN_NONE] = None + spans: Annotated[tuple[tuple[int, ...], ...] | None, _OMITTED_WHEN_NONE] = None + mean: OptionalJson = None + median: OptionalJson = None + p10: OptionalJson = None + p90: OptionalJson = None + p_below: OptionalFloat = None + p_above: OptionalFloat = None + quantiles: Annotated[tuple[JsonValue, ...] | None, _OMITTED_WHEN_NONE] = None + steps: Annotated[tuple[JsonValue, ...] | None, _OMITTED_WHEN_NONE] = None + + +class ExtractionUsage(_ClairicObject): + input_tokens: int + output_tokens: int + state_tokens: OptionalInt = None + image_tokens: OptionalInt = None + schema_tokens: OptionalInt = None + passes: OptionalInt = None + + +class ExtractionResponse(_ClairicObject): + model: str + answers: Mapping[str, ExtractionAnswer] + truncated: bool + state_text: Annotated[str | None, _OMITTED_WHEN_NONE] = None + usage: ExtractionUsage + + _hidden_params: dict[str, object] = PrivateAttr( # mutable-ok: @client and the proxy write cost into it in place + default_factory=dict + ) + + @property + def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation + return self._hidden_params + + def set_hidden_params(self, params: Mapping[str, object]) -> None: + self._hidden_params.update(params) + + +@with_config(ConfigDict(extra="allow")) +class ExtractionCallParams(TypedDict, total=False): + model: Required[ReadOnly[str]] + state: Required[ReadOnly[str]] + fields: Required[ReadOnly[Mapping[str, Mapping[str, object]]]] + images: ReadOnly[Sequence[str | Mapping[str, object]] | None] + reference_date: ReadOnly[str | None] + api_key: ReadOnly[str | None] + api_base: ReadOnly[str | None] + timeout: ReadOnly[float | None] + custom_llm_provider: ReadOnly[str | None] + extra_headers: ReadOnly[Mapping[str, str] | None] diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 6b08920e042..0ed156524be 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -54,6 +54,7 @@ RoutingStrategyName: TypeAlias = Literal[ TRUSTED_CALLBACK_VARS_FIELD: Final = "litellm_trusted_callback_vars" ADDRESSED_RESPONSE_ID_FIELD: Final = "_litellm_addressed_response_id" +STREAM_USAGE_INJECTED_FIELD: Final = "_litellm_stream_usage_injected" WIRE_NAME: Final = "wire_name" @@ -91,6 +92,7 @@ class ProviderConnection: configurable_clientside_auth_params: "Sequence[str | ConfigurableClientsideParamsCustomAuth] | None" = None use_xai_oauth: bool | None = None github_copilot_auth_type: str | None = None + github_copilot_host: str | None = None github_copilot_user_session: object | None = None token_exchange_endpoint: str | None = None token_exchange_profile: str | None = None @@ -396,6 +398,7 @@ class ProxyRequestState: trusted_callback_vars: Mapping[str, str] | None = field(default=None, metadata=wire(TRUSTED_CALLBACK_VARS_FIELD)) addressed_response_id: str | None = field(default=None, metadata=wire(ADDRESSED_RESPONSE_ID_FIELD)) strip_stream_usage: bool | None = field(default=None, metadata=wire("_litellm_strip_stream_usage")) + stream_usage_injected: bool | None = field(default=None, metadata=wire(STREAM_USAGE_INJECTED_FIELD)) client_side_timeout: bool | None = None model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None diff --git a/litellm/types/proxy/litellm_pre_call_utils.py b/litellm/types/proxy/litellm_pre_call_utils.py index 0b980d38320..f0dfd50bfbe 100644 --- a/litellm/types/proxy/litellm_pre_call_utils.py +++ b/litellm/types/proxy/litellm_pre_call_utils.py @@ -24,5 +24,5 @@ class SecretFields(TypedDict): """ raw_headers: dict - user_provider_credentials: NotRequired[ReadOnly[Mapping[str, str]]] + user_provider_credentials: NotRequired[ReadOnly[Mapping[str, Mapping[str, str]]]] user_provider_credentials_user_id: NotRequired[ReadOnly[str]] diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index 3cc85a9427e..6a4160c3c50 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from typing import Any @@ -13,6 +14,8 @@ class ModelGroupInfoProxy(ModelGroupInfo): health_status: str | None = Field(default=None) health_response_time: float | None = Field(default=None) health_checked_at: str | None = Field(default=None) + custom_tags: tuple[str, ...] | None = None + custom_metadata: Mapping[str, str | int | float | bool] | None = None class UpdateUsefulLinksRequest(LiteLLMBaseModel): diff --git a/litellm/types/router.py b/litellm/types/router.py index 97b0d81d33b..cd1389aea9d 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1184,6 +1184,8 @@ class RouterModelGroupAliasItem(TypedDict): class RetryAttemptRecord(TypedDict): model_group: ReadOnly[str | None] deployment_id: ReadOnly[str | None] + model: ReadOnly[str | None] + base_model: ReadOnly[str | None] exception_type: ReadOnly[str] exception_string: ReadOnly[str] attempted_retries: ReadOnly[int | None] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e268cc90ffd..5939a350970 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -212,6 +212,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_parallel_function_calling: bool | None supports_web_search: bool | None supports_reasoning: bool | None + supports_sampling_params: ReadOnly[bool | None] supports_adaptive_thinking: bool | None supports_legacy_thinking: ReadOnly[bool | None] thinking_always_on: ReadOnly[bool | None] @@ -458,6 +459,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): "responses", "decisions", "evaluation", + "extraction", "ocr", "realtime", ] @@ -518,6 +520,8 @@ class CallTypes(str, Enum): adecisions = "adecisions" systemone = "systemone" asystemone = "asystemone" + extraction = "extraction" + aextraction = "aextraction" arealtime = "_arealtime" aresponses_websocket = "_aresponses_websocket" create_batch = "create_batch" @@ -708,6 +712,8 @@ CallTypesLiteral = Literal[ "adecisions", "systemone", "asystemone", + "extraction", + "aextraction", "_arealtime", "_aresponses_websocket", "create_batch", @@ -821,6 +827,8 @@ API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { "/v1/decisions": [CallTypes.adecisions, CallTypes.decisions], "/systemone": [CallTypes.asystemone, CallTypes.systemone], "/v1/systemone": [CallTypes.asystemone, CallTypes.systemone], + "/extraction": [CallTypes.aextraction, CallTypes.extraction], + "/v1/extraction": [CallTypes.aextraction, CallTypes.extraction], # Batches "/batches": [CallTypes.acreate_batch, CallTypes.create_batch], "/v1/batches": [CallTypes.acreate_batch, CallTypes.create_batch], @@ -4211,7 +4219,7 @@ oauth_token_exchange_litellm_params: Final = ( "token_exchange_profile", "token_exchange_scope", ) -github_copilot_oauth_litellm_params: Final = ("github_copilot_auth_type",) +github_copilot_oauth_litellm_params: Final = ("github_copilot_auth_type", "github_copilot_host") OIDC_SECRET_REFERENCE_PREFIX: Final = "oidc/" API_KEY_SECRET_REFERENCE_REJECTION: Final = ( "Rejected Request: api_key holds an oidc/ secret reference, which only a deployment or credential a proxy " diff --git a/litellm/utils.py b/litellm/utils.py index 6b16fc398bd..1a114ffe9f0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -71,6 +71,7 @@ from litellm.constants import ( DEFAULT_TRIM_RATIO, FUNCTION_DEFINITION_TOKEN_COUNT, HF_CONFIG_FETCH_TIMEOUT_SECONDS, + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, INITIAL_RETRY_DELAY, JITTER, MAX_RETRY_DELAY, @@ -390,6 +391,7 @@ if TYPE_CHECKING: from litellm.llms.base_llm.containers.transformation import BaseContainerConfig from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig + from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig from litellm.llms.base_llm.files.transformation import BaseFilesConfig from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, @@ -1124,83 +1126,90 @@ def function_setup( dynamic_callbacks: Final[list[str | Callable | CustomLogger] | None] = kwargs.pop("callbacks", None) all_callbacks: Final = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) - if len(all_callbacks) > 0: - for callback in all_callbacks: - # check if callback is a string - e.g. "lago", "openmeter" - if isinstance(callback, str): - callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( - callback, - internal_usage_cache=None, - llm_router=None, - ) - if callback is None or any( - type(cb) is type(callback) for cb in litellm._async_success_callback - ): # don't double add a callback - continue - if callback not in litellm.input_callback: - litellm.input_callback.append(callback) - if callback not in litellm.success_callback: - litellm.logging_callback_manager.add_litellm_success_callback(callback) - if callback not in litellm.failure_callback: - litellm.logging_callback_manager.add_litellm_failure_callback(callback) - if callback not in litellm._async_success_callback: - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - if callback not in litellm._async_failure_callback: - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - print_verbose(f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}") + with litellm.logging_callback_manager.callback_lists_lock: + if len(all_callbacks) > 0: + for callback in all_callbacks: + # check if callback is a string - e.g. "lago", "openmeter" + if isinstance(callback, str): + callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( + callback, + internal_usage_cache=None, + llm_router=None, + ) + if callback is None or any( + type(cb) is type(callback) for cb in litellm._async_success_callback + ): # don't double add a callback + continue + if callback not in litellm.input_callback: + litellm.input_callback.append(callback) + if callback not in litellm.success_callback: + litellm.logging_callback_manager.add_litellm_success_callback(callback) + if callback not in litellm.failure_callback: + litellm.logging_callback_manager.add_litellm_failure_callback(callback) + if callback not in litellm._async_success_callback: + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + if callback not in litellm._async_failure_callback: + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + print_verbose( + f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}" + ) - if ( - len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0 - ) and len(callback_list) == 0: - callback_list = list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) + callbacks_to_set_up: Final = ( + list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) + if len(callback_list) == 0 + else [] + ) + if callbacks_to_set_up: + callback_list = callbacks_to_set_up + ## ASYNC CALLBACKS - safety net for callbacks added via direct append + if len(litellm.input_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.input_callback): + if coroutine_checker.is_async_callable(callback): + litellm._async_input_callback.append(callback) + removed_async_items.append(index) + + # Pop the async items from input_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.input_callback.pop(index) + if len(litellm.success_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.success_callback): + if coroutine_checker.is_async_callable(callback): + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + removed_async_items.append(index) + elif callback == "dynamodb" or callback == "openmeter": + # dynamo is an async callback, it's used for the proxy and needs to be async + # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + removed_async_items.append(index) + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): + add_custom_logger_callback_to_specific_event(callback, "success") + + # Pop the async items from success_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.success_callback.pop(index) + + if len(litellm.failure_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.failure_callback): + if coroutine_checker.is_async_callable(callback): + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + removed_async_items.append(index) + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): + add_custom_logger_callback_to_specific_event(callback, "failure") + + # Pop the async items from failure_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.failure_callback.pop(index) + if callbacks_to_set_up: get_set_callbacks: Final[Callable[[], Callable[..., None]]] = litellm_utils.get_set_callbacks - get_set_callbacks()(callback_list=callback_list, function_id=function_id) - ## ASYNC CALLBACKS - safety net for callbacks added via direct append - if len(litellm.input_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.input_callback): - if coroutine_checker.is_async_callable(callback): - litellm._async_input_callback.append(callback) - removed_async_items.append(index) - - # Pop the async items from input_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.input_callback.pop(index) - if len(litellm.success_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.success_callback): - if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - removed_async_items.append(index) - elif callback == "dynamodb" or callback == "openmeter": - # dynamo is an async callback, it's used for the proxy and needs to be async - # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - removed_async_items.append(index) - elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): - add_custom_logger_callback_to_specific_event(callback, "success") - - # Pop the async items from success_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.success_callback.pop(index) - - if len(litellm.failure_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.failure_callback): - if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - removed_async_items.append(index) - elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): - add_custom_logger_callback_to_specific_event(callback, "failure") - - # Pop the async items from failure_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.failure_callback.pop(index) + get_set_callbacks()(callback_list=callbacks_to_set_up, function_id=function_id) ### DYNAMIC CALLBACKS ### - dynamic_success_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_async_success_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_async_failure_callbacks: Final[list[str | Callable | CustomLogger] | None] = None + dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = None if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list): removed_async_items = [] for index, callback in enumerate(kwargs["success_callback"]): @@ -1500,6 +1509,7 @@ def _schedule_async_success_logging( return def _enqueue_async_logging() -> None: + logging_obj.mark_async_success_log_enqueued() asyncio.create_task( _client_async_logging_helper( logging_obj=logging_obj, @@ -1578,6 +1588,22 @@ def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: return kwargs.get("timeout") if stream_timeout is None else stream_timeout +def _runs_as_async_success_callback(callback: str | Callable[..., object] | CustomLogger) -> bool: + get_coroutine_checker: Final = litellm_utils.get_coroutine_checker + return get_coroutine_checker().is_async_callable(callback) or callback in ("dynamodb", "s3") + + +def split_dynamic_success_callbacks( + callbacks: Sequence[str | Callable[..., object] | CustomLogger], +) -> tuple[ + tuple[str | Callable[..., object] | CustomLogger, ...], tuple[str | Callable[..., object] | CustomLogger, ...] +]: + return ( + tuple(callback for callback in callbacks if not _runs_as_async_success_callback(callback)), + tuple(callback for callback in callbacks if _runs_as_async_success_callback(callback)), + ) + + def check_coroutine(value) -> bool: get_coroutine_checker: Final = litellm_utils.get_coroutine_checker return get_coroutine_checker().is_async_callable(value) @@ -2165,7 +2191,11 @@ def client(original_function): # Sync logging_obj.stream after deployment hooks (they may convert it). _hook_stream: Final = kwargs.get("stream") - if _hook_stream is not None and logging_obj.stream != _hook_stream: + if ( + _hook_stream is not None + and logging_obj.stream != _hook_stream + and call_type not in _NEVER_STREAMING_CALL_TYPES + ): logging_obj.stream = _hook_stream kwargs["litellm_logging_obj"] = logging_obj @@ -2370,7 +2400,7 @@ def client(original_function): unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) try: logging_obj.failure_handler( - e, traceback_exception, start_time, end_time + e, traceback_exception, start_time, end_time, for_async_request=True ) # DO NOT MAKE THREADED - router retry fallback relies on this! except Exception as e: raise e @@ -2509,6 +2539,15 @@ _STREAMING_CALL_TYPES: Final = frozenset( } ) +_NEVER_STREAMING_CALL_TYPES: Final = frozenset( + { + CallTypes.extraction, + CallTypes.aextraction, + CallTypes.extraction.value, + CallTypes.aextraction.value, + } +) + def _is_streaming_request( kwargs: dict[str, object], @@ -2520,6 +2559,8 @@ def _is_streaming_request( - if "stream=True" in kwargs (litellm chat completion, litellm text completion, litellm messages) - if call_type is generate_content_stream or agenerate_content_stream (litellm google genai) """ + if call_type in _NEVER_STREAMING_CALL_TYPES: + return False if "stream" in kwargs and kwargs["stream"] is True: return True return call_type in _STREAMING_CALL_TYPES @@ -2564,19 +2605,14 @@ def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse: try: result: Final = _return_huggingface_tokenizer(model) - if result is not None: + if result is not None and not isinstance(result, _HuggingFaceTokenizerLoadFailure): return result except Exception as e: from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking if isinstance(e, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): raise - verbose_logger.warning( - "Falling back to tiktoken for %s; token counts may be approximate. " - "For Python Hugging Face tokenization, install tokenizers and huggingface-hub. Error: %s", - json.dumps(redact_secrets(model)), - json.dumps(redact_secrets(str(e))), - ) + _warn_huggingface_tokenizer_fallback(model, str(e)) # default - tiktoken return _return_openai_tokenizer(model) @@ -2616,15 +2652,41 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None: return None -def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None: +@dataclass(frozen=True, slots=True) +class _HuggingFaceTokenizerLoadFailure: + reason: str + + +def _warn_huggingface_tokenizer_fallback(model: str, reason: str) -> None: + verbose_logger.warning( + "Falling back to tiktoken for %s; token counts may be approximate. " + "For Python Hugging Face tokenization, install tokenizers and huggingface-hub. Error: %s", + json.dumps(redact_secrets(model)), + json.dumps(redact_secrets(reason)), + ) + + +def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _HuggingFaceTokenizerLoadFailure | None: kind: Final = huggingface_tokenizer_kind(model) if kind is None: return None + backend: Final = _huggingface_tokenizer_backend() + key: Final = _huggingface_tokenizer_failure_key(kind, backend) with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()): - tokenizer: Final = _load_huggingface_tokenizer(kind, _huggingface_tokenizer_backend()) + cached_failure: Final = _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.get_cache(key) + if isinstance(cached_failure, _HuggingFaceTokenizerLoadFailure): + return cached_failure + tokenizer: Final = _load_huggingface_tokenizer_recording_failure(kind, backend, key) + if isinstance(tokenizer, _HuggingFaceTokenizerLoadFailure): + _warn_huggingface_tokenizer_fallback(model, tokenizer.reason) + return tokenizer return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} +def _huggingface_tokenizer_failure_key(kind: HuggingFaceTokenizerKind, backend: Decision) -> str: + return f"{kind}:{backend.name}" + + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decision) -> HuggingFace: """One tokenizer per kind and backend; `backend` is the cache key, the dispatch re-derives it.""" @@ -2647,6 +2709,21 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio assert_never(kind) +def _load_huggingface_tokenizer_recording_failure( + kind: HuggingFaceTokenizerKind, backend: Decision, key: str +) -> HuggingFace | _HuggingFaceTokenizerLoadFailure: + try: + return _load_huggingface_tokenizer(kind, backend) + except Exception as error: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking + + if isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): + raise + failure: Final = _HuggingFaceTokenizerLoadFailure(reason=str(error)) + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.set_cache(key, failure, ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS) + return failure + + def encode(model="", text="", custom_tokenizer: dict | None = None): """ Encodes the given text using the specified model. @@ -7892,12 +7969,16 @@ class TextCompletionStreamWrapper: text_choices: Final = TextChoices() if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices raise Exception - delta: Final = chunk["choices"][0]["delta"] - text_choices["text"] = delta["content"] - text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] - response["choices"] = [text_choices] + + if chunk["choices"]: + delta: Final = chunk["choices"][0]["delta"] + text_choices["text"] = delta["content"] + text_choices["reasoning_content"] = delta.get("reasoning_content") + text_choices["index"] = chunk["choices"][0]["index"] + text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + response["choices"] = [text_choices] + else: + response["choices"] = [] # only pass usage when stream_options["include_usage"] is True if self.stream_options and self.stream_options.get("include_usage", False) is True: @@ -8262,6 +8343,26 @@ def trim_messages( from litellm.caching.in_memory_cache import InMemoryCache +class _HuggingFaceTokenizerLoadFailureCache(Protocol): + def get_cache(self, key: str) -> object | None: ... + + def set_cache(self, key: str, value: object, ttl: float | None = None) -> None: ... + + def flush_cache(self) -> None: ... + + +_HUGGINGFACE_TOKENIZER_LOAD_FAILURES: Final[_HuggingFaceTokenizerLoadFailureCache] = ( + cast( # cast-ok: cache methods are untyped + _HuggingFaceTokenizerLoadFailureCache, + InMemoryCache( + max_size_in_memory=16, + default_ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, + clock=lambda: time.monotonic(), + ), + ) +) + + class AvailableModelsCache(InMemoryCache): def __init__(self, ttl_seconds: int = 300, max_size: int = 1000): super().__init__(ttl_seconds, max_size) @@ -9538,6 +9639,12 @@ class ProviderConfigManager: return litellm.AzureAIDecisionsConfig() return None + @staticmethod + def get_provider_extraction_config(model: str, provider: LlmProviders) -> BaseExtractionConfig | None: + if provider == LlmProviders.CLOUDFLARE: + return litellm.CloudflareExtractionConfig() + return None + @staticmethod def get_provider_anthropic_messages_config( model: str, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 15eba7f8af8..3d89af9d691 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16580,6 +16580,24 @@ "source": "https://developers.cloudflare.com/workers-ai/platform/pricing/", "supports_vision": true }, + "cloudflare/clairic": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 65536, + "mode": "extraction", + "output_cost_per_token": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/clairic/", + "supports_vision": true + }, + "cloudflare/@cf/cloudflare/clairic": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 65536, + "mode": "extraction", + "output_cost_per_token": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/clairic/", + "supports_vision": true + }, "cloudflare/@cf/meta/llama-2-7b-chat-fp16": { "input_cost_per_token": 1.923e-06, "litellm_provider": "cloudflare", @@ -21903,6 +21921,7 @@ ], "supports_function_calling": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true }, @@ -21927,6 +21946,7 @@ ], "supports_function_calling": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true }, @@ -21958,6 +21978,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true }, @@ -22021,6 +22042,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true }, @@ -29417,6 +29439,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -30139,6 +30162,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -30200,6 +30224,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -30268,6 +30293,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -30317,6 +30343,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": false, "supports_url_context": false, @@ -31743,6 +31770,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -31927,6 +31955,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -31989,6 +32018,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32051,6 +32081,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32420,6 +32451,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32481,6 +32513,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32549,6 +32582,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -32598,6 +32632,7 @@ "supports_reasoning": true, "supports_minimal_reasoning_effort": false, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": false, "supports_url_context": false, @@ -33192,6 +33227,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_vision": true }, "github_copilot/gemini-3.8-flash": { @@ -33210,6 +33246,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_vision": true }, "github_copilot/gpt-3.5-turbo": { @@ -52605,6 +52642,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, @@ -65823,6 +65861,7 @@ "supports_tool_choice": true, "supports_response_schema": true, "supports_reasoning": true, + "supports_sampling_params": false, "supports_vision": true, "source": "https://deepinfra.com/pricing", "supports_audio_input": true @@ -69264,6 +69303,7 @@ "perplexity/google/gemini-3.5-flash-lite": { "litellm_provider": "perplexity", "mode": "responses", + "supports_sampling_params": false, "supports_web_search": true, "supports_function_calling": true, "input_cost_per_token": 3e-07, @@ -69274,6 +69314,7 @@ "perplexity/google/gemini-3.6-flash": { "litellm_provider": "perplexity", "mode": "responses", + "supports_sampling_params": false, "supports_web_search": true, "supports_function_calling": true, "input_cost_per_token": 1.5e-06, @@ -69284,6 +69325,7 @@ "perplexity/google/gemini-3.7-flash": { "litellm_provider": "perplexity", "mode": "responses", + "supports_sampling_params": false, "supports_web_search": true, "supports_function_calling": true, "input_cost_per_token": 7.5e-07, @@ -71375,6 +71417,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_vision": true, "supports_web_search": true }, @@ -71392,6 +71435,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_vision": true, "supports_web_search": true }, @@ -74322,6 +74366,7 @@ "perplexity/google/gemini-3.8-flash": { "litellm_provider": "perplexity", "mode": "responses", + "supports_sampling_params": false, "input_cost_per_token": 7.5e-07, "output_cost_per_token": 3.75e-06, "cache_read_input_token_cost": 7.5e-08, @@ -75064,12 +75109,10 @@ "gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, "cache_read_input_token_cost_batches": 7.5e-08, - "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_image": 0.00168, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, - "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, @@ -75080,7 +75123,6 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, - "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", @@ -75113,12 +75155,10 @@ "vertex_ai/gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, "cache_read_input_token_cost_batches": 7.5e-08, - "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_image": 0.00168, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, - "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, @@ -75129,7 +75169,6 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, - "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", @@ -76929,5 +76968,55 @@ "supported_output_modalities": [ "text" ] + }, + "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash-1017-us": { + "cache_read_input_token_cost": 9e-09, + "cache_read_input_token_cost_priority": 1.125e-08, + "input_cost_per_token": 4.5e-07, + "input_cost_per_token_priority": 5.625e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.8e-06, + "output_cost_per_token_priority": 2.25e-06, + "source": "https://docs.fireworks.ai/serverless/us-only-serverless", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "wandb/reflection/Beam-1.0-501B-A23B": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "wandb/zai-org/GLM-5.3": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 1.2e-06, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false } } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json deleted file mode 100644 index 0e8f4f596c0..00000000000 --- a/model_prices_and_context_window.schema.json +++ /dev/null @@ -1,1284 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "title": "LiteLLM model_prices_and_context_window.json", - "description": "Schema for LiteLLM's model price and context window registry (https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. All costs are USD per unit. New optional fields are added regularly, so consumers should ignore unknown fields rather than reject them.", - "type": "object", - "properties": { - "sample_spec": { - "type": "object", - "description": "Documentation placeholder illustrating the entry shape; not a real model and not schema-conformant (several values are prose)." - }, - "fallback_generalizations": { - "type": "object", - "description": "Regex rules that generalize unknown model ids to known families; not a model entry.", - "properties": { - "rules": { - "type": "array", - "items": { - "type": "object", - "properties": { - "name": { - "type": "string" - }, - "pattern": { - "type": "string" - }, - "description": { - "type": "string" - } - }, - "required": [ - "name", - "pattern" - ], - "additionalProperties": true - } - } - }, - "additionalProperties": false - } - }, - "additionalProperties": { - "$ref": "#/$defs/modelEntry" - }, - "$defs": { - "modelEntry": { - "type": "object", - "description": "Pricing, limits, and capability flags for one model. Fields other than litellm_provider are optional; boolean capability flags are simply omitted when unknown or false.", - "required": [ - "litellm_provider" - ], - "properties": { - "annotation_cost_per_page": { - "type": "number", - "minimum": 0 - }, - "annotation_cost_per_page_batches": { - "type": "number", - "minimum": 0 - }, - "audio_transcription_config": { - "type": "string" - }, - "bedrock_converse_supports_strict_tools": { - "type": "boolean" - }, - "bedrock_output_config_effort_ceiling": { - "type": "string", - "description": "Highest reasoning effort the Bedrock output_config accepts for this model.", - "enum": [ - "low", - "medium", - "high", - "max", - "xhigh" - ] - }, - "cache_creation_input_audio_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost": { - "type": "number", - "minimum": 0, - "description": "USD per token written to the provider's prompt cache." - }, - "cache_creation_input_token_cost_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_100k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_1hr": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_1hr_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_1hr_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_272k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_272k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_272k_tokens_decisions": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_above_272k_tokens_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "cache_creation_input_token_cost_above_272k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_creation_input_token_cost_above_272k_tokens_ultrafast": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_creation_input_token_cost_batches": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost_decisions": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "cache_creation_input_token_cost_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_creation_input_token_cost_ultrafast": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_audio_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_image_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost": { - "type": "number", - "minimum": 0, - "description": "USD per prompt token served from the provider's prompt cache." - }, - "cache_read_input_token_cost_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_100k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_200k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_200k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_above_272k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_272k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_272k_tokens_decisions": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_272k_tokens_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_above_272k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_above_272k_tokens_ultrafast": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_above_512k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "cache_read_input_token_cost_balanced": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost_batches": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost_decisions": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "cache_read_input_token_cost_ultrafast": { - "type": "number", - "minimum": 0 - }, - "citation_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "code_interpreter_cost_per_session": { - "type": "number", - "minimum": 0 - }, - "comment": { - "type": "string" - }, - "cost_per_second": { - "type": "number", - "minimum": 0 - }, - "default_reasoning_effort": { - "type": "string", - "description": "Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'.", - "enum": [ - "none", - "minimal", - "low", - "medium", - "high", - "xhigh" - ] - }, - "deprecation_date": { - "type": "string", - "description": "Date the provider deprecates the model, YYYY-MM-DD.", - "format": "date", - "pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$" - }, - "explicit_cache_read_input_token_cost": { - "type": "number", - "minimum": 0 - }, - "gemini_audio_only_live": { - "type": "boolean" - }, - "gemini_native_audio": { - "type": "boolean" - }, - "google_maps_grounding_cost_per_query": { - "type": "number", - "minimum": 0, - "description": "USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit." - }, - "guardrail_cost_per_unit": { - "type": "object", - "description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).", - "additionalProperties": { - "type": "number", - "minimum": 0 - } - }, - "input_cost_per_audio_per_second": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_audio_per_second_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_audio_token": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_audio_token_batches": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_audio_token_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "input_cost_per_character": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_character_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_image": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_image_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_image_token": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_image_token_batches": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_pixel": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_query": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_request": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_second": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_text_message": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_token": { - "type": "number", - "minimum": 0, - "description": "USD per prompt token." - }, - "input_cost_per_token_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_100k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_200k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_200k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "input_cost_per_token_above_272k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_272k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_272k_tokens_decisions": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_272k_tokens_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "input_cost_per_token_above_272k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "input_cost_per_token_above_272k_tokens_ultrafast": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_above_512k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_token_balanced": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_token_batches": { - "type": "number", - "minimum": 0, - "description": "USD per prompt token via the provider's batch API." - }, - "input_cost_per_token_cache_hit": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_token_decisions": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_token_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "input_cost_per_token_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "input_cost_per_token_ultrafast": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_video_per_second": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_video_per_second_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_video_per_second_above_15s_interval": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_video_per_second_above_8s_interval": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "input_cost_per_video_token": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_video_token_batches": { - "type": "number", - "minimum": 0 - }, - "input_dbu_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "litellm_provider": { - "type": "string", - "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers." - }, - "max_input_tokens": { - "type": "integer", - "minimum": 0, - "description": "Maximum prompt/context tokens the model accepts." - }, - "max_output_tokens": { - "type": "integer", - "minimum": 0, - "description": "Maximum tokens the model can generate in one response." - }, - "max_tokens": { - "type": "integer", - "minimum": 0, - "description": "Legacy field: max output tokens if the provider specifies it, else max input tokens." - }, - "metadata": { - "type": "object", - "description": "Free-form notes about the entry (e.g. pricing derivation)." - }, - "mode": { - "type": "string", - "description": "Primary API surface / task type of the model.", - "enum": [ - "audio_speech", - "audio_transcription", - "chat", - "completion", - "embedding", - "evaluation", - "guardrail", - "image_edit", - "image_generation", - "moderation", - "ocr", - "realtime", - "rerank", - "responses", - "search", - "vector_store", - "video_generation" - ] - }, - "ocr_cost_per_credit": { - "type": "number", - "minimum": 0 - }, - "ocr_cost_per_page": { - "type": "number", - "minimum": 0 - }, - "ocr_cost_per_page_batches": { - "type": "number", - "minimum": 0 - }, - "off_peak_pricing": { - "type": "object", - "description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.", - "properties": { - "hours_utc": { - "description": "UTC \"HH:MM-HH:MM\" window, or a list of them; a window may wrap past midnight.", - "oneOf": [ - { - "type": "string", - "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" - }, - { - "type": "array", - "items": { - "type": "string", - "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" - }, - "minItems": 1 - } - ] - }, - "windows": { - "type": "array", - "items": { - "type": "object", - "properties": { - "hours_utc": { - "description": "UTC \"HH:MM-HH:MM\" window, or a list of them; a window may wrap past midnight.", - "oneOf": [ - { - "type": "string", - "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" - }, - { - "type": "array", - "items": { - "type": "string", - "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" - }, - "minItems": 1 - } - ] - }, - "weekdays": { - "type": "array", - "description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.", - "items": { - "oneOf": [ - { - "type": "integer", - "minimum": 1, - "maximum": 7 - }, - { - "type": "string", - "pattern": "(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" - } - ] - }, - "minItems": 1 - } - }, - "required": [ - "hours_utc" - ], - "additionalProperties": false - }, - "minItems": 1 - }, - "weekday_timezone": { - "type": "string", - "description": "IANA zone the weekdays of each window are read on; defaults to UTC." - }, - "input_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_reasoning_token": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost": { - "type": "number", - "minimum": 0 - } - }, - "anyOf": [ - { - "required": [ - "hours_utc" - ] - }, - { - "required": [ - "windows" - ] - } - ], - "additionalProperties": false - }, - "output_cost_per_audio_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_character": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_character_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_image": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_0.5K": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_1024": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_1536": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_1K": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_2K": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_4K": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_512": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_image_token_batches": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_pixel": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_reasoning_token": { - "type": "number", - "minimum": 0, - "description": "USD per reasoning/thinking token, when billed separately." - }, - "output_cost_per_second": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_1080p": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_2k": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_480p": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_4k": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_720p": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_second_768p": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token": { - "type": "number", - "minimum": 0, - "description": "USD per generated token." - }, - "output_cost_per_token_above_100k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_100k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_128k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_200k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_200k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_200k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "output_cost_per_token_above_272k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_272k_tokens_batches": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_272k_tokens_decisions": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_272k_tokens_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "output_cost_per_token_above_272k_tokens_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "output_cost_per_token_above_272k_tokens_ultrafast": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_above_512k_tokens": { - "type": "number", - "minimum": 0, - "description": "Rate applied once the prompt exceeds the token threshold in the field name." - }, - "output_cost_per_token_balanced": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token_batches": { - "type": "number", - "minimum": 0, - "description": "USD per generated token via the provider's batch API." - }, - "output_cost_per_token_decisions": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token_flex": { - "type": "number", - "minimum": 0, - "description": "Flex service-tier rate for the same-named base field." - }, - "output_cost_per_token_priority": { - "type": "number", - "minimum": 0, - "description": "Priority service-tier rate for the same-named base field." - }, - "output_cost_per_token_ultrafast": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_video_per_second": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_video_token": { - "type": "number", - "minimum": 0 - }, - "output_dbu_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_vector_size": { - "type": "integer", - "minimum": 0, - "description": "Embedding dimension for embedding models." - }, - "prompt_cache_min_tokens": { - "type": "integer", - "minimum": 0, - "description": "Smallest prefix the provider will actually cache; absent means the provider default applies." - }, - "provider_specific_entry": { - "type": "object", - "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema)." - }, - "reasoning_effort_levels": { - "type": "array", - "description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.", - "items": { - "type": "string", - "enum": [ - "none", - "minimal", - "low", - "medium", - "high", - "xhigh", - "max" - ] - } - }, - "regional_endpoint_uplift_multiplier": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%)." - }, - "regional_processing_uplift_multiplier_eu": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%)." - }, - "regional_processing_uplift_multiplier_us": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%)." - }, - "rpm": { - "type": "integer", - "minimum": 0, - "description": "Provider default requests-per-minute limit." - }, - "search_context_cost_per_query": { - "type": "object", - "description": "USD cost per web search query, keyed by search context size.", - "properties": { - "search_context_size_low": { - "type": "number", - "minimum": 0 - }, - "search_context_size_medium": { - "type": "number", - "minimum": 0 - }, - "search_context_size_high": { - "type": "number", - "minimum": 0 - } - }, - "additionalProperties": false - }, - "source": { - "type": "string", - "description": "URL of the provider pricing/model page this entry was taken from." - }, - "supported_audio_formats": { - "type": "array", - "description": "Audio container formats the model can return.", - "items": { - "type": "string", - "enum": [ - "mp3", - "wav" - ] - } - }, - "supported_endpoints": { - "type": "array", - "description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.", - "items": { - "type": "string" - } - }, - "supported_modalities": { - "type": "array", - "description": "Input modalities the model accepts.", - "items": { - "type": "string", - "enum": [ - "text", - "image", - "audio", - "video" - ] - } - }, - "supported_output_modalities": { - "type": "array", - "description": "Output modalities the model can produce.", - "items": { - "type": "string", - "enum": [ - "text", - "image", - "audio", - "video", - "code" - ] - } - }, - "supported_regions": { - "type": "array", - "description": "Cloud regions the model is available in ('global' or region ids).", - "items": { - "type": "string" - } - }, - "supports_adaptive_thinking": { - "type": "boolean" - }, - "supports_anthropic_compaction": { - "type": "boolean" - }, - "supports_anthropic_thinking_payload": { - "type": "boolean" - }, - "supports_assistant_prefill": { - "type": "boolean" - }, - "supports_audio_input": { - "type": "boolean" - }, - "supports_audio_output": { - "type": "boolean" - }, - "supports_bedrock_converse_tool_result_images": { - "type": "boolean" - }, - "supports_bedrock_native_structured_output": { - "type": "boolean" - }, - "supports_bedrock_runtime_chat_completions_inline_reasoning": { - "type": "boolean" - }, - "supports_bedrock_runtime_chat_completions_response_format": { - "type": "boolean" - }, - "supports_bedrock_runtime_chat_completions_tools_with_reasoning": { - "type": "boolean" - }, - "supports_box_questions": { - "type": "boolean" - }, - "supports_computer_use": { - "type": "boolean" - }, - "supports_embedding_image_input": { - "type": "boolean" - }, - "supports_fast_mode": { - "type": "boolean" - }, - "supports_forced_tool_use": { - "type": "boolean" - }, - "supports_function_calling": { - "type": "boolean" - }, - "supports_image_input": { - "type": "boolean" - }, - "supports_image_size": { - "type": "boolean" - }, - "supports_legacy_thinking": { - "type": "boolean" - }, - "supports_low_reasoning_effort": { - "type": "boolean" - }, - "supports_max_reasoning_effort": { - "type": "boolean" - }, - "supports_mid_conversation_system": { - "type": "boolean" - }, - "supports_minimal_reasoning_effort": { - "type": "boolean" - }, - "supports_multimodal": { - "type": "boolean" - }, - "supports_native_streaming": { - "type": "boolean" - }, - "supports_native_structured_output": { - "type": "boolean" - }, - "supports_none_reasoning_effort": { - "type": "boolean" - }, - "supports_nova_canvas_image_edit": { - "type": "boolean" - }, - "supports_output_config": { - "type": "boolean" - }, - "supports_parallel_function_calling": { - "type": "boolean" - }, - "supports_parallel_tool_use_config": { - "type": "boolean" - }, - "supports_pdf_input": { - "type": "boolean" - }, - "supports_prompt_cache_breakpoint": { - "type": "boolean" - }, - "supports_prompt_caching": { - "type": "boolean" - }, - "supports_reasoning": { - "type": "boolean" - }, - "supports_regex_lookaround": { - "type": "boolean" - }, - "supports_response_schema": { - "type": "boolean" - }, - "supports_sampling_params": { - "type": "boolean" - }, - "supports_speed": { - "type": "boolean" - }, - "supports_system_messages": { - "type": "boolean" - }, - "supports_thinking_cache_preservation": { - "type": "boolean" - }, - "supports_tool_choice": { - "type": "boolean" - }, - "supports_tool_search": { - "type": "boolean" - }, - "supports_url_context": { - "type": "boolean" - }, - "supports_video_input": { - "type": "boolean" - }, - "supports_vision": { - "type": "boolean" - }, - "supports_web_search": { - "type": "boolean" - }, - "supports_xhigh_reasoning_effort": { - "type": "boolean" - }, - "thinking_always_on": { - "type": "boolean" - }, - "tiered_pricing": { - "type": "array", - "description": "Context-length or result-count tiered rates; each tier's costs apply within its range.", - "items": { - "type": "object", - "properties": { - "range": { - "type": "array", - "description": "[min, max] prompt-token span this tier applies to.", - "items": { - "type": "number", - "minimum": 0 - }, - "minItems": 2, - "maxItems": 2 - }, - "max_results_range": { - "type": "array", - "description": "[min, max] result-count span this tier applies to (search models).", - "items": { - "type": "number", - "minimum": 0 - }, - "minItems": 2, - "maxItems": 2 - }, - "input_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_token": { - "type": "number", - "minimum": 0 - }, - "output_cost_per_reasoning_token": { - "type": "number", - "minimum": 0 - }, - "cache_read_input_token_cost": { - "type": "number", - "minimum": 0 - }, - "cache_creation_input_token_cost": { - "type": "number", - "minimum": 0 - }, - "input_cost_per_query": { - "type": "number", - "minimum": 0 - } - }, - "additionalProperties": false - } - }, - "token_count_framing": { - "type": "integer", - "minimum": 0, - "description": "Tokens added once per request to a local count estimate, for the provider's chat framing." - }, - "token_count_model": { - "type": "object", - "required": [ - "model" - ], - "additionalProperties": false, - "properties": { - "model": { - "type": "string", - "pattern": "\\S" - }, - "region": { - "type": "string", - "pattern": "\\S" - }, - "litellm_provider": { - "type": "string", - "enum": [ - "bedrock", - "bedrock_mantle" - ] - } - }, - "description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in." - }, - "token_count_tokenizer": { - "type": "string", - "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", - "enum": [ - "kimi_k3" - ] - }, - "tpm": { - "type": "integer", - "minimum": 0, - "description": "Provider default tokens-per-minute limit." - }, - "use_openai_responses_path": { - "type": "boolean" - }, - "uses_embed_content": { - "type": "boolean" - }, - "vertex_ai_audio_api": { - "type": "string", - "enum": [ - "lyria_predict", - "lyria_interactions" - ] - }, - "web_search_billing_unit": { - "type": "string", - "description": "Whether web search is billed per query or per prompt.", - "enum": [ - "per_query", - "per_prompt" - ] - } - }, - "additionalProperties": true - } - } -} diff --git a/osv-scanner.toml b/osv-scanner.toml index 60f163392a8..bb802137f6e 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -7,3 +7,16 @@ reason = "diskcache has no fixed release published; remove this entry once one e id = "GHSA-vfj7-8cjw-p6xm" ignoreUntil = 2026-11-03 reason = "braces has no fixed release published (3.0.3, 2024-05-21, is latest and last_affected); dev-only via knip > fast-glob > micromatch; remove this entry once one exists" + +[[IgnoredVulns]] +id = "GHSA-6q48-6c7w-9gw2" +ignoreUntil = 2026-10-23 +reason = "certifi 2026.10.19 (published 2026-10-20) is inside the 3-day uv exclude-newer window until 2026-10-23 02:24 UTC; bump certifi and remove this entry once it clears" + +[[PackageOverrides]] +name = "pillow" +version = "12.3.0" +ecosystem = "PyPI" +vulnerability.ignore = true +effectiveUntil = 2026-10-27 +reason = "pillow 13.0.0 requires Python >= 3.11, so the Python 3.10 resolution of the ci and dev groups stays on 12.3.0; 3.11+ is on 13.0.0" diff --git a/pyproject.toml b/pyproject.toml index 9871a83b421..9076760153a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.107.0" +version = "1.108.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -75,11 +75,11 @@ proxy = [ "boto3>=1.43.1,<2.0", "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", - "mcp>=2.2.0,<3", - "httpx2>=2.5.0,<3", + "mcp>=2.3.0,<3", + "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3", - "litellm-proxy-extras==0.4.110", - "litellm-enterprise==0.1.78", + "litellm-proxy-extras==0.4.111", + "litellm-enterprise==0.1.79", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -125,7 +125,7 @@ utils = [ "numpydoc>=1.8.0,<2.0", ] caching = ["diskcache>=5.6.3,<6.0"] -mcp = ["mcp>=2.2.0,<3", "httpx2>=2.5.0,<3", "pydantic>=2.12.0,<3"] +mcp = ["mcp>=2.3.0,<3", "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3"] # Driver for the MongoDB Atlas vector store; Atlas Vector Search has no HTTP query API. # The floor is 4.9 because that is the release AsyncMongoClient landed in. # SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels @@ -243,7 +243,7 @@ e2e-dev = [ "locust==2.45.0", "anthropic==0.84.0", "psutil==7.2.2", - "mcp>=2.2.0,<3", + "mcp>=2.3.0,<3", ] admin-mcp = [ "litellm-admin-mcp @ https://github.com/BerriAI/liteadmin-mcp/archive/d35ec9c19c117d4c50cc1eccf6ce1296aac25a1a.tar.gz#sha256=86b930f6706fb2da10d53d1798ee7fd14e44fb0afdef0da122cf0e6bb88fa2c3 ; python_version >= '3.12'", @@ -269,7 +269,8 @@ ci = [ # protobuf, Pillow is a compiled C extension). "tenacity==8.5.0", "google-generativeai==0.8.6", - "Pillow==12.3.0", + "Pillow==12.3.0; python_version == '3.10'", + "Pillow==13.0.0; python_version >= '3.11'", # Azure batch E2E tests still import psycopg2 directly. "psycopg2-binary==2.9.11", "pytest-codspeed==4.3.0", @@ -303,7 +304,7 @@ healthcheck = [ benchmarks = [ "pytest==9.0.3", "pytest-codspeed==4.3.0", - "mcp>=2.2.0,<3", + "mcp>=2.3.0,<3", "a2a-sdk==1.1.0", ] @@ -363,6 +364,7 @@ override-dependencies = [ # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0. "packaging>=24.0", "cryptography>=50.0.0,<51.0", + "mcp>=2.3.0,<3", ] default-groups = ["dev"] required-version = ">=0.10.9" @@ -376,7 +378,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.107.0" +version = "1.108.0" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/benchmarks/test_extraction_guardrail_benchmarks.py b/tests/benchmarks/test_extraction_guardrail_benchmarks.py new file mode 100644 index 00000000000..0d8a3db42ce --- /dev/null +++ b/tests/benchmarks/test_extraction_guardrail_benchmarks.py @@ -0,0 +1,61 @@ +import asyncio +from typing import Final, Literal + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.cloudflare.extraction.guardrail_translation.handler import ExtractionHandler +from litellm.types.extraction import ExtractionResponse +from litellm.types.utils import GenericGuardrailAPIInputs + +MANY_FIELDS: Final = 2_000 + + +class _MaskingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="masking") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs(texts=[text.replace("secret", "[MASKED]") for text in inputs.get("texts", [])]) + + +@pytest.mark.benchmark +def test_extraction_request_guardrail_over_many_field_instructions() -> None: + data: Final = { + "model": "invoice-reader", + "fields": { + f"field_{index}": {"type": "text", "instructions": f"secret {index}"} for index in range(MANY_FIELDS) + }, + } + + result: Final = asyncio.run(ExtractionHandler().process_input_messages(data, _MaskingGuardrail())) + + assert result["fields"] == { + f"field_{index}": {"type": "text", "instructions": f"[MASKED] {index}"} for index in range(MANY_FIELDS) + } + + +@pytest.mark.benchmark +def test_extraction_response_guardrail_over_many_text_answers() -> None: + response: Final = ExtractionResponse.model_validate( + { + "model": "invoice-reader", + "answers": {f"field_{index}": {"type": "text", "value": f"secret {index}"} for index in range(MANY_FIELDS)}, + "truncated": False, + "usage": {"input_tokens": 52, "output_tokens": 0}, + } + ) + + result: Final = asyncio.run(ExtractionHandler().process_output_response(response, _MaskingGuardrail())) + + assert isinstance(result, ExtractionResponse) + assert {name: answer.value for name, answer in result.answers.items()} == { + f"field_{index}": f"[MASKED] {index}" for index in range(MANY_FIELDS) + } diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 7809a23e955..af2f386a362 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -93,6 +93,7 @@ ignored_function_names = [ "_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name) "_wait_for_scheduler_turn", # Tested through prioritized acompletion and atext_completion in test_router.py "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 "_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 @@ -145,6 +146,18 @@ ignored_function_names = [ "async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py "_deployment_as_configured_now", # Read by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py "_admit_turn_through_routing_strategy", # Run by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py + "_callback_selectors", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "bind_to_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py + "serve_through_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py + "_request_adaptive_router_hooks", # Run by serve_through_request_logging, tested through the user_config adaptive test + "_release_global_callbacks", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "_routing_callbacks", # Read by the pre-call checks, tested through the user_config limit tests in test_route_llm_request.py + "_request_policies_not_set_globally", # Read by bind_to_request_logging, tested through the user_config budget test in test_route_llm_request.py + "_optional_loggers", # Read by discard and serve_through_request_logging, tested through the budget limiter discard test + "retire_routing_state", # Run by discard, tested through the budget limiter discard test + "_slack_alerting_loggers", # Run by discard, tested through the Slack alerting discard test in test_router_helper_utils.py + "_adaptive_router_hooks", # Run by discard, tested through the adaptive router discard tests in test_router_helper_utils.py + "_global_callback_candidates", # Run by discard, tested through the survivor discard tests in test_router_helper_utils.py "without_deployments_that_cannot_batch", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py "batch_create_routing", # Tested through the mixed-group acreate_batch tests in test_router.py "is_batch_create_routing", # Tested through the chat-after-batch-create test in test_router.py diff --git a/tests/code_coverage_tests/unbounded_in_baseline.txt b/tests/code_coverage_tests/unbounded_in_baseline.txt index c7e73bf921c..1286ea25c4e 100644 --- a/tests/code_coverage_tests/unbounded_in_baseline.txt +++ b/tests/code_coverage_tests/unbounded_in_baseline.txt @@ -80,14 +80,12 @@ litellm/proxy/management_endpoints/team_endpoints.py _get_keys_count_by_team pri litellm/proxy/management_endpoints/team_endpoints.py _hydrate_member_user_details prisma user_id.in `sorted(user_ids)` 0 litellm/proxy/management_endpoints/team_endpoints.py _resolve_existing_member_user_ids prisma user_id.in `sorted(requested_user_ids)` 0 litellm/proxy/management_endpoints/team_endpoints.py resolve_team_daily_activity_scope prisma team_id.in `list(team_ids_list)` 0 -litellm/proxy/management_endpoints/team_endpoints.py _sweep_deleted_team_references prisma team_id.in `tuple(team_ids)` 0 litellm/proxy/management_endpoints/team_endpoints.py _sweep_deleted_team_references_tx prisma team_id.in `tuple(team_ids)` 0 litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(addressed_user_ids)` 0 litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(user_ids_to_delete)` 0 litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(user_ids_to_delete)` 1 litellm/proxy/management_endpoints/team_endpoints.py _team_user_spend_sql raw-sql sl.team_id.IN `IN ({team_placeholders})` 0 -litellm/proxy/management_endpoints/team_endpoints.py delete_team prisma team_id.in `data.team_ids` 0 -litellm/proxy/management_endpoints/team_endpoints.py delete_team prisma team_id.in `data.team_ids` 1 +litellm/proxy/management_endpoints/team_endpoints.py delete_validated_teams prisma team_id.in `team_ids` 0 litellm/proxy/management_endpoints/team_endpoints.py get_all_team_memberships prisma team_id.in `team_ids` 0 litellm/proxy/management_endpoints/team_endpoints.py list_available_teams prisma team_id.in `available_teams` 0 litellm/proxy/management_endpoints/tool_management_endpoints.py get_tool_spend prisma tool_name.in `[row.tool_name for row in top_tools]` 0 diff --git a/tests/e2e/gateway/perf_ci_config.yml b/tests/e2e/gateway/perf_ci_config.yml index 7ab6637e985..767db313847 100644 --- a/tests/e2e/gateway/perf_ci_config.yml +++ b/tests/e2e/gateway/perf_ci_config.yml @@ -46,6 +46,11 @@ model_list: model: openai/gpt-6-luna api_base: os.environ/E2E_PERF_MOCK_OPENAI_URL api_key: perf-mock-not-checked + - model_name: gpt-6-luna-responses + litellm_params: + model: openai/responses/gpt-6-luna + api_base: os.environ/E2E_PERF_MOCK_OPENAI_URL + api_key: perf-mock-not-checked - model_name: claude-sonnet-5 litellm_params: model: anthropic/claude-sonnet-5 @@ -58,6 +63,13 @@ model_list: aws_region_name: us-east-1 aws_access_key_id: perf-mock-not-checked aws_secret_access_key: perf-mock-not-checked + - model_name: claude-sonnet-5-bedrock-converse + litellm_params: + model: bedrock/converse/anthropic.claude-sonnet-5 + aws_bedrock_runtime_endpoint: os.environ/E2E_PERF_MOCK_BEDROCK_URL + aws_region_name: us-east-1 + aws_access_key_id: perf-mock-not-checked + aws_secret_access_key: perf-mock-not-checked guardrails: - guardrail_name: perf-pre litellm_params: diff --git a/tests/e2e/perf/cases/chat_completions_2rps.py b/tests/e2e/perf/cases/chat_completions_2rps.py new file mode 100644 index 00000000000..9792d99c14f --- /dev/null +++ b/tests/e2e/perf/cases/chat_completions_2rps.py @@ -0,0 +1,15 @@ +"""`/v1/chat/completions` at 2 rps: `gpt-6-luna` is served by the OpenAI mock's chat completions +route, `gpt-6-luna-responses` (`openai/responses/gpt-6-luna`) is bridged to its Responses route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "chat_completions_2rps", + "chat_completions", + direct="gpt-6-luna", + bridge="gpt-6-luna-responses", +) diff --git a/tests/e2e/perf/cases/direct_vs_bridge.py b/tests/e2e/perf/cases/direct_vs_bridge.py new file mode 100644 index 00000000000..13f609c8299 --- /dev/null +++ b/tests/e2e/perf/cases/direct_vs_bridge.py @@ -0,0 +1,31 @@ +"""The mixed case's traffic on one public API, half of it on a deployment the gateway serves on +that provider's own API (direct) and half on a deployment it has to translate to another provider +API (bridge). Same 2 rps, request count, shapes, stream share, callers, guardrails and limits as +`mixed_size_requests_2rps`, so a gap between the two halves and the mixed baseline is the +translation's cost.""" + +from __future__ import annotations + +from dataclasses import replace + +from setup.performance_test_case import Endpoint, PerformanceTestCase, VirtualKey + +from .mixed_size_requests_2rps import CASE as MIXED + + +def _scoped(key: VirtualKey, direct: str) -> VirtualKey: + return replace(key, models=(direct,)) if key.models else key + + +def direct_vs_bridge(name: str, endpoint: Endpoint, *, direct: str, bridge: str) -> PerformanceTestCase: + return replace( + MIXED, + name=name, + teams=( + replace(MIXED.team("team-a"), models=(direct, bridge)), + replace(MIXED.team("team-b"), models=(bridge,)), + ), + virtual_keys=tuple(_scoped(key, direct) for key in MIXED.virtual_keys), + model_groups=(direct, bridge), + traffic=replace(MIXED.traffic, endpoints={endpoint: 1.0}), + ) diff --git a/tests/e2e/perf/cases/event_loop_stalls.py b/tests/e2e/perf/cases/event_loop_stalls.py index 0d67a98257f..8a5654cbf7b 100644 --- a/tests/e2e/perf/cases/event_loop_stalls.py +++ b/tests/e2e/perf/cases/event_loop_stalls.py @@ -8,8 +8,8 @@ new hook that deep-copies the request inline all land here as a new row. Runs on `PYTHONASYNCIODEBUG=1` and `LITELLM_EVENT_LOOP=asyncio` (uvloop reports nothing and crashed under debug locally), and the runner `E2E_PERF_EVENT_LOOP=1`, so the nightly build, which sets neither, skips it. The allowlist counts are twice the highest count any of ten back-to-back local runs showed on that row (rows that showed nothing in -those ten but did in the earlier before/after legs for 46925 and 46967 get 2), so each run sits at or under half -of every cap and a listed location doubling in frequency is a failure, not a flake. The +those ten but did in the earlier before/after legs for 46925 and 46967, or in the first scheduled-lane run, get 2), +so each run sits at or under half of every cap and a listed location doubling in frequency is a failure, not a flake. The `garbage collection` bucket is the gen-2 passes the collector logs in debug mode (0.6 to 0.7 s each, one every minute or two of this traffic, plus the ramp while the heap fills on a local whole-file run), which asyncio would otherwise charge to a random task; its count is the signal for a change that makes much more garbage.""" @@ -44,6 +44,7 @@ CASE: Final = replace( KnownStall("run_coroutine_job", max_count=2), KnownStall("user_api_key_auth.py", max_count=2), KnownStall("route_streaming_logging_to_handler", max_count=2), + KnownStall("TaskHandle._run_coro", max_count=2), ), ), ) diff --git a/tests/e2e/perf/cases/messages_2rps.py b/tests/e2e/perf/cases/messages_2rps.py new file mode 100644 index 00000000000..df5d221d430 --- /dev/null +++ b/tests/e2e/perf/cases/messages_2rps.py @@ -0,0 +1,16 @@ +"""`/v1/messages` at 2 rps: `claude-sonnet-5` is served by the Anthropic mock's Messages route, +`claude-sonnet-5-bedrock-converse` (`bedrock/converse/anthropic.claude-sonnet-5`) is bridged to the +Bedrock mock's Converse route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "messages_2rps", + "messages", + direct="claude-sonnet-5", + bridge="claude-sonnet-5-bedrock-converse", +) diff --git a/tests/e2e/perf/cases/responses_2rps.py b/tests/e2e/perf/cases/responses_2rps.py new file mode 100644 index 00000000000..dd55d536eb6 --- /dev/null +++ b/tests/e2e/perf/cases/responses_2rps.py @@ -0,0 +1,15 @@ +"""`/v1/responses` at 2 rps: `gpt-6-luna` is served by the OpenAI mock's Responses route, +`claude-sonnet-5` is bridged to the Anthropic mock's Messages route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "responses_2rps", + "responses", + direct="gpt-6-luna", + bridge="claude-sonnet-5", +) diff --git a/tests/e2e/perf/setup/event_loop.py b/tests/e2e/perf/setup/event_loop.py index dee45c374d9..58144424544 100644 --- a/tests/e2e/perf/setup/event_loop.py +++ b/tests/e2e/perf/setup/event_loop.py @@ -14,10 +14,11 @@ a hold that such a pass accounts for is reported under `garbage collection` inst from __future__ import annotations +import json import os import re from collections.abc import Iterable -from dataclasses import dataclass, replace +from dataclasses import asdict, dataclass, replace from datetime import datetime from itertools import groupby from pathlib import Path @@ -79,11 +80,24 @@ class StallRow: total_seconds: float max_seconds: float allowed_by: KnownStall | None + """Every hold in the row, longest first.""" + holds: tuple[float, ...] = () @property def offending(self) -> bool: return self.allowed_by is None or self.count > self.allowed_by.max_count + @property + def over_allowance(self) -> int: + return self.count if self.allowed_by is None else max(0, self.count - self.allowed_by.max_count) + + @property + def note(self) -> str: + allowed: Final = ( + f"{self.count} of {self.allowed_by.max_count} allowed" if self.allowed_by else "not on the allowlist" + ) + return f"{allowed}, max {self.max_seconds:.2f} s" + def _short(path: str) -> str: if "site-packages/" in path: @@ -160,6 +174,7 @@ def summarize(spec: EventLoopStalls, stalls: Iterable[Stall]) -> tuple[StallRow, total_seconds=sum(stall.seconds for stall in group), max_seconds=max(stall.seconds for stall in group), allowed_by=_allowed_by(spec, coroutine, where), + holds=tuple(sorted((stall.seconds for stall in group), reverse=True)), ) for (coroutine, where), group in ( (key, tuple(items)) for key, items in groupby(over, key=lambda s: (s.coroutine, s.where)) @@ -181,6 +196,68 @@ def offenders(rows: Iterable[StallRow]) -> tuple[str, ...]: ) +TOP_HOLDS: Final = 10 + + +@dataclass(frozen=True, slots=True) +class Slice: + """One `_breakdown` row in the shape the other metrics use (label, value, share) plus the allowlist verdict the + dashboard shows under the label and colours by. `share` is None where the rows are not parts of a total.""" + + label: str + value: float + share: float | None + note: str | None + over: bool + + +def _label(row: StallRow) -> str: + return f"{row.coroutine} {row.where}".strip() + + +def _slice(row: StallRow, value: float, total: float) -> Slice: + return Slice(label=_label(row), value=value, share=value / total, note=row.note, over=row.offending) + + +def _breakdown(name: str, slices: Iterable[Slice]) -> tuple[tuple[str, str], ...]: + rows: Final = tuple(asdict(row) for row in slices) + return ((f"perf_event_loop_{name}_breakdown", json.dumps(rows)),) if rows else () + + +def _longest_holds(table: tuple[StallRow, ...]) -> tuple[Slice, ...]: + holds: Final = sorted(((seconds, row) for row in table for seconds in row.holds), key=lambda hold: -hold[0]) + return tuple( + Slice(label=_label(row), value=seconds, share=None, note=None, over=row.offending) + for seconds, row in holds[:TOP_HOLDS] + ) + + +def junit_properties(spec: EventLoopStalls, rows: Iterable[StallRow]) -> tuple[tuple[str, float | str], ...]: + """The stall table as `perf_event_loop_*` properties for the dashboard: every hold over the threshold with the + allowlist's total as its limit, the holds past their allowance (limit 0, the case's one assertion), the longest + and the summed hold, each with a `_breakdown` the dashboard opens on click: holds and seconds per location, the + locations past their cap by how far over they are, and the ten longest single holds with where they landed.""" + table: Final = tuple(rows) + holds: Final = sum(row.count for row in table) + offending: Final = sum(row.over_allowance for row in table) + total_seconds: Final = sum(row.total_seconds for row in table) + return ( + ("perf_event_loop_holds", holds), + ("perf_event_loop_holds_limit", sum(known.max_count for known in spec.allowlist)), + ("perf_event_loop_offending_holds", offending), + ("perf_event_loop_offending_holds_limit", 0), + ("perf_event_loop_hold_seconds_max", max((row.max_seconds for row in table), default=0.0)), + ("perf_event_loop_hold_seconds_total", total_seconds), + *_breakdown("holds", (_slice(row, row.count, holds) for row in table)), + *_breakdown( + "offending_holds", + (_slice(row, row.over_allowance, offending) for row in table if row.over_allowance), + ), + *_breakdown("hold_seconds_total", (_slice(row, row.total_seconds, total_seconds) for row in table)), + *_breakdown("hold_seconds_max", _longest_holds(table)), + ) + + def render_stalls(spec: EventLoopStalls, rows: Iterable[StallRow]) -> str: header: Final = ( f" event-loop holds >= {spec.max_hold_seconds:.2f} s in the gateway log (count total_s max_s where)" diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index ccdae6398a9..1532a78606e 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -565,6 +565,7 @@ def _diagnostics(observation: Observation) -> list[str]: return [ " diagnostics (not limits):", f" bodies off reply {observation.bodies_off_reply()}", + f" mock routes {_mock_routes(observation)}", f" guardrail calls per request {len(observation.guardrail_calls) / count:.3f}", f" guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}", f" retried share {len(observation.retried_markers) / count:.4f}", @@ -575,6 +576,12 @@ def _diagnostics(observation: Observation) -> list[str]: ] +def _mock_routes(observation: Observation) -> str: + """Which provider route each request landed on, so a direct-vs-bridge case shows its split.""" + counts: Final = Counter(f"{seen.provider}/{seen.route}" for seen in observation.mock_requests) + return ", ".join(f"{route} {count}" for route, count in sorted(counts.items())) or "n/a" + + def _datastore_lines(observation: Observation) -> list[str]: """What the per-request Postgres and Redis counts are made of, so a red run names the statement or command that grew. `plumbing` (BEGIN, COMMIT, session settings) and `unclassified` count in neither limit.""" diff --git a/tests/e2e/perf/setup/ui_journeys.py b/tests/e2e/perf/setup/ui_journeys.py index 2b7bda77733..4de97d8c717 100644 --- a/tests/e2e/perf/setup/ui_journeys.py +++ b/tests/e2e/perf/setup/ui_journeys.py @@ -10,8 +10,8 @@ from dataclasses import dataclass from datetime import datetime from itertools import chain from types import MappingProxyType -from urllib.parse import urlsplit from typing import Final +from urllib.parse import urlsplit from playwright.sync_api import Error as PlaywrightError from playwright.sync_api import Locator, Page, Response, expect, sync_playwright @@ -189,7 +189,9 @@ def _login(page: Page, dashboard: Dashboard) -> None: def _is_data_fetch(response: Response) -> bool: - return response.request.resource_type in ("fetch", "xhr") and "/_next/" not in response.url + """API calls only. The Admin UI's own assets and Next.js prefetches live under /ui/ and /_next/.""" + path: Final = urlsplit(response.url).path + return response.request.resource_type in ("fetch", "xhr") and not path.startswith(("/ui/", "/_next/")) def _worst(responses: list[Response]) -> Response | None: diff --git a/tests/e2e/perf/setup/ui_query_fixture.py b/tests/e2e/perf/setup/ui_query_fixture.py index cfffa0822e8..ed941a096bd 100644 --- a/tests/e2e/perf/setup/ui_query_fixture.py +++ b/tests/e2e/perf/setup/ui_query_fixture.py @@ -3,7 +3,9 @@ from __future__ import annotations import hashlib import os import time +from collections.abc import Iterable from datetime import datetime, timedelta +from itertools import chain from typing import Final import psycopg @@ -18,6 +20,33 @@ class UiFixture(BaseModel): anchor: AwareDatetime api_key: str + """`perf_ui_db_*` JUnit properties naming how much data the journeys ran against.""" + volume: tuple[tuple[str, int], ...] = () + + +class TableVolume(BaseModel): + model_config = ConfigDict(frozen=True) + + table: str + rows: int + bytes: int + + +QUERIED_TABLES: Final = ( + "LiteLLM_SpendLogs", + "LiteLLM_DailyUserSpend", + "LiteLLM_DailyTeamSpend", + "LiteLLM_DailyTagSpend", + "LiteLLM_VerificationToken", + "LiteLLM_TeamTable", + "LiteLLM_UserTable", + "LiteLLM_ProxyModelTable", +) +_VOLUME_SQL: Final = """ + SELECT c.relname, GREATEST(c.reltuples, 0)::bigint, pg_total_relation_size(c.oid) + FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' AND c.relname = ANY(%(tables)s) +""" SEEDED_TEAMS: Final = 100 @@ -59,6 +88,30 @@ _SEED_SQL: Final = ( ) +def volume_properties(volumes: Iterable[TableVolume]) -> tuple[tuple[str, int], ...]: + """One `perf_ui_db_rows_` and one `perf_ui_db_bytes_
` property per table the journeys read, + so the dashboard can say what the latencies were measured against.""" + return tuple( + chain.from_iterable( + ((f"perf_ui_db_rows_{v.table}", v.rows), (f"perf_ui_db_bytes_{v.table}", v.bytes)) for v in volumes + ) + ) + + +def table_volumes(connection: psycopg.Connection, spend_log_rows: int) -> tuple[TableVolume, ...]: + """Planner row estimates and on-disk size (indexes included) of the queried tables. The restore's + `LiteLLM_SpendLogs` count comes from the seed checkpoint, exact where the estimate can be stale.""" + adapter: Final = TypeAdapter(tuple[str, int, int]) + found: Final = tuple( + adapter.validate_python(row) + for row in connection.execute(_VOLUME_SQL, {"tables": list(QUERIED_TABLES)}).fetchall() + ) + return tuple( + TableVolume(table=table, rows=spend_log_rows if table == "LiteLLM_SpendLogs" else rows, bytes=size) + for table, rows, size in sorted(found) + ) + + def prepare_fixture() -> UiFixture: validate_target(os.environ["DATABASE_URL"]) _finish_upgrade() @@ -71,6 +124,7 @@ def prepare_fixture() -> UiFixture: anchor=TypeAdapter(datetime).validate_python(row[0]), api_key=hashlib.md5(b"perf-key:0", usedforsecurity=False).hexdigest() + hashlib.md5(b"perf-key-tail:0", usedforsecurity=False).hexdigest(), + volume=volume_properties(table_volumes(connection, TypeAdapter(int).validate_python(row[1]))), ) return fixture diff --git a/tests/e2e/perf/setup/upgrade.py b/tests/e2e/perf/setup/upgrade.py index 37e5c7bc904..36efc16e1fe 100644 --- a/tests/e2e/perf/setup/upgrade.py +++ b/tests/e2e/perf/setup/upgrade.py @@ -199,6 +199,18 @@ class UpgradeMonitor: assert not state.pending_indexes or self.overlap_writes > 0, "No persisted requests observed during index work" +def upgrade_properties(state: UpgradeState, monitor: UpgradeMonitor) -> tuple[tuple[str, float], ...]: + """The upgrade's own numbers under the `perf_` prefix the dashboard ingests, next to the perf rows.""" + return ( + ("perf_upgrade_schema_seconds", state.schema_seconds), + ("perf_upgrade_index_seconds", state.index_seconds), + ("perf_upgrade_overlap_writes", float(monitor.overlap_writes)), + ("perf_upgrade_max_lock_waiters", float(monitor.max_lock_waiters)), + ("perf_upgrade_pending_indexes", float(len(state.pending_indexes))), + ("perf_upgrade_unmatched_requests", float(len(monitor.pending))), + ) + + def monitor_for(connection: psycopg.Connection[tuple[object, ...]], urls: tuple[str, ...]) -> UpgradeMonitor: return UpgradeMonitor( state=lambda: read_state(connection), diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py index 8079aa0274d..fba6af30014 100644 --- a/tests/e2e/perf/test_database_upgrade.py +++ b/tests/e2e/perf/test_database_upgrade.py @@ -12,19 +12,29 @@ from lifecycle import ResourceManager from memory_readings import read_rss_everywhere from proxy_client import ProxyClient from setup.datastore import read_db_counters, read_redis_counters -from setup.observe import Observation, evaluate, read_mock_requests, render, reset_mock, set_mock_latency +from setup.observe import ( + Observation, + evaluate, + junit_properties, + read_mock_requests, + render, + reset_mock, + set_mock_latency, +) from setup.plan import build_plan from setup.podcpu import gateway_cpu_sampler from setup.provision import MOCK_PROVIDERS, build_setup, mock_url from setup.run import run_plan -from setup.upgrade import assert_finished, monitor_for, upgrade_case +from setup.upgrade import assert_finished, monitor_for, upgrade_case, upgrade_properties from setup.upgrade_worker import connect, read_state pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] @pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only") -def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: ResourceManager) -> None: +def test_database_upgrade( + proxy: ProxyClient, idp: Keycloak, resources: ResourceManager, request: pytest.FixtureRequest +) -> None: case: Final = upgrade_case() urls: Final = tuple(mock_url(provider) for provider in MOCK_PROVIDERS) for url in urls: @@ -52,12 +62,7 @@ def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: Resource f"successful_requests={monitor.successes} persisted_requests={monitor.persisted_count} " f"unmatched_requests={len(monitor.pending)}\n" ) - monitor.assert_complete(state) - assert_finished(connection, state) rss_after: Final = read_rss_everywhere(proxy) - assert {(r.replica, r.hostname) for r in rss_before.readings} == { - (r.replica, r.hostname) for r in rss_after.readings - }, "Gateway pods replaced during the upgrade" observed: Final = Observation( records=outcome.records, mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)), @@ -73,6 +78,13 @@ def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: Resource gateway_cpu_seconds=cpu.seconds, ) rows, exceeded = evaluate(case, observed) + item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped + assert isinstance(item, pytest.Item) + item.user_properties[:] = [ + entry for entry in item.user_properties if not entry[0].startswith("perf_") + ] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set + item.user_properties.extend(junit_properties(rows)) + item.user_properties.extend(upgrade_properties(state, monitor)) _ = sys.stdout.write(render(case, observed, rows) + "\n") _ = sys.stdout.write( f"schema_seconds={state.schema_seconds:.3f} index_seconds={state.index_seconds:.3f} " @@ -82,4 +94,9 @@ def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: Resource artifact: Final = os.environ.get("E2E_ARTIFACT_DIR") if artifact: _ = (Path(artifact) / "database-upgrade.json").write_text(state.model_dump_json(), encoding="utf-8") + monitor.assert_complete(state) + assert_finished(connection, state) + assert {(r.replica, r.hostname) for r in rss_before.readings} == { + (r.replica, r.hostname) for r in rss_after.readings + }, "Gateway pods replaced during the upgrade" assert not exceeded, "\n".join(exceeded) diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index 49b23ddfc57..3e5c3c522c8 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -29,6 +29,7 @@ from setup.event_loop import ( render_stalls, summarize, ) +from setup.event_loop import junit_properties as stall_properties from setup.observe import ( Observation, evaluate, @@ -158,7 +159,9 @@ def test_perf_case( item.user_properties[:] = [ entry for entry in item.user_properties if not entry[0].startswith("perf_") ] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set - item.user_properties.extend(junit_properties(rows)) + item.user_properties.extend( + stall_properties(case.event_loop, stall_rows) if case.event_loop is not None else junit_properties(rows) + ) print(render(case, observation, rows)) if case.event_loop is not None: print(render_stalls(case.event_loop, stall_rows)) diff --git a/tests/e2e/perf/test_ui_journeys.py b/tests/e2e/perf/test_ui_journeys.py index db25608a00d..63eefedec68 100644 --- a/tests/e2e/perf/test_ui_journeys.py +++ b/tests/e2e/perf/test_ui_journeys.py @@ -11,14 +11,14 @@ from types import FrameType from typing import Final import pytest -from playwright.sync_api import Page from e2e_config import MASTER_KEY, UI_BASE_URL, UI_PASSWORD, UI_USERNAME from e2e_http import send_text_external, without_retries +from playwright.sync_api import Page from setup.ui_journeys import ( - assert_within_ceiling, Dashboard, Journey, Sample, + assert_within_ceiling, dashboard_page, database_failures, error_code, @@ -111,6 +111,7 @@ def test_ui_journey_latency(journey_name: str, ui_fixture: UiFixture, request: p item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped assert isinstance(item, pytest.Item) item.user_properties.extend(ui_junit_properties(samples)) + item.user_properties.extend(ui_fixture.volume) after: Final = settled_counter(lambda: _metrics(transport), evidence) item.user_properties.extend( ( diff --git a/tests/e2e/ui/tests/budgets/budgets.spec.ts b/tests/e2e/ui/tests/budgets/budgets.spec.ts index 89691c05605..866c7688426 100644 --- a/tests/e2e/ui/tests/budgets/budgets.spec.ts +++ b/tests/e2e/ui/tests/budgets/budgets.spec.ts @@ -60,13 +60,11 @@ test.describe("Budgets", () => { await expect(modal).toBeVisible({ timeout: 10_000 }); await modal.getByRole("textbox", { name: "Budget ID" }).fill(budgetId); - await modal.getByRole("spinbutton", { name: "Max Tokens per minute" }).fill("5000"); - await modal.getByRole("spinbutton", { name: "Max Requests per minute" }).fill("60"); - - await modal.getByRole("button", { name: "Optional Settings" }).click(); - await modal.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("25.5"); - await modal.getByRole("combobox", { name: "Reset Budget" }).click(); - await page.getByRole("option", { name: "weekly" }).click(); + await modal.getByRole("spinbutton", { name: "Tokens per minute" }).fill("5000"); + await modal.getByRole("spinbutton", { name: "Requests per minute" }).fill("60"); + await modal.getByRole("spinbutton", { name: "Spend limit" }).fill("25.5"); + await modal.getByRole("combobox", { name: "Reset period" }).click(); + await page.getByRole("option", { name: "per week" }).click(); await modal.getByRole("button", { name: "Create Budget" }).click(); await expect(modal).not.toBeVisible({ timeout: 10_000 }); @@ -100,15 +98,12 @@ test.describe("Budgets", () => { const modal = page.getByRole("dialog", { name: "Edit Budget" }); await expect(modal).toBeVisible({ timeout: 10_000 }); - await modal.getByRole("button", { name: "Optional Settings" }).click(); - await modal.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("99"); + await modal.getByRole("spinbutton", { name: "Spend limit" }).fill("99"); await modal.getByRole("button", { name: "Save", exact: true }).click(); await expect(modal).not.toBeVisible({ timeout: 10_000 }); await expect(page.getByRole("row").filter({ hasText: budgetId })).toContainText("$99.00", { timeout: 10_000 }); - // Not hypothetical: the edit form posts the whole budget, so a field it fails to - // seed from the existing row goes to the server as null and silently clears. const stored = await findBudget(page, budgetId); expect(stored?.max_budget, "spend cap raised").toBe(99); expect(stored?.tpm_limit, "TPM limit untouched by a spend-cap edit").toBe(1000); diff --git a/tests/e2e_harness/perf/setup/test_event_loop.py b/tests/e2e_harness/perf/setup/test_event_loop.py index 57868305716..75989b38275 100644 --- a/tests/e2e_harness/perf/setup/test_event_loop.py +++ b/tests/e2e_harness/perf/setup/test_event_loop.py @@ -1,8 +1,10 @@ from __future__ import annotations +import math from typing import Final -from setup.event_loop import EventLoopStalls, KnownStall, offenders, parse_stalls, summarize +from pydantic import BaseModel, ConfigDict, TypeAdapter +from setup.event_loop import EventLoopStalls, KnownStall, junit_properties, offenders, parse_stalls, summarize _HANDLER: Final = ( "Executing list[_Slice]: + return TypeAdapter(list[_Slice]).validate_json(str(properties[f"perf_event_loop_{name}_breakdown"])) + + def test_parse_keeps_the_coroutine_its_location_and_the_hold() -> None: (stall,) = parse_stalls(("INFO: plain log line", _HANDLER)) assert stall.coroutine == "RequestResponseCycle.run_asgi()" @@ -95,3 +111,72 @@ def test_a_collector_pass_explains_only_the_next_hold_and_only_one_at_least_as_l ("RequestResponseCycle.run_asgi()", False), ] assert summarize(_SPEC, parse_stalls((_GC_PASS,))) == () + + +def test_junit_properties_count_every_hold_and_only_the_holds_past_their_allowance_as_offending() -> None: + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, _STREAM, _HANDLER, _HANDLER, _HANDLER))) + properties: Final = dict(junit_properties(_SPEC, rows)) + assert properties["perf_event_loop_holds"] == 6 + assert properties["perf_event_loop_holds_limit"] == 6 + assert properties["perf_event_loop_offending_holds"] == 3 + assert properties["perf_event_loop_offending_holds_limit"] == 0 + assert properties["perf_event_loop_hold_seconds_max"] == 0.731 + assert math.isclose(float(properties["perf_event_loop_hold_seconds_total"]), 0.302 + 2 * 0.412 + 3 * 0.731) + breakdown: Final = _breakdown(properties, "holds") + assert [(slice.label, slice.value, slice.note, slice.over) for slice in breakdown] == [ + ( + "RequestResponseCycle.run_asgi() uvicorn/protocols/http/h11_impl.py:416", + 3, + "3 of 2 allowed, max 0.73 s", + True, + ), + ("TaskHandle._run_coro() anyio/_core/_tasks.py:275", 2, "not on the allowlist, max 0.41 s", True), + ("LifespanOn.main() uvicorn/lifespan/on.py:86", 1, "1 of 3 allowed, max 0.30 s", False), + ] + assert math.isclose(sum(slice.share or 0.0 for slice in breakdown), 1.0) + + +def test_the_offending_breakdown_lists_only_the_locations_past_their_cap_by_how_far_over_they_are() -> None: + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, _STREAM, _HANDLER, _HANDLER, _HANDLER))) + offending: Final = _breakdown(dict(junit_properties(_SPEC, rows)), "offending_holds") + assert [(slice.label, slice.value, slice.share, slice.over) for slice in offending] == [ + ("RequestResponseCycle.run_asgi() uvicorn/protocols/http/h11_impl.py:416", 1, 1 / 3, True), + ("TaskHandle._run_coro() anyio/_core/_tasks.py:275", 2, 2 / 3, True), + ] + clean: Final = dict(junit_properties(_SPEC, summarize(_SPEC, parse_stalls((_HANDLER, _HANDLER))))) + assert clean["perf_event_loop_offending_holds"] == 0 + assert "perf_event_loop_offending_holds_breakdown" not in clean + + +def test_the_seconds_breakdowns_carry_the_time_per_location_and_the_longest_single_holds() -> None: + long_handler: Final = _HANDLER.replace("0.731", "0.950") + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, long_handler, _HANDLER))) + properties: Final = dict(junit_properties(_SPEC, rows)) + per_location: Final = _breakdown(properties, "hold_seconds_total") + assert [(slice.label.split()[0], slice.value) for slice in per_location] == [ + ("RequestResponseCycle.run_asgi()", 0.950 + 0.731), + ("TaskHandle._run_coro()", 0.412), + ("LifespanOn.main()", 0.302), + ] + assert math.isclose(sum(slice.share or 0.0 for slice in per_location), 1.0) + longest: Final = _breakdown(properties, "hold_seconds_max") + assert [(slice.label.split()[0], slice.value, slice.share, slice.over) for slice in longest] == [ + ("RequestResponseCycle.run_asgi()", 0.950, None, False), + ("RequestResponseCycle.run_asgi()", 0.731, None, False), + ("TaskHandle._run_coro()", 0.412, None, True), + ("LifespanOn.main()", 0.302, None, False), + ] + + +def test_the_longest_holds_breakdown_stops_at_ten() -> None: + rows: Final = summarize(_SPEC, parse_stalls(tuple(_STARTUP for _ in range(12)))) + longest: Final = _breakdown(dict(junit_properties(_SPEC, rows)), "hold_seconds_max") + assert len(longest) == 10 + + +def test_junit_properties_without_holds_report_zeros_and_no_breakdown() -> None: + properties: Final = dict(junit_properties(_SPEC, ())) + assert properties["perf_event_loop_holds"] == 0 + assert properties["perf_event_loop_offending_holds"] == 0 + assert properties["perf_event_loop_hold_seconds_max"] == 0 + assert not [name for name in properties if name.endswith("_breakdown")] diff --git a/tests/e2e_harness/perf/setup/test_ui_journeys.py b/tests/e2e_harness/perf/setup/test_ui_journeys.py index dd2c1fc0301..d008368052c 100644 --- a/tests/e2e_harness/perf/setup/test_ui_journeys.py +++ b/tests/e2e_harness/perf/setup/test_ui_journeys.py @@ -6,9 +6,9 @@ import pytest from setup.ui_journeys import ( UI_JOURNEY_LIMITS, Sample, - error_code, assert_within_ceiling, database_failures, + error_code, journeys, summarize, ui_junit_properties, @@ -153,3 +153,21 @@ def test_failed_data_request_status_becomes_the_error_code() -> None: metrics: Final = dict(ui_junit_properties((gateway_error, slow, hung))) assert metrics["perf_ui_error_http_500"] == 1 and metrics["perf_ui_error_wall_clock"] == 1 assert "perf_ui_error_ui" not in metrics + + +def test_volume_properties_name_rows_and_bytes_per_queried_table() -> None: + from setup.ui_query_fixture import TableVolume, volume_properties + + properties: Final = volume_properties( + ( + TableVolume(table="LiteLLM_SpendLogs", rows=850_000_000, bytes=1_000_000_000_000), + TableVolume(table="LiteLLM_VerificationToken", rows=5_000, bytes=4_194_304), + ) + ) + + assert properties == ( + ("perf_ui_db_rows_LiteLLM_SpendLogs", 850_000_000), + ("perf_ui_db_bytes_LiteLLM_SpendLogs", 1_000_000_000_000), + ("perf_ui_db_rows_LiteLLM_VerificationToken", 5_000), + ("perf_ui_db_bytes_LiteLLM_VerificationToken", 4_194_304), + ) diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py index 1f2bc8c46cb..9a6c9d2af14 100644 --- a/tests/e2e_harness/perf/test_upgrade.py +++ b/tests/e2e_harness/perf/test_upgrade.py @@ -7,7 +7,14 @@ import pytest from cases.small_bodies_5rps import CASE from memory_readings import RssCapture, RssReading from setup.observe import Observation, evaluate -from setup.upgrade import MAX_TRAFFIC_SECONDS, UpgradeMonitor, matched_response_ids, postgres_failures, upgrade_case +from setup.upgrade import ( + MAX_TRAFFIC_SECONDS, + UpgradeMonitor, + matched_response_ids, + postgres_failures, + upgrade_case, + upgrade_properties, +) from setup.upgrade_worker import UpgradeState from setup.wire import PlanProgress, WireRecord, response_id @@ -332,3 +339,35 @@ def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples( rows, exceeded = evaluate(upgrade_case(), observed) assert tuple(row.label for row in rows if row.over or row.missing) == expected assert len(exceeded) == len(expected) + + +def test_upgrade_properties_carry_the_upgrade_timings_and_what_the_monitor_saw() -> None: + state: Final = UpgradeState( + phase="complete", + started_at=0, + schema_seconds=12.5, + index_seconds=240.25, + required_indexes=("idx_a", "idx_b"), + required_migrations=("20260101_a",), + pending_indexes=("idx_b",), + baseline_rows={}, + ) + monitor: Final = UpgradeMonitor( + state=lambda: state, + check_metrics=lambda: None, + read_activity=lambda: ([], 0), + read_persisted=lambda records: frozenset(), + start_indexes=lambda: None, + ) + monitor.overlap_writes = 37 + monitor.max_lock_waiters = 3 + monitor.pending = (record(0), record(1)) + + assert dict(upgrade_properties(state, monitor)) == { + "perf_upgrade_schema_seconds": 12.5, + "perf_upgrade_index_seconds": 240.25, + "perf_upgrade_overlap_writes": 37, + "perf_upgrade_max_lock_waiters": 3, + "perf_upgrade_pending_indexes": 1, + "perf_upgrade_unmatched_requests": 2, + } diff --git a/tests/integration/_support/hook_count_probe.py b/tests/integration/_support/hook_count_probe.py new file mode 100644 index 00000000000..0d4056e8fb9 --- /dev/null +++ b/tests/integration/_support/hook_count_probe.py @@ -0,0 +1,91 @@ +import json +import os +from collections.abc import Mapping +from datetime import datetime +from typing import Final + +import httpx +from fastapi import HTTPException + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger + + +def _record(hook: str, kwargs: Mapping[str, object]) -> bytes: + return json.dumps( + {"hook": hook, "litellm_call_id": kwargs.get("litellm_call_id"), "call_type": kwargs.get("call_type")} + ).encode() + + +def _send(hook: str, kwargs: Mapping[str, object]) -> None: + with httpx.Client(trust_env=False) as client: + client.post(os.environ["HOOK_COUNT_PROBE_URL"], content=_record(hook, kwargs)).raise_for_status() + + +async def _send_async(hook: str, kwargs: Mapping[str, object]) -> None: + async with httpx.AsyncClient(trust_env=False) as client: + response: Final = await client.post(os.environ["HOOK_COUNT_PROBE_URL"], content=_record(hook, kwargs)) + response.raise_for_status() + + +class HookCountProbe(CustomLogger): + """Reports every sync and async success and failure hook it receives, one record per call.""" + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + _send("sync_success", kwargs) + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + await _send_async("async_success", kwargs) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + _send("sync_failure", kwargs) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + await _send_async("async_failure", kwargs) + + +class SyncOnlyHookCountProbe(CustomLogger): + """Overrides only the sync success and failure hooks.""" + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + _send("sync_success", kwargs) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + _send("sync_failure", kwargs) + + +class BlockingPostCallGuardrail(CustomGuardrail): + """Rejects every response with a 400 after the provider answered.""" + + async def async_post_call_success_hook( + self, data: Mapping[str, object], user_api_key_dict: object, response: object + ) -> None: + raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail") + + +probe: Final = HookCountProbe() +sync_only_probe: Final = SyncOnlyHookCountProbe() + + +def sync_success_function( + kwargs: Mapping[str, object], completion_response: object, start_time: datetime, end_time: datetime +) -> None: + _send("function_success", kwargs) + + +def sync_failure_function( + kwargs: Mapping[str, object], completion_response: object, start_time: datetime, end_time: datetime +) -> None: + _send("function_failure", kwargs) diff --git a/tests/integration/_support/hook_counts.py b/tests/integration/_support/hook_counts.py new file mode 100644 index 00000000000..6b7f482bd85 --- /dev/null +++ b/tests/integration/_support/hook_counts.py @@ -0,0 +1,329 @@ +"""Scripted provider, proxy config and request helpers shared by the CustomLogger hook count tests.""" + +import base64 +import json +import threading +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Final + +import httpx +import yaml +from integration._support.client import Gateway, held +from integration._support.wire import Reply, Request, Wire + +WAV_BYTES: Final = ( + b"RIFF\x30\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00" + b"\x40\x1f\x00\x00\x40\x1f\x00\x00\x01\x00\x08\x00data\x08\x00\x00\x00" + b"\x00\x01\x02\x03\x04\x05\x06\x07" +) +PNG_BYTES: Final = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00" + b"\x1f\x15\xc4\x89\x00\x00\x00\rIDAT\x08\xd7c\xf8\xcf\xc0\xf0\x1f\x00\x05\x00\x01\xff" + b"\x89\x99=\x1d\x00\x00\x00\x00IEND\xaeB`\x82" +) +_IMAGE_REPLY: Final = json.dumps({"created": 1700000000, "data": [{"b64_json": base64.b64encode(PNG_BYTES).decode()}]}) +_CHAT_REPLY: Final = json.dumps( + { + "id": "chatcmpl-hook-count", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } +) +_EMBEDDING_REPLY: Final = json.dumps( + { + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}], + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + } +) +_MESSAGE_REPLY: Final = json.dumps( + { + "id": "msg_hook_count", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + } +) +_RESPONSE_REPLY: Final = json.dumps( + { + "id": "resp_hook_count", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_hook_count", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + } +) +_OCR_REPLY: Final = json.dumps( + { + "pages": [ + {"index": 0, "markdown": "page", "images": [], "dimensions": {"dpi": 200, "height": 1100, "width": 850}} + ], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 64}, + "document_annotation": None, + } +) +_SEARCH_REPLY: Final = json.dumps( + {"results": [{"title": "hook count", "url": "https://result.invalid/a", "snippet": "hit"}]} +) +_CHAT_FRAMES: Final = ( + 'data: {"id":"chatcmpl-hook-count","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-hook-count","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],' + '"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}', + "data: [DONE]", +) +_MESSAGE_FRAMES: Final = ( + "event: message_start\n" + 'data: {"type":"message_start","message":{"id":"msg_hook_count","type":"message","role":"assistant",' + '"model":"claude-haiku-4-5","content":[],"stop_reason":null,"stop_sequence":null,' + '"usage":{"input_tokens":3,"output_tokens":0}}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + "event: content_block_delta\n" + 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + "event: message_delta\n" + 'data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":1}}', + 'event: message_stop\ndata: {"type":"message_stop"}', +) +_RESPONSE_FRAMES: Final = ( + "event: response.created\n" + 'data: {"type":"response.created","sequence_number":0,"response":{"id":"resp_hook_count","object":"response",' + '"created_at":1,"status":"in_progress","model":"gpt-4o-mini","output":[],"usage":null}}', + "event: response.output_text.delta\n" + 'data: {"type":"response.output_text.delta","sequence_number":1,"item_id":"msg_hook_count","output_index":0,' + '"content_index":0,"delta":"ok"}', + "event: response.completed\n" + 'data: {"type":"response.completed","sequence_number":2,"response":{"id":"resp_hook_count","object":"response",' + '"created_at":1,"status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_hook_count",' + '"status":"completed","role":"assistant","content":[{"type":"output_text","text":"ok","annotations":[]}]}],' + '"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}', +) +_STREAM_FRAMES: Final = { + "/v1/chat/completions": _CHAT_FRAMES, + "/v1/messages": _MESSAGE_FRAMES, + "/v1/responses": _RESPONSE_FRAMES, +} +SURFACE_CALL_TYPES: Final = { + "chat": "acompletion", + "embedding": "aembedding", + "messages": "anthropic_messages", + "responses": "aresponses", + "speech": "aspeech", + "transcription": "atranscription", + "image_generation": "aimage_generation", + "image_edit": "aimage_edit", + "ocr": "aocr", + "search": "asearch", +} +IN_THE_ASYNC_FLAG_LIST: Final = frozenset({"chat", "embedding", "messages", "responses"}) +STREAMING_SURFACES: Final = ("chat", "messages", "responses") + + +def _streams(request: Request) -> bool: + if request.headers.get("content-type", "").startswith("application/json"): + return json.loads(request.body or b"{}").get("stream") is True + return False + + +def provider(request: Request) -> Reply: + """Answers every surface the hook count tests drive; any target under /fail/ is a provider 400.""" + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + if request.target.startswith("/fail/"): + return Reply( + status=400, + body=b'{"error":{"message":"scripted provider rejection","type":"invalid_request_error","code":null}}', + ) + if request.target in _STREAM_FRAMES and _streams(request): + return Reply( + chunks=tuple(f"{frame}\n\n".encode() for frame in _STREAM_FRAMES[request.target]), + content_type="text/event-stream", + ) + replies: Final[dict[str, Callable[[], Reply]]] = { + "/v1/chat/completions": lambda: Reply(body=_CHAT_REPLY.encode()), + "/v1/embeddings": lambda: Reply(body=_EMBEDDING_REPLY.encode()), + "/v1/messages": lambda: Reply(body=_MESSAGE_REPLY.encode()), + "/v1/responses": lambda: Reply(body=_RESPONSE_REPLY.encode()), + "/v1/audio/speech": lambda: Reply(body=WAV_BYTES, content_type="audio/wav"), + "/v1/audio/transcriptions": lambda: Reply(body=b'{"text":"hook count transcript"}'), + "/v1/images/generations": lambda: Reply(body=_IMAGE_REPLY.encode()), + "/v1/images/edits": lambda: Reply(body=_IMAGE_REPLY.encode()), + "/v1/ocr": lambda: Reply(body=_OCR_REPLY.encode()), + "/perplexity/search": lambda: Reply(body=_SEARCH_REPLY.encode()), + } + return replies[request.target]() + + +@dataclass +class HookRecords: + """Every hook record the probe reported, kept across drains so cells can read their own call id.""" + + wire: Wire + function_callbacks: bool + seen: list[dict[str, str]] = field(default_factory=list) + lock: threading.Lock = field(default_factory=threading.Lock) + + def drained(self) -> tuple[dict[str, str], ...]: + with self.lock: + self.seen.extend(json.loads(request.body) for request in self.wire.drain()) + return tuple(self.seen) + + def for_call(self, call_id: str) -> tuple[dict[str, str], ...]: + return tuple(record for record in self.drained() if record["litellm_call_id"] == call_id) + + +def _model(name: str, model: str, api_base: str, **extra: object) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": "synthetic-provider-key"}, + **extra, + } + + +def write_hook_count_config( + path: Path, + provider_url: str, + litellm_settings: Mapping[str, object], + guardrails: tuple[Mapping[str, object], ...] = (), +) -> Path: + """One success and one failure deployment per surface, all answered by `provider`.""" + models: Final = [] + tools: Final = [] + for outcome, prefix in (("success", ""), ("failure", "/fail")): + base: str = provider_url + prefix + models.extend( + ( + _model(f"hooks-chat-{outcome}", "openai/gpt-4o-mini", base + "/v1"), + _model(f"hooks-embedding-{outcome}", "openai/text-embedding-3-small", base + "/v1"), + _model(f"hooks-messages-{outcome}", "anthropic/claude-haiku-4-5", base), + _model(f"hooks-responses-{outcome}", "openai/gpt-4o-mini", base + "/v1"), + _model(f"hooks-speech-{outcome}", "openai/gpt-4o-mini-tts", base + "/v1"), + _model(f"hooks-transcription-{outcome}", "openai/whisper-1", base + "/v1"), + _model(f"hooks-image_generation-{outcome}", "openai/gpt-image-1", base + "/v1"), + _model(f"hooks-image_edit-{outcome}", "openai/gpt-image-1", base + "/v1"), + _model(f"hooks-ocr-{outcome}", "mistral/mistral-ocr-latest", base + "/v1", model_info={"mode": "ocr"}), + ) + ) + tools.append( + { + "search_tool_name": f"hooks-search-{outcome}", + "litellm_params": { + "search_provider": "perplexity", + "api_key": "synthetic-provider-key", + "api_base": base + "/perplexity", + }, + } + ) + path.write_text( + yaml.safe_dump( + { + "model_list": models, + "search_tools": tools, + "litellm_settings": dict(litellm_settings), + **({"guardrails": [dict(guardrail) for guardrail in guardrails]} if guardrails else {}), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_model_info_refresh": True, + }, + "router_settings": {"num_retries": 0}, + } + ) + ) + return path + + +def send( + proxy: Gateway, surface: str, outcome: str, call_id: str, *, key: str | None = None, stream: bool = False +) -> httpx.Response: + model: Final = f"hooks-{surface}-{outcome}" + nonce: Final = f"hook-count-{call_id}" + headers: Final = {"Authorization": f"Bearer {key or proxy.key}", "x-litellm-call-id": call_id} + if surface == "transcription": + return proxy.client.post( + "/v1/audio/transcriptions", + data={"model": model, "prompt": nonce}, + files={"file": ("a.wav", WAV_BYTES + call_id.encode(), "audio/wav")}, + headers=headers, + ) + if surface == "image_edit": + return proxy.client.post( + "/v1/images/edits", + data={"model": model, "prompt": nonce}, + files={"image": ("a.png", PNG_BYTES, "image/png")}, + headers=headers, + ) + streaming: Final[dict[str, object]] = {"stream": True} if stream else {} + routes: Final[dict[str, tuple[str, dict[str, object]]]] = { + "chat": ( + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": nonce}], **streaming}, + ), + "embedding": ("/v1/embeddings", {"model": model, "input": nonce}), + "messages": ( + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": nonce}], **streaming}, + ), + "responses": ("/v1/responses", {"model": model, "input": nonce, **streaming}), + "speech": ("/v1/audio/speech", {"model": model, "input": nonce, "voice": "alloy", "response_format": "wav"}), + "image_generation": ("/v1/images/generations", {"model": model, "prompt": nonce, "n": 1}), + "ocr": ( + "/v1/ocr", + { + "model": model, + "document": { + "type": "document_url", + "document_url": "data:application/pdf;base64," + base64.b64encode(nonce.encode()).decode(), + }, + }, + ), + "search": ("/v1/search", {"search_tool_name": f"hooks-search-{outcome}", "query": nonce}), + } + path, body = routes[surface] + return proxy.client.post(path, json=body, headers=headers) + + +def counts_function_hooks(records: HookRecords, surface: str) -> bool: + return records.function_callbacks and surface not in IN_THE_ASYNC_FLAG_LIST + + +def probe_hooks(records: HookRecords, surface: str, found: tuple[dict[str, str], ...]) -> list[str]: + counted: Final = counts_function_hooks(records, surface) + return sorted(record["hook"] for record in found if counted or not record["hook"].startswith("function_")) + + +def one_async_hook(records: HookRecords, surface: str, outcome: str) -> list[str]: + function_hook: Final = (f"function_{outcome}",) if counts_function_hooks(records, surface) else () + return sorted((f"async_{outcome}", *function_hook)) + + +def assert_hooks(records: HookRecords, surface: str, call_id: str, expected: list[str]) -> None: + settled: Final = held( + lambda: records.for_call(call_id), + lambda found: probe_hooks(records, surface, found) == expected, + holding=3, + seconds=20, + ) + assert {record["call_type"] for record in settled} == {SURFACE_CALL_TYPES[surface]}, settled diff --git a/tests/integration/_support/mcp.py b/tests/integration/_support/mcp.py index f69079ce90b..99bf71cd459 100644 --- a/tests/integration/_support/mcp.py +++ b/tests/integration/_support/mcp.py @@ -282,6 +282,7 @@ class ScriptedTool: respond: Callable[[JsonRpc], Reply | JsonRpc] description: str | Callable[[Mapping[str, str]], str] | None = None input_schema: JsonRpc = field(default_factory=lambda: {"type": "object"}) + event_stream: bool = False def listing(self, headers: Mapping[str, str]) -> JsonRpc: described: Final = self.description(headers) if callable(self.description) else self.description @@ -334,7 +335,10 @@ def scripted_peer(*tools: ScriptedTool) -> Iterator[McpPeer]: if tool is None: return jsonrpc_error(identity, -32602, "unknown tool") produced: Final = tool.respond(body["params"]) - return produced if isinstance(produced, Reply) else jsonrpc_reply(identity, produced) + reply: Final = produced if isinstance(produced, Reply) else jsonrpc_reply(identity, produced) + if not tool.event_stream: + return reply + return Reply(content_type="text/event-stream", body=b"event: message\ndata: " + reply.body + b"\n\n") with wire_server(provider) as wire: yield McpPeer(wire.url + "/mcp", observed) diff --git a/tests/integration/_support/scim.py b/tests/integration/_support/scim.py new file mode 100644 index 00000000000..11eea030ba3 --- /dev/null +++ b/tests/integration/_support/scim.py @@ -0,0 +1,207 @@ +"""SCIM v2 group helpers for the management cells: seeded users, group pushes, the team lock, and what a group +leaves behind in Postgres.""" + +import os +import uuid +from collections.abc import Generator, Mapping, Sequence +from contextlib import contextmanager +from typing import Final + +import httpx +import psycopg +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + delete_key_if_present, + object_value, + string_value, +) +from integration._support.database import read_rows, write_rows +from pydantic import JsonValue + +GROUP_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:Group" +PATCH_OP_SCHEMA: Final = "urn:ietf:params:scim:api:messages:2.0:PatchOp" +SEED_USERS_SQL: Final = """ +INSERT INTO "LiteLLM_UserTable" (user_id, user_role, teams, models) +SELECT %s || '-' || lpad(n::text, 3, '0'), 'internal_user', '{}'::text[], '{}'::text[] +FROM generate_series(1, %s::int) AS n +""" +TAKE_TEAM_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext(%s))" +WAITERS_ON_HELD_LOCK_SQL: Final = """ +SELECT count(*)::text AS waiting +FROM pg_locks waiter +JOIN pg_stat_activity session ON session.pid = waiter.pid +WHERE waiter.locktype = 'advisory' + AND NOT waiter.granted + AND session.wait_event_type = 'Lock' + AND (waiter.classid, waiter.objid, waiter.objsubid) IN ( + SELECT held.classid, held.objid, held.objsubid + FROM pg_locks held + WHERE held.locktype = 'advisory' AND held.granted AND held.pid = %s::int + ) +""" +TAKE_USER_ROW_LOCK_SQL: Final = 'SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s FOR UPDATE' +WAITERS_ON_ROWS_LOCKED_BY_SQL: Final = """ +SELECT count(*)::text AS waiting +FROM pg_locks waiter +WHERE waiter.locktype = 'transactionid' + AND NOT waiter.granted + AND waiter.transactionid IN ( + SELECT held.transactionid + FROM pg_locks held + WHERE held.locktype = 'transactionid' AND held.granted AND held.pid = %s::int + ) +""" + + +def seed_users(scenario: Scenario, count: int) -> tuple[str, ...]: + prefix: Final = f"integration-scim-{uuid.uuid4().hex}" + write_rows(SEED_USERS_SQL, (prefix, str(count))) + scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_UserTable" WHERE user_id LIKE %s', (f"{prefix}-%",)) + return tuple(f"{prefix}-{index:03d}" for index in range(1, count + 1)) + + +def members(users: Sequence[str]) -> list[JsonValue]: + return [{"value": user} for user in users] + + +def group_name() -> str: + return f"integration-{uuid.uuid4().hex}" + + +def create_group(candidate: Gateway, users: Sequence[str], *, display_name: str | None = None) -> httpx.Response: + return candidate.request( + "POST", + "/scim/v2/Groups", + {"schemas": [GROUP_SCHEMA], "displayName": display_name or group_name(), "members": members(users)}, + ) + + +def replace_group( + candidate: Gateway, team: str, users: Sequence[str], *, display_name: str | None = None +) -> httpx.Response: + return candidate.request( + "PUT", + f"/scim/v2/Groups/{team}", + {"schemas": [GROUP_SCHEMA], "id": team, "displayName": display_name or group_name(), "members": members(users)}, + ) + + +def add_to_group(candidate: Gateway, team: str, users: Sequence[str]) -> httpx.Response: + return candidate.request( + "PATCH", + f"/scim/v2/Groups/{team}", + {"schemas": [PATCH_OP_SCHEMA], "Operations": [{"op": "add", "path": "members", "value": members(users)}]}, + ) + + +def body(response: httpx.Response) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(response.content) + + +def member_ids(group: Mapping[str, JsonValue]) -> frozenset[str]: + listed: Final = group.get("members") or [] + assert isinstance(listed, list), listed + return frozenset(string_value(object_value(member)["value"]) for member in listed) + + +def delete_group(candidate: Gateway, team: str, *, key: str | None = None) -> httpx.Response: + return candidate.request("DELETE", f"/scim/v2/Groups/{team}", key=key) + + +def delete_team_if_present(scenario: Scenario, team: str) -> None: + if read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,)): + scenario.delete_team(team) + + +def created_team(scenario: Scenario, created: httpx.Response) -> str: + assert created.status_code == 201, created.text + team: Final = string_value(body(created)["id"]) + scenario.cleanups.callback(delete_team_if_present, scenario, team) + return team + + +def team_key(candidate: Gateway, team: str, user: str) -> str: + created: Final = candidate.post( + "/key/generate", {"team_id": team, "user_id": user, "key_alias": f"integration-{uuid.uuid4().hex}"} + ) + return string_value(created["key"]) + + +def service_account_key(candidate: Gateway, scenario: Scenario, team: str) -> str: + created: Final = candidate.post("/key/generate", {"team_id": team, "key_alias": f"integration-{uuid.uuid4().hex}"}) + key: Final = string_value(created["key"]) + scenario.cleanups.callback(delete_key_if_present, candidate, key) + return key + + +def membership_user_ids(team: str) -> frozenset[str]: + rows: Final = read_rows('SELECT user_id FROM "LiteLLM_TeamMembership" WHERE team_id = %s', (team,)) + return frozenset(string_value(row["user_id"]) for row in rows) + + +def keys_of(team: str) -> frozenset[str]: + rows: Final = read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,)) + return frozenset(string_value(row["token"]) for row in rows) + + +def users_referencing(team: str) -> frozenset[str]: + rows: Final = read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE %s = ANY(teams)', (team,)) + return frozenset(string_value(row["user_id"]) for row in rows) + + +def user_role(candidate: Gateway, user: str) -> str: + return string_value(object_value(candidate.get("/user/info", {"user_id": user})["user_info"])["user_role"]) + + +def groups_of(candidate: Gateway, user: str) -> frozenset[str]: + groups: Final = candidate.get(f"/scim/v2/Users/{user}").get("groups") or [] + assert isinstance(groups, list), groups + return frozenset(string_value(object_value(group)["value"]) for group in groups) + + +def assert_landed(candidate: Gateway, team: str, users: Sequence[str]) -> None: + assert member_ids(candidate.get(f"/scim/v2/Groups/{team}")) == frozenset(users) + assert membership_user_ids(team) == frozenset(users) + assert users_referencing(team) == frozenset(users) + + +def assert_gone(candidate: Gateway, team: str, users: Sequence[str], key: str) -> None: + assert candidate.request("GET", f"/scim/v2/Groups/{team}").status_code == 404 + assert team not in groups_of(candidate, users[0]) + assert membership_user_ids(team) == frozenset() + assert users_referencing(team) == frozenset() + assert keys_of(team) == frozenset() + assert candidate.request("GET", "/v1/models", key=key).status_code == 401 + + +@contextmanager +def held_team_lock(team: str) -> Generator[int]: + """Hold the team's advisory lock from a session of our own and yield its backend pid; leaving the block + commits, which releases the lock to whoever is queued on it.""" + with psycopg.connect(os.environ["DATABASE_URL"]) as holder: + holder.execute(TAKE_TEAM_LOCK_SQL, (team,)) + yield holder.info.backend_pid + + +@contextmanager +def held_user_row(user: str) -> Generator[int]: + """Hold a row lock on the user's row from a session of our own and yield its backend pid, so a role write + that reaches the row waits until leaving the block commits.""" + with psycopg.connect(os.environ["DATABASE_URL"]) as holder: + holder.execute(TAKE_USER_ROW_LOCK_SQL, (user,)) + yield holder.info.backend_pid + + +def waiters_on_lock_held_by(backend_pid: int) -> int: + return _waiting(WAITERS_ON_HELD_LOCK_SQL, backend_pid) + + +def waiters_on_rows_locked_by(backend_pid: int) -> int: + return _waiting(WAITERS_ON_ROWS_LOCKED_BY_SQL, backend_pid) + + +def _waiting(query: str, backend_pid: int) -> int: + rows: Final = read_rows(query, (str(backend_pid),)) + return int(string_value(rows[0]["waiting"])) diff --git a/tests/integration/compatibility/test_openai_consumer.py b/tests/integration/compatibility/test_openai_consumer.py index a49f20f4711..6150995723b 100644 --- a/tests/integration/compatibility/test_openai_consumer.py +++ b/tests/integration/compatibility/test_openai_consumer.py @@ -459,7 +459,6 @@ async def test_openai_clients_parse_text_completion_streams_and_choices(gateway: def test_text_completion_usage_chunk_has_empty_choices(gateway: Gateway) -> None: - pytest.skip("BUG: the /v1/completions include_usage chunk carries choices [{index: 0}] instead of []") identity: Final = f"cmpl-usage-{uuid.uuid4().hex}" prompt: Final = f"raw-stream {identity}" with wire_server(_text_provider(identity, {prompt})) as wire, gateway.scenario() as scenario: diff --git a/tests/integration/database/test_cleanup_transactions.py b/tests/integration/database/test_cleanup_transactions.py index 72da2bb48ee..202657d180e 100644 --- a/tests/integration/database/test_cleanup_transactions.py +++ b/tests/integration/database/test_cleanup_transactions.py @@ -113,7 +113,7 @@ async def test_size_rule_deletes_the_oldest_rows_until_under_the_target_and_does """ Deleting rows never shrinks the table file, so the size pass must stop on live bytes, not on pg_total_relation_size, taking the oldest rows first and keeping the tool index on the same cutoff. Once - autovacuum has turned the deleted rows into free space the catalog no longer shows them as dead, and the + VACUUM has turned the deleted rows into free space the catalog no longer shows them as dead, and the files still measure the same; the next check must read what the last run freed from its saved count and leave the table alone instead of calling the whole file live again. """ @@ -140,13 +140,15 @@ async def test_size_rule_deletes_the_oldest_rows_until_under_the_target_and_does setup.execute( sql.SQL( 'CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL, payload text, ' - 'PRIMARY KEY (request_id, "startTime"))' + 'PRIMARY KEY (request_id, "startTime")) ' + 'WITH (autovacuum_enabled = false, toast.autovacuum_enabled = false)' ).format(logs) ) setup.execute( sql.SQL( "CREATE TABLE {} (request_id text, tool_name text, start_time timestamp NOT NULL, " - "PRIMARY KEY (request_id, tool_name))" + "PRIMARY KEY (request_id, tool_name)) " + "WITH (autovacuum_enabled = false, toast.autovacuum_enabled = false)" ).format(index) ) setup.execute(sql.SQL("CREATE TABLE {} (param_name text PRIMARY KEY, param_value jsonb)").format(config)) diff --git a/tests/integration/management/test_model_hub_custom_metadata.py b/tests/integration/management/test_model_hub_custom_metadata.py new file mode 100644 index 00000000000..6f007b6ebac --- /dev/null +++ b/tests/integration/management/test_model_hub_custom_metadata.py @@ -0,0 +1,276 @@ +import math +import uuid +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import yaml +from pydantic import JsonValue + +from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS +from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.process import graceful_stop_seconds, owned_proxy_process + +LIFECYCLE_METADATA: Final = {"retirement_date": "2027-03-01", "ptu_max_tpm": 300000, "provisioned_ptu": 100} + + +def _model_group(gateway: Gateway, model_group: str) -> dict[str, JsonValue]: + data: Final = gateway.get("/model_group/info", params={"model_group": model_group})["data"] + assert isinstance(data, list) and len(data) == 1, data + return object_value(data[0]) + + +def _deployment(gateway: Gateway, scenario: Scenario, model_name: str, model_info: dict[str, JsonValue]) -> str: + created: Final = gateway.post( + "/model/new", + { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": f"{gateway.upstream_url}/v1", + }, + "model_info": model_info, + }, + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return identity + + +def _custom_fields(row: dict[str, JsonValue]) -> tuple[JsonValue, JsonValue]: + assert "custom_tags" in row and "custom_metadata" in row, row + return row["custom_tags"], row["custom_metadata"] + + +def _public_row(gateway: Gateway, path: str, model_group: str) -> dict[str, JsonValue] | None: + response: Final = gateway.client.get(path) + assert response.status_code == 200, response.text + body: Final = response.json() + rows: Final = body if isinstance(body, list) else object_value(body)["data"] + assert isinstance(rows, list), response.text + return next((object_value(row) for row in rows if object_value(row).get("model_group") == model_group), None) + + +def test_model_group_info_serves_the_custom_tags_and_metadata_set_on_the_model(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model( + model_info={"custom_tags": ["deprecated", "ptu"], "custom_metadata": LIFECYCLE_METADATA} + ) + group: Final = _model_group(gateway, model) + assert group.get("custom_tags") == ["deprecated", "ptu"], group + assert group.get("custom_metadata") == LIFECYCLE_METADATA, group + + +def test_model_group_info_merges_custom_tags_and_metadata_across_the_groups_deployments(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model_name: Final = f"integration-{uuid.uuid4().hex}" + _deployment( + gateway, + scenario, + model_name, + {"custom_tags": ["ptu", "eu"], "custom_metadata": {"provisioned_ptu": 100, "region": "westeurope"}}, + ) + _deployment( + gateway, + scenario, + model_name, + { + "custom_tags": ["eu", "deprecated"], + "custom_metadata": {"provisioned_ptu": 50, "retirement_date": "2027-03-01"}, + }, + ) + group: Final = _model_group(gateway, model_name) + assert sorted(string_value(tag) for tag in group.get("custom_tags") or []) == ["deprecated", "eu", "ptu"], group + metadata: Final = object_value(group.get("custom_metadata")) + assert metadata.get("region") == "westeurope", metadata + assert metadata.get("retirement_date") == "2027-03-01", metadata + assert metadata.get("provisioned_ptu") in (100, 50), metadata + + +def test_model_group_info_ignores_custom_metadata_values_that_are_not_scalars(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model( + model_info={ + "custom_tags": ["ok", 7, {"nested": True}], + "custom_metadata": {"owner": "platform", "nested": {"a": 1}, "list": [1, 2]}, + } + ) + group: Final = _model_group(gateway, model) + assert group.get("custom_tags") == ["ok"], group + assert group.get("custom_metadata") == {"owner": "platform"}, group + + +@pytest.mark.parametrize("path", ["/public/model_hub", "/public/v1/model_hub"]) +def test_public_model_hub_serves_the_custom_tags_and_metadata_of_a_public_model_group( + gateway: Gateway, path: str +) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model( + model_info={"custom_tags": ["deprecated", "ptu"], "custom_metadata": LIFECYCLE_METADATA} + ) + gateway.post("/model_group/make_public", {"model_groups": [model]}) + row: Final = eventually( + lambda: _public_row(gateway, path, model), + lambda found: found is not None, + seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ) + assert row is not None + assert _custom_fields(row) == (["deprecated", "ptu"], LIFECYCLE_METADATA), row + + +def test_patching_a_model_replaces_the_custom_tags_and_metadata_served_for_its_group(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model_name: Final = f"integration-{uuid.uuid4().hex}" + identity: Final = _deployment( + gateway, scenario, model_name, {"custom_tags": ["ptu"], "custom_metadata": LIFECYCLE_METADATA} + ) + patched: Final = gateway.request( + "PATCH", + f"/model/{identity}/update", + {"model_info": {"custom_tags": ["retiring"], "custom_metadata": {"retirement_date": "2027-06-01"}}}, + ) + assert patched.status_code == 200, patched.text + group: Final = eventually( + lambda: _model_group(gateway, model_name), + lambda found: found.get("custom_tags") == ["retiring"], + seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ) + assert _custom_fields(group) == (["retiring"], {"retirement_date": "2027-06-01"}), group + + +@pytest.mark.parametrize( + "model_info", + [ + pytest.param({"description": "plain model"}, id="missing"), + pytest.param({"custom_tags": [], "custom_metadata": {}}, id="empty"), + pytest.param({"custom_tags": None, "custom_metadata": None}, id="null"), + pytest.param({"custom_tags": "ptu", "custom_metadata": "retiring"}, id="wrong-shape"), + pytest.param( + {"custom_tags": [1, None, {"a": 1}], "custom_metadata": {"nested": {"a": 1}}}, id="no-valid-entry" + ), + ], +) +def test_model_group_info_serves_null_custom_fields_when_the_model_sets_none_usable( + gateway: Gateway, model_info: dict[str, JsonValue] +) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(model_info=model_info) + group: Final = _model_group(gateway, model) + assert _custom_fields(group) == (None, None), group + + +def test_model_group_info_keeps_long_values_whole_and_serves_a_repeated_tag_once(gateway: Gateway) -> None: + long_value: Final = "x" * 5000 + with gateway.scenario() as scenario: + model: Final = scenario.model( + model_info={"custom_tags": ["ptu", "ptu", "eu", "ptu"], "custom_metadata": {"notes": long_value}} + ) + group: Final = _model_group(gateway, model) + assert _custom_fields(group) == (["ptu", "eu"], {"notes": long_value}), group + + +def test_custom_tags_and_metadata_are_never_forwarded_to_the_provider(gateway: Gateway) -> None: + marker: Final = f"custom-fields-{uuid.uuid4().hex}" + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + model: Final = scenario.model( + model_info={"custom_tags": ["deprecated"], "custom_metadata": {"retirement_date": "2027-03-01"}} + ) + reply: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + assert reply.status_code == 200, reply.text + observed: Final = upstream.get("/__observations") + assert observed.status_code == 200, observed.text + bodies: Final = [ + str(object_value(value)["body"]) + for value in object_value(observed.json())["requests"] or [] + if marker in str(object_value(value)["body"]) + ] + assert len(bodies) == 1, bodies + assert "custom_tags" not in bodies[0] and "2027-03-01" not in bodies[0], bodies[0] + + +def _yaml_config( + directory: Path, + upstream_url: str, + model_name: str, + model_info: Mapping[str, object] = MappingProxyType( + {"custom_tags": ["deprecated", "ptu"], "custom_metadata": LIFECYCLE_METADATA} + ), + model_group_alias: Mapping[str, str] = MappingProxyType({}), +) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + deployment: Final = { + "model_name": model_name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_base": f"{upstream_url}/v1", "api_key": "yaml-key"}, + "model_info": dict(model_info), + } + path: Final = directory / "custom_fields.yaml" + router_settings: Final = {**config.get("router_settings", {}), "model_group_alias": dict(model_group_alias)} + path.write_text( + yaml.safe_dump( + {**config, "model_list": [*config["model_list"], deployment], "router_settings": router_settings} + ) + ) + return path + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_yaml_custom_fields_survive_a_killed_worker_on_every_worker(gateway: Gateway, tmp_path: Path) -> None: + model_name: Final = f"integration-yaml-{uuid.uuid4().hex}" + config: Final = _yaml_config(tmp_path, gateway.upstream_url, model_name) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + victim: Final = min(owned.worker_pids()) + owned.kill_worker(victim) + with ThreadPoolExecutor(max_workers=8) as pool: + groups: Final = tuple(pool.map(lambda _: _model_group(owned.gateway, model_name), range(24))) + assert all(_custom_fields(group) == (["deprecated", "ptu"], LIFECYCLE_METADATA) for group in groups), groups + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize( + "unservable_metadata", + [ + pytest.param({2027: "retirement"}, id="non-string-key"), + pytest.param({"nan": math.nan, "inf": math.inf, "neg_inf": -math.inf}, id="non-finite-number"), + ], +) +def test_yaml_custom_metadata_that_json_cannot_carry_is_dropped_without_breaking_completions_or_the_hub( + gateway: Gateway, tmp_path: Path, unservable_metadata: Mapping[object, object] +) -> None: + model_name: Final = f"integration-yaml-{uuid.uuid4().hex}" + config: Final = _yaml_config( + tmp_path, + gateway.upstream_url, + model_name, + {"custom_tags": ["ptu"], "custom_metadata": {**unservable_metadata, "retirement_date": "2027-03-01"}}, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + reply: Final = owned.gateway.request( + "POST", "/v1/chat/completions", {"model": model_name, "messages": [{"role": "user", "content": "hi"}]} + ) + assert reply.status_code == 200, reply.text + every_group: Final = owned.gateway.request("GET", "/model_group/info") + assert every_group.status_code == 200, every_group.text + group: Final = _model_group(owned.gateway, model_name) + assert _custom_fields(group) == (["ptu"], {"retirement_date": "2027-03-01"}), group + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_an_alias_group_serves_the_custom_tags_and_metadata_of_the_group_it_points_to( + gateway: Gateway, tmp_path: Path +) -> None: + model_name: Final = f"integration-yaml-{uuid.uuid4().hex}" + alias: Final = f"{model_name}-alias" + config: Final = _yaml_config(tmp_path, gateway.upstream_url, model_name, model_group_alias={alias: model_name}) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + group: Final = _model_group(owned.gateway, alias) + assert _custom_fields(group) == (["deprecated", "ptu"], LIFECYCLE_METADATA), group diff --git a/tests/integration/management/test_scim_group_bulk_roster.py b/tests/integration/management/test_scim_group_bulk_roster.py index 6d8e657c5ce..c61aa62f52e 100644 --- a/tests/integration/management/test_scim_group_bulk_roster.py +++ b/tests/integration/management/test_scim_group_bulk_roster.py @@ -3,7 +3,10 @@ POST /scim/v2/Groups, PUT and PATCH /scim/v2/Groups/{id} used to run /team/member_add once per member, so a 500-member push was thousands of statements and outlasted the edge in front of the proxy. One locked transaction now resolves every member, writes the roster and the membership rows and tags the users, so a -push of 500 members runs as many statements as a push of 5. +push of 500 members runs as many statements as a push of 5. DELETE /scim/v2/Groups/{id} used to detach the +members one at a time and leave their membership rows and keys behind; it now detaches them in one locked pass, +so deleting 500 members who each hold a key costs as many statements as deleting 5 and leaves no membership rows +or member keys. """ import os @@ -13,6 +16,7 @@ from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass +from functools import partial from hashlib import sha256 from pathlib import Path from typing import Final, LiteralString @@ -22,7 +26,6 @@ import psycopg import pytest from integration._support.budget import case_timeout_seconds from integration._support.client import ( - JSON_OBJECT, Gateway, Scenario, delete_key_if_present, @@ -32,14 +35,30 @@ from integration._support.client import ( ) from integration._support.database import read_rows, write_rows from integration._support.database_relay import StatementCountingRelay, statement_counting_relay -from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.scim import ( + GROUP_SCHEMA, + PATCH_OP_SCHEMA, + add_to_group, + assert_gone, + assert_landed, + body, + create_group, + created_team, + delete_group, + member_ids, + members, + membership_user_ids, + replace_group, + seed_users, + team_key, + users_referencing, +) from psycopg import sql from pydantic import JsonValue pytestmark: Final = pytest.mark.shares_worker_state -GROUP_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:Group" -PATCH_OP_SCHEMA: Final = "urn:ietf:params:scim:api:messages:2.0:PatchOp" SMALL_PUSH: Final = 5 LARGE_PUSH: Final = 500 # Background work on an idle proxy (the cron-job leader poll, spend flushes) lands a few statements inside a @@ -48,7 +67,9 @@ NOISE_ALLOWANCE: Final = 100 UNFIXED_PUSH_SECONDS: Final = 120 SWAP_GROUP: Final = 200 SWAP_ROUNDS: Final = 10 +OWNED_PROXY_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 2 * UNFIXED_PUSH_SECONDS) CHAOS_GROUP: Final = 100 +KEY_WRITERS: Final = 16 ROLLBACK_SECONDS: Final = 90 TEAM_ROW_SQL: Final = 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = %s' LOCK_USER_ROW_SQL: Final = 'SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s FOR UPDATE' @@ -56,16 +77,12 @@ LOCK_KEY_ROW_SQL: Final = 'SELECT token FROM "LiteLLM_VerificationToken" WHERE t SESSIONS_BLOCKED_BY_SQL: Final = ( "SELECT count(*)::int AS waiting FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))" ) +BACKENDS_BLOCKED_BY_SQL: Final = "SELECT pid FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))" TRANSACTIONS_OPEN_SINCE_SQL: Final = """ SELECT count(*)::int AS open FROM pg_stat_activity WHERE datname = current_database() AND backend_type = 'client backend' AND pid <> pg_backend_pid() AND xact_start < %s::timestamptz """ -SEED_USERS_SQL: Final = """ -INSERT INTO "LiteLLM_UserTable" (user_id, user_role, teams, models) -SELECT %s || '-' || lpad(n::text, 3, '0'), 'internal_user', '{}'::text[], '{}'::text[] -FROM generate_series(1, %s::int) AS n -""" @dataclass(frozen=True, slots=True) @@ -82,71 +99,6 @@ def _measured(relay: StatementCountingRelay, send: Callable[[], httpx.Response]) return Measured(response=response, statements=relay.statements - before, seconds=time.perf_counter() - started) -def _seed_users(scenario: Scenario, count: int) -> tuple[str, ...]: - prefix: Final = f"integration-scim-{uuid.uuid4().hex}" - write_rows(SEED_USERS_SQL, (prefix, str(count))) - scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_UserTable" WHERE user_id LIKE %s', (f"{prefix}-%",)) - return tuple(f"{prefix}-{index:03d}" for index in range(1, count + 1)) - - -def _members(users: Sequence[str]) -> list[JsonValue]: - return [{"value": user} for user in users] - - -def _create_group(candidate: Gateway, users: Sequence[str]) -> httpx.Response: - return candidate.request( - "POST", - "/scim/v2/Groups", - {"schemas": [GROUP_SCHEMA], "displayName": f"integration-{uuid.uuid4().hex}", "members": _members(users)}, - ) - - -def _replace_group(candidate: Gateway, team: str, users: Sequence[str]) -> httpx.Response: - return candidate.request( - "PUT", - f"/scim/v2/Groups/{team}", - { - "schemas": [GROUP_SCHEMA], - "id": team, - "displayName": f"integration-{uuid.uuid4().hex}", - "members": _members(users), - }, - ) - - -def _add_to_group(candidate: Gateway, team: str, users: Sequence[str]) -> httpx.Response: - return candidate.request( - "PATCH", - f"/scim/v2/Groups/{team}", - { - "schemas": [PATCH_OP_SCHEMA], - "Operations": [{"op": "add", "path": "members", "value": _members(users)}], - }, - ) - - -def _body(response: httpx.Response) -> dict[str, JsonValue]: - return JSON_OBJECT.validate_json(response.content) - - -def _member_ids(body: Mapping[str, JsonValue]) -> frozenset[str]: - members: Final = body.get("members") or [] - assert isinstance(members, list), members - return frozenset(string_value(object_value(member)["value"]) for member in members) - - -def _created_team(scenario: Scenario, created: httpx.Response) -> str: - assert created.status_code == 201, created.text - team: Final = string_value(_body(created)["id"]) - scenario.cleanups.callback(scenario.delete_team, team) - return team - - -def _membership_user_ids(team: str) -> frozenset[str]: - rows: Final = read_rows('SELECT user_id FROM "LiteLLM_TeamMembership" WHERE team_id = %s', (team,)) - return frozenset(string_value(row["user_id"]) for row in rows) - - _UTC_MICROSECONDS: Final = "YYYY-MM-DD HH24:MI:SS.US" @@ -163,18 +115,7 @@ def _database_utc_now() -> str: return string_value(rows[0]["utc_now"]) -def _users_referencing(team: str) -> frozenset[str]: - rows: Final = read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE %s = ANY(teams)', (team,)) - return frozenset(string_value(row["user_id"]) for row in rows) - - -def _assert_landed(candidate: Gateway, team: str, users: Sequence[str]) -> None: - assert _member_ids(candidate.get(f"/scim/v2/Groups/{team}")) == frozenset(users) - assert _membership_user_ids(team) == frozenset(users) - assert _users_referencing(team) == frozenset(users) - - -@pytest.mark.timeout(300) # owned proxy boot plus three 500-member pushes +@pytest.mark.timeout(OWNED_PROXY_TIMEOUT) def test_group_pushes_of_500_members_run_as_many_statements_as_pushes_of_5( gateway: Gateway, tmp_path: Path, record_property: Callable[[str, object], None] ) -> None: @@ -191,39 +132,50 @@ def test_group_pushes_of_500_members_run_as_many_statements_as_pushes_of_5( ) as owned, ): candidate: Final = owned.gateway - small_a: Final = _seed_users(scenario, SMALL_PUSH) - small_b: Final = _seed_users(scenario, SMALL_PUSH) - large: Final = _seed_users(scenario, LARGE_PUSH) + small_a: Final = seed_users(scenario, SMALL_PUSH) + small_b: Final = seed_users(scenario, SMALL_PUSH) + large: Final = seed_users(scenario, LARGE_PUSH) - created_small: Final = _measured(relay, lambda: _create_group(candidate, small_a)) - small_team: Final = _created_team(scenario, created_small.response) - created_large: Final = _measured(relay, lambda: _create_group(candidate, large)) - large_team: Final = _created_team(scenario, created_large.response) - assert _member_ids(_body(created_large.response)) == frozenset(large) - _assert_landed(candidate, small_team, small_a) - _assert_landed(candidate, large_team, large) + created_small: Final = _measured(relay, lambda: create_group(candidate, small_a)) + small_team: Final = created_team(scenario, created_small.response) + created_large: Final = _measured(relay, lambda: create_group(candidate, large)) + large_team: Final = created_team(scenario, created_large.response) + assert member_ids(body(created_large.response)) == frozenset(large) + assert_landed(candidate, small_team, small_a) + assert_landed(candidate, large_team, large) detach_window_open: Final = _database_utc_now() - replaced_small: Final = _measured(relay, lambda: _replace_group(candidate, small_team, small_b)) - replaced_large: Final = _measured(relay, lambda: _replace_group(candidate, large_team, small_a)) + replaced_small: Final = _measured(relay, lambda: replace_group(candidate, small_team, small_b)) + replaced_large: Final = _measured(relay, lambda: replace_group(candidate, large_team, small_a)) assert replaced_small.response.status_code == 200, replaced_small.response.text assert replaced_large.response.status_code == 200, replaced_large.response.text - assert _member_ids(_body(replaced_small.response)) == frozenset(small_b) - _assert_landed(candidate, small_team, small_b) - _assert_landed(candidate, large_team, small_a) + assert member_ids(body(replaced_small.response)) == frozenset(small_b) + assert_landed(candidate, small_team, small_b) + assert_landed(candidate, large_team, small_a) assert detach_window_open <= _updated_at(large[0]) <= _database_utc_now() - patched_small: Final = _measured(relay, lambda: _add_to_group(candidate, small_team, small_a)) - patched_large: Final = _measured(relay, lambda: _add_to_group(candidate, large_team, large)) + patched_small: Final = _measured(relay, lambda: add_to_group(candidate, small_team, small_a)) + patched_large: Final = _measured(relay, lambda: add_to_group(candidate, large_team, large)) assert patched_small.response.status_code == 200, patched_small.response.text assert patched_large.response.status_code == 200, patched_large.response.text - _assert_landed(candidate, small_team, [*small_a, *small_b]) - _assert_landed(candidate, large_team, [*small_a, *large]) + assert_landed(candidate, small_team, [*small_a, *small_b]) + assert_landed(candidate, large_team, [*small_a, *large]) + + with ThreadPoolExecutor(max_workers=KEY_WRITERS) as pool: + small_keys: Final = tuple(pool.map(partial(team_key, candidate, small_team), (*small_a, *small_b))) + large_keys: Final = tuple(pool.map(partial(team_key, candidate, large_team), (*small_a, *large))) + deleted_small: Final = _measured(relay, lambda: delete_group(candidate, small_team)) + deleted_large: Final = _measured(relay, lambda: delete_group(candidate, large_team)) + assert deleted_small.response.status_code == 204, deleted_small.response.text + assert deleted_large.response.status_code == 204, deleted_large.response.text + assert_gone(candidate, small_team, [*small_a, *small_b], small_keys[0]) + assert_gone(candidate, large_team, [*small_a, *large], large_keys[-1]) pushes: Final = ( ("POST", created_small, created_large), ("PUT", replaced_small, replaced_large), ("PATCH", patched_small, patched_large), + ("DELETE", deleted_small, deleted_large), ) for verb, small, big in pushes: record_property( @@ -238,13 +190,13 @@ def test_group_pushes_of_500_members_run_as_many_statements_as_pushes_of_5( def _trade_rosters(candidate: Gateway, teams: tuple[str, str], rosters: tuple[Sequence[str], Sequence[str]]) -> None: with ThreadPoolExecutor(max_workers=2) as pool: - left: Final = pool.submit(_replace_group, candidate, teams[0], rosters[1]) - right: Final = pool.submit(_replace_group, candidate, teams[1], rosters[0]) + left: Final = pool.submit(replace_group, candidate, teams[0], rosters[1]) + right: Final = pool.submit(replace_group, candidate, teams[1], rosters[0]) for response in (left.result(), right.result()): assert response.status_code == 200, response.text -@pytest.mark.timeout(300) +@pytest.mark.timeout(OWNED_PROXY_TIMEOUT) def test_two_groups_trading_rosters_in_parallel_both_land(gateway: Gateway, tmp_path: Path) -> None: with ( gateway.scenario() as scenario, @@ -258,15 +210,18 @@ def test_two_groups_trading_rosters_in_parallel_both_land(gateway: Gateway, tmp_ ) as owned, ): candidate: Final = owned.gateway - left_users: Final = _seed_users(scenario, SWAP_GROUP) - right_users: Final = _seed_users(scenario, SWAP_GROUP) - left: Final = _created_team(scenario, _create_group(candidate, left_users)) - right: Final = _created_team(scenario, _create_group(candidate, right_users)) - for round_index in range(SWAP_ROUNDS): - held: Final = (left_users, right_users) if round_index % 2 == 0 else (right_users, left_users) + left_users: Final = seed_users(scenario, SWAP_GROUP) + right_users: Final = seed_users(scenario, SWAP_GROUP) + left: Final = created_team(scenario, create_group(candidate, left_users)) + right: Final = created_team(scenario, create_group(candidate, right_users)) + rounds: Final = tuple( + (left_users, right_users) if round_index % 2 == 0 else (right_users, left_users) + for round_index in range(SWAP_ROUNDS) + ) + for held in rounds: _trade_rosters(candidate, (left, right), held) - _assert_landed(candidate, left, left_users) - _assert_landed(candidate, right, right_users) + assert_landed(candidate, left, left_users) + assert_landed(candidate, right, right_users) def _int_value(rows: Sequence[Mapping[str, JsonValue]], column: str) -> int: @@ -286,6 +241,10 @@ def _sessions_blocked_by(backend_pid: int) -> int: return _int_value(read_rows(SESSIONS_BLOCKED_BY_SQL, (str(backend_pid),)), "waiting") +def _backend_blocked_by(backend_pid: int) -> int: + return _int_value(read_rows(BACKENDS_BLOCKED_BY_SQL, (str(backend_pid),)), "pid") + + def _database_now() -> str: return string_value(read_rows("SELECT now()::text AS now", ())[0]["now"]) @@ -302,36 +261,25 @@ def _roster_ids(team: str) -> tuple[str, ...]: return tuple(string_value(object_value(member)["user_id"]) for member in roster) -def _delete_team_if_present(candidate: Gateway, team: str) -> None: - if read_rows(TEAM_ROW_SQL, (team,)): - _ = candidate.post("/team/delete", {"team_ids": [team]}) - - -def _group_the_test_deletes(candidate: Gateway, scenario: Scenario) -> str: - created: Final = _create_group(candidate, ()) - assert created.status_code == 201, created.text - team: Final = string_value(_body(created)["id"]) - scenario.cleanups.callback(_delete_team_if_present, candidate, team) - return team - - -def test_a_group_deleted_while_its_push_waits_on_a_member_row_answers_404_and_writes_nothing( +def test_a_group_deleted_while_its_push_waits_on_a_member_row_waits_for_the_push_then_removes_every_trace( gateway: Gateway, ) -> None: - with gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=1) as pool: - users: Final = _seed_users(scenario, SMALL_PUSH) - team: Final = _group_the_test_deletes(gateway, scenario) + with gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=2) as pool: + users: Final = seed_users(scenario, SMALL_PUSH) + team: Final = created_team(scenario, create_group(gateway, ())) with _holding_row_lock(LOCK_USER_ROW_SQL, min(users)) as holder_pid: - push: Final = pool.submit(_replace_group, gateway, team, users) + push: Final = pool.submit(replace_group, gateway, team, users) _ = eventually(lambda: _sessions_blocked_by(holder_pid), lambda waiting: waiting >= 1, seconds=30) - deleted: Final = gateway.request("DELETE", f"/scim/v2/Groups/{team}") - assert deleted.status_code == 204, deleted.text + push_pid: Final = _backend_blocked_by(holder_pid) + deleted: Final = pool.submit(delete_group, gateway, team) + _ = eventually(lambda: _sessions_blocked_by(push_pid), lambda waiting: waiting >= 1, seconds=30) assert not push.done(), push.result().text - pushed: Final = push.result() - assert pushed.status_code == 404, pushed.text + assert not deleted.done(), deleted.result().text + assert push.result().status_code == 200, push.result().text + assert deleted.result().status_code == 204, deleted.result().text assert read_rows(TEAM_ROW_SQL, (team,)) == [] - assert _membership_user_ids(team) == frozenset() - assert _users_referencing(team) == frozenset() + assert membership_user_ids(team) == frozenset() + assert users_referencing(team) == frozenset() @contextmanager @@ -361,17 +309,17 @@ def test_a_push_whose_membership_write_fails_answers_5xx_keeps_the_old_roster_an gateway: Gateway, ) -> None: with gateway.scenario() as scenario: - kept: Final = _seed_users(scenario, SMALL_PUSH) - arriving: Final = _seed_users(scenario, SMALL_PUSH) - team: Final = _created_team(scenario, _create_group(gateway, kept)) + kept: Final = seed_users(scenario, SMALL_PUSH) + arriving: Final = seed_users(scenario, SMALL_PUSH) + team: Final = created_team(scenario, create_group(gateway, kept)) replacement: Final = (*kept[:2], *arriving) with _refusing_memberships_for(max(arriving)): - refused: Final = _replace_group(gateway, team, replacement) + refused: Final = replace_group(gateway, team, replacement) assert 500 <= refused.status_code < 600, refused.text - _assert_landed(gateway, team, kept) - retried: Final = _replace_group(gateway, team, replacement) + assert_landed(gateway, team, kept) + retried: Final = replace_group(gateway, team, replacement) assert retried.status_code == 200, retried.text - _assert_landed(gateway, team, replacement) + assert_landed(gateway, team, replacement) def _five_kilobyte_member(scenario: Scenario) -> str: @@ -394,8 +342,8 @@ def test_a_push_with_an_odd_member_value_never_answers_500_and_keeps_the_roster_ gateway: Gateway, shape: str ) -> None: with gateway.scenario() as scenario: - kept: Final = _seed_users(scenario, SMALL_PUSH) - team: Final = _created_team(scenario, _create_group(gateway, kept)) + kept: Final = seed_users(scenario, SMALL_PUSH) + team: Final = created_team(scenario, create_group(gateway, kept)) odd: Final = ODD_MEMBER_VALUES[shape](scenario, kept) pushed: Final = gateway.request( "PUT", @@ -404,14 +352,14 @@ def test_a_push_with_an_odd_member_value_never_answers_500_and_keeps_the_roster_ "schemas": [GROUP_SCHEMA], "id": team, "displayName": f"integration-{uuid.uuid4().hex}", - "members": [*_members(kept), {"value": odd}], + "members": [*members(kept), {"value": odd}], }, ) assert pushed.status_code < 500, pushed.text roster: Final = _roster_ids(team) assert len(roster) == len(frozenset(roster)), roster assert frozenset(kept) <= frozenset(roster), roster - _assert_landed(gateway, team, roster) + assert_landed(gateway, team, roster) if pushed.status_code >= 400: assert frozenset(roster) == frozenset(kept), (pushed.text, roster) @@ -443,15 +391,15 @@ PATCH_SHAPES: Final[Mapping[str, PatchShape]] = { @pytest.mark.parametrize("shape", tuple(PATCH_SHAPES)) def test_patch_member_operation_shapes_land_the_expected_roster(gateway: Gateway, shape: str) -> None: with gateway.scenario() as scenario: - users: Final = _seed_users(scenario, 4) - team: Final = _created_team(scenario, _create_group(gateway, users[:3])) + users: Final = seed_users(scenario, 4) + team: Final = created_team(scenario, create_group(gateway, users[:3])) patched: Final = gateway.request( "PATCH", f"/scim/v2/Groups/{team}", {"schemas": [PATCH_OP_SCHEMA], "Operations": PATCH_SHAPES[shape].operations(users)}, ) assert patched.status_code == 200, patched.text - _assert_landed(gateway, team, PATCH_SHAPES[shape].expected(users)) + assert_landed(gateway, team, PATCH_SHAPES[shape].expected(users)) def _team_key_for(scenario: Scenario, team: str, user_id: str) -> str: @@ -462,7 +410,7 @@ def _team_key_for(scenario: Scenario, team: str, user_id: str) -> str: def _push_outcome(candidate: Gateway, team: str, users: Sequence[str]) -> str: try: - return str(_replace_group(candidate, team, users).status_code) + return str(replace_group(candidate, team, users).status_code) except httpx.TransportError as error: return type(error).__name__ @@ -484,9 +432,9 @@ def test_a_worker_killed_mid_push_leaves_the_old_roster_and_a_retry_lands_the_ne ThreadPoolExecutor(max_workers=1) as pool, ): candidate: Final = owned.gateway - leaving: Final = _seed_users(scenario, CHAOS_GROUP) - arriving: Final = _seed_users(scenario, CHAOS_GROUP) - team: Final = _created_team(scenario, _create_group(candidate, leaving)) + leaving: Final = seed_users(scenario, CHAOS_GROUP) + arriving: Final = seed_users(scenario, CHAOS_GROUP) + team: Final = created_team(scenario, create_group(candidate, leaving)) pinned_key: Final = _team_key_for(scenario, team, max(leaving)) with _holding_row_lock(LOCK_KEY_ROW_SQL, sha256(pinned_key.encode()).hexdigest()) as holder_pid: push: Final = pool.submit(_push_outcome, candidate, team, arriving) @@ -508,11 +456,11 @@ def test_a_worker_killed_mid_push_leaves_the_old_roster_and_a_retry_lands_the_ne {"leaving": len(settled & frozenset(leaving)), "arriving": len(settled & frozenset(arriving))}, ) assert settled in (frozenset(leaving), frozenset(arriving)), sorted(settled) - _assert_landed(candidate, team, sorted(settled)) + assert_landed(candidate, team, sorted(settled)) - retried: Final = _replace_group(candidate, team, arriving) + retried: Final = replace_group(candidate, team, arriving) assert retried.status_code == 200, retried.text - _assert_landed(candidate, team, arriving) + assert_landed(candidate, team, arriving) assert ( read_rows( 'SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s', diff --git a/tests/integration/management/test_scim_group_delete.py b/tests/integration/management/test_scim_group_delete.py new file mode 100644 index 00000000000..4a4a93f8bd2 --- /dev/null +++ b/tests/integration/management/test_scim_group_delete.py @@ -0,0 +1,431 @@ +"""DELETE /scim/v2/Groups/{id} deletes the group's team the way /team/delete does. + +The members' group entries, their membership rows and the keys the members hold on the team go in one locked +pass whatever the roster size, and the configured admin group's members are demoted in that same pass, so they stay +admins until the group is gone. A member a concurrent group write added while the delete waited on the team lock is +demoted with the rest, and a group write that arrives once the delete holds the lock finds the group gone. The group +is 404 after, on every worker and on the peer proxy, and an unrelated key keeps serving. +A service-account key and a model on the team stay, so re-creating the group with the same id brings them back. +""" + +import os +import re +import uuid +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from functools import partial +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import GATEWAY_LIMITS, JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.scim import ( + GROUP_SCHEMA, + add_to_group, + assert_gone, + assert_landed, + body, + create_group, + created_team, + delete_group, + groups_of, + held_team_lock, + held_user_row, + keys_of, + member_ids, + membership_user_ids, + replace_group, + seed_users, + service_account_key, + team_key, + user_role, + waiters_on_lock_held_by, + waiters_on_rows_locked_by, +) +from pydantic import JsonValue, TypeAdapter + +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + +pytestmark: Final = pytest.mark.shares_worker_state + +PROXY_ADMIN: Final = "proxy_admin" +DEMOTED_ROLE: Final = "internal_user_viewer" +MASTER_KEY_USER: Final = "default_user_id" +BURST_GROUPS: Final = 20 +OWNED_PROXY_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +DELETED_TEAM_RECORDS_SQL: Final = 'SELECT team_id FROM "LiteLLM_DeletedTeamTable" WHERE team_id = %s' +DELETED_KEY_RECORDS_SQL: Final = 'SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s' +TEAM_MODEL_IDS_SQL: Final = """SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_info->>'team_id' = %s""" +BACKENDS_BLOCKED_BY_SQL: Final = "SELECT pid FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))" +USER_ROLE_SQL: Final = 'SELECT user_role FROM "LiteLLM_UserTable" WHERE user_id = %s' +DELETED_AUDIT_ROWS_SQL: Final = """ +SELECT table_name, action, object_id, changed_by, changed_by_api_key +FROM "LiteLLM_AuditLog" +WHERE object_id IN (%s, %s) AND action = 'deleted' +ORDER BY table_name +""" +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_PIDS: Final = TypeAdapter(tuple[int, ...]) + + +def _hashed(key: str) -> str: + return sha256(key.encode()).hexdigest() + + +def _admin_group_config(tmp_path: Path, admin_alias: str) -> Path: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + settings: Final = {**object_value(config["litellm_settings"]), "scim_admin_group": admin_alias} + path: Final = tmp_path / "scim-admin-group.yaml" + path.write_text(yaml.safe_dump({**config, "litellm_settings": settings})) + return path + + +def _roles_of(candidate: Gateway, users: tuple[str, ...]) -> frozenset[str]: + return frozenset(user_role(candidate, user) for user in users) + + +def _backends_blocked_by(holder: int) -> tuple[int, ...]: + return _PIDS.validate_python([row["pid"] for row in read_rows(BACKENDS_BLOCKED_BY_SQL, (str(holder),))]) + + +def _stored_roles(users: tuple[str, ...]) -> frozenset[str]: + return frozenset(string_value(read_rows(USER_ROLE_SQL, (user,))[0]["user_role"]) for user in users) + + +def _models_status(candidate: Gateway, key: str) -> int: + return candidate.request("GET", "/v1/models", key=key).status_code + + +def _team_model_ids(team: str) -> frozenset[str]: + return frozenset(string_value(row["model_id"]) for row in read_rows(TEAM_MODEL_IDS_SQL, (team,))) + + +def _assert_revoked_everywhere(proxies: tuple[Gateway, ...], key: str) -> None: + for candidate in proxies: + eventually(partial(_models_status, candidate, key), lambda code: code == 401, seconds=30) + + +def _deleted_audit_rows(team: str, hashed_key: str) -> list[dict[str, JsonValue]]: + return read_rows(DELETED_AUDIT_ROWS_SQL, (team, hashed_key)) + + +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]: + with httpx.Client( + base_url=candidate.client.base_url, timeout=candidate.client.timeout, trust_env=False, limits=GATEWAY_LIMITS + ) as client: + yield Gateway(client, candidate.key, candidate.upstream_url) + + +def test_deleting_one_group_leaves_the_members_other_group_intact(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + users: Final = seed_users(scenario, 3) + doomed: Final = created_team(scenario, create_group(gateway, users)) + kept: Final = created_team(scenario, create_group(gateway, users)) + kept_key: Final = team_key(gateway, kept, users[0]) + + deleted: Final = delete_group(gateway, doomed) + assert deleted.status_code == 204, deleted.text + + assert gateway.request("GET", f"/scim/v2/Groups/{doomed}").status_code == 404 + assert_landed(gateway, kept, users) + assert groups_of(gateway, users[0]) == frozenset({kept}) + assert keys_of(kept) == frozenset({_hashed(kept_key)}) + assert _models_status(gateway, kept_key) == 200 + + +@pytest.mark.timeout(OWNED_PROXY_TIMEOUT) +def test_deleting_the_admin_group_demotes_its_members_and_writes_an_audit_row_naming_the_scim_caller( + gateway: Gateway, tmp_path: Path +) -> None: + admin_alias: Final = f"integration-admins-{uuid.uuid4().hex}" + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"}, + config=_admin_group_config(tmp_path, admin_alias), + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + ) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + users: Final = seed_users(scenario, 3) + team: Final = created_team(scenario, create_group(candidate, users, display_name=admin_alias)) + assert _roles_of(candidate, users) == frozenset({PROXY_ADMIN}) + key: Final = team_key(candidate, team, users[0]) + + deleted: Final = delete_group(candidate, team) + assert deleted.status_code == 204, deleted.text + + assert_gone(candidate, team, users, key) + assert _roles_of(candidate, users) == frozenset({DEMOTED_ROLE}) + rows: Final = eventually(lambda: _deleted_audit_rows(team, _hashed(key)), lambda found: len(found) >= 2, 30) + caller: Final = {"changed_by": MASTER_KEY_USER, "changed_by_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS} + assert rows == [ + {"table_name": "LiteLLM_TeamTable", "action": "deleted", "object_id": team, **caller}, + {"table_name": "LiteLLM_VerificationToken", "action": "deleted", "object_id": _hashed(key), **caller}, + ] + + +@pytest.mark.timeout(OWNED_PROXY_TIMEOUT) +def test_a_member_added_while_the_admin_group_delete_waits_on_the_team_lock_is_demoted_too( + gateway: Gateway, tmp_path: Path +) -> None: + admin_alias: Final = f"integration-admins-{uuid.uuid4().hex}" + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": os.environ["DATABASE_URL"]}, + config=_admin_group_config(tmp_path, admin_alias), + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + client_timeout=60, + ) as owned, + owned.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=2) as pool, + ): + candidate: Final = owned.gateway + users: Final = seed_users(scenario, 5) + (late,) = seed_users(scenario, 1) + team: Final = created_team(scenario, create_group(candidate, users, display_name=admin_alias)) + assert _roles_of(candidate, users) == frozenset({PROXY_ADMIN}) + key: Final = team_key(candidate, team, users[0]) + + with held_team_lock(team) as holder: + put: Final = pool.submit(replace_group, candidate, team, (*users, late), display_name=admin_alias) + eventually(lambda: waiters_on_lock_held_by(holder), lambda waiting: waiting == 1, seconds=30) + delete: Final = pool.submit(delete_group, candidate, team) + eventually(lambda: waiters_on_lock_held_by(holder), lambda waiting: waiting == 2, seconds=30) + assert _stored_roles(users) == frozenset({PROXY_ADMIN}) + replaced: Final = put.result() + deleted: Final = delete.result() + + assert replaced.status_code == 200, replaced.text + assert member_ids(body(replaced)) == frozenset((*users, late)) + assert deleted.status_code == 204, deleted.text + assert_gone(candidate, team, (*users, late), key) + assert _roles_of(candidate, (*users, late)) == frozenset({DEMOTED_ROLE}) + + +@pytest.mark.timeout(OWNED_PROXY_TIMEOUT) +def test_a_member_added_while_the_admin_group_delete_holds_the_team_lock_is_refused_and_never_made_admin( + gateway: Gateway, tmp_path: Path +) -> None: + admin_alias: Final = f"integration-admins-{uuid.uuid4().hex}" + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": os.environ["DATABASE_URL"]}, + config=_admin_group_config(tmp_path, admin_alias), + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + client_timeout=60, + ) as owned, + owned.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=2) as pool, + ): + candidate: Final = owned.gateway + users: Final = seed_users(scenario, 5) + (late,) = seed_users(scenario, 1) + team: Final = created_team(scenario, create_group(candidate, users, display_name=admin_alias)) + assert _roles_of(candidate, users) == frozenset({PROXY_ADMIN}) + key: Final = team_key(candidate, team, users[0]) + + with held_user_row(users[0]) as holder: + delete: Final = pool.submit(delete_group, candidate, team) + eventually(lambda: waiters_on_rows_locked_by(holder), lambda waiting: waiting == 1, seconds=30) + (deleting,) = _backends_blocked_by(holder) + added: Final = pool.submit(add_to_group, candidate, team, (late,)) + eventually(lambda: _backends_blocked_by(deleting), lambda blocked: len(blocked) == 1, seconds=30) + assert _stored_roles(users) == frozenset({PROXY_ADMIN}) + deleted: Final = delete.result() + + assert deleted.status_code == 204, deleted.text + assert added.result().status_code == 404, added.result().text + assert_gone(candidate, team, users, key) + assert _roles_of(candidate, users) == frozenset({DEMOTED_ROLE}) + assert user_role(candidate, late) != PROXY_ADMIN + + +def test_a_group_delete_revokes_member_keys_and_a_re_created_group_gets_its_service_account_key_and_model_back( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + users: Final = seed_users(scenario, 2) + team: Final = created_team(scenario, create_group(gateway, users)) + member_key: Final = team_key(gateway, team, users[0]) + service_key: Final = service_account_key(gateway, scenario, team) + scenario.model(model_info={"team_id": team}) + team_models: Final = _team_model_ids(team) + assert team_models != frozenset() + + deleted: Final = delete_group(gateway, team) + assert deleted.status_code == 204, deleted.text + + assert gateway.request("GET", f"/scim/v2/Groups/{team}").status_code == 404 + assert membership_user_ids(team) == frozenset() + assert keys_of(team) == frozenset({_hashed(service_key)}) + assert _team_model_ids(team) == team_models + _assert_revoked_everywhere((gateway,), member_key) + + recreated: Final = gateway.request( + "POST", + "/scim/v2/Groups", + {"schemas": [GROUP_SCHEMA], "externalId": team, "displayName": f"integration-{uuid.uuid4().hex}"}, + ) + assert recreated.status_code == 201, recreated.text + assert string_value(body(recreated)["id"]) == team + eventually(partial(_models_status, gateway, service_key), lambda code: code == 200, seconds=90) + assert _models_status(gateway, member_key) == 401 + + +def test_deleting_an_unknown_group_is_404(gateway: Gateway) -> None: + missing: Final = f"missing-{uuid.uuid4().hex}" + response: Final = delete_group(gateway, missing) + assert response.status_code == 404, response.text + assert "error" in body(response), response.text + assert missing in response.text + + +def test_deleting_a_group_twice_is_404_the_second_time(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + users: Final = seed_users(scenario, 2) + team: Final = created_team(scenario, create_group(gateway, users)) + first: Final = delete_group(gateway, team) + assert first.status_code == 204, first.text + second: Final = delete_group(gateway, team) + assert second.status_code == 404, second.text + assert "error" in body(second), second.text + assert gateway.request("GET", f"/scim/v2/Groups/{team}").status_code == 404 + assert membership_user_ids(team) == frozenset() + + +def test_two_deletes_racing_on_one_group_answer_204_then_404_and_record_the_team_once(gateway: Gateway) -> None: + with gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=2) as pool: + users: Final = seed_users(scenario, 2) + team: Final = created_team(scenario, create_group(gateway, users)) + key: Final = team_key(gateway, team, users[0]) + + with held_team_lock(team) as holder: + deletes: Final = (pool.submit(delete_group, gateway, team), pool.submit(delete_group, gateway, team)) + eventually(lambda: waiters_on_lock_held_by(holder), lambda waiting: waiting == 2, seconds=30) + + assert sorted(delete.result().status_code for delete in deletes) == [204, 404] + assert len(read_rows(DELETED_TEAM_RECORDS_SQL, (team,))) == 1 + assert len(read_rows(DELETED_KEY_RECORDS_SQL, (_hashed(key),))) == 1 + assert gateway.request("GET", f"/scim/v2/Groups/{team}").status_code == 404 + + +def test_deleting_a_group_with_a_non_admin_key_is_refused(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + users: Final = seed_users(scenario, 2) + team: Final = created_team(scenario, create_group(gateway, users)) + outsider: Final = scenario.key() + response: Final = delete_group(gateway, team, key=outsider) + assert response.status_code in (401, 403), response.text + assert "error" in body(response), response.text + assert_landed(gateway, team, users) + + +@pytest.mark.parametrize("group_id", ("123", "x" * 5000, "a%2Fb"), ids=("int_like", "five_kb", "encoded_slash")) +def test_deleting_a_group_with_a_hostile_id_is_404(gateway: Gateway, group_id: str) -> None: + response: Final = delete_group(gateway, group_id) + assert response.status_code == 404, response.text + assert isinstance(response.json(), dict), response.text + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def test_deleting_an_empty_group_is_204_then_404(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + team: Final = created_team(scenario, create_group(gateway, ())) + assert member_ids(gateway.get(f"/scim/v2/Groups/{team}")) == frozenset() + deleted: Final = delete_group(gateway, team) + assert deleted.status_code == 204, deleted.text + assert gateway.request("GET", f"/scim/v2/Groups/{team}").status_code == 404 + assert membership_user_ids(team) == frozenset() + + +def test_group_delete_on_one_proxy_revokes_the_team_key_and_the_group_on_the_peer( + gateway: Gateway, peer: Gateway +) -> None: + with gateway.scenario() as scenario: + users: Final = seed_users(scenario, 3) + team: Final = created_team(scenario, create_group(gateway, users)) + key: Final = team_key(gateway, team, users[0]) + assert _models_status(peer, key) == 200 + assert member_ids(peer.get(f"/scim/v2/Groups/{team}")) == frozenset(users) + + deleted: Final = delete_group(gateway, team) + assert deleted.status_code == 204, deleted.text + + assert peer.request("GET", f"/scim/v2/Groups/{team}").status_code == 404 + _assert_revoked_everywhere((peer, gateway), key) + assert_gone(peer, team, users, key) + + +def test_twenty_group_deletes_split_across_two_proxies_all_land(gateway: Gateway, peer: Gateway) -> None: + proxies: Final = (gateway, peer) + with gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=10) as pool: + bystander: Final = scenario.key() + rosters: Final = tuple(seed_users(scenario, 3) for _ in range(BURST_GROUPS)) + teams: Final = tuple(created_team(scenario, create_group(gateway, users)) for users in rosters) + keys: Final = tuple(team_key(gateway, team, users[0]) for team, users in zip(teams, rosters)) + + assigned: Final = tuple(proxies[index % 2] for index in range(BURST_GROUPS)) + responses: Final = tuple(pool.map(delete_group, assigned, teams)) + + for response in responses: + assert response.status_code == 204, response.text + for team, users, key in zip(teams, rosters, keys): + _assert_revoked_everywhere(proxies, key) + assert_gone(gateway, team, users, key) + assert peer.request("GET", f"/scim/v2/Groups/{team}").status_code == 404 + assert _models_status(gateway, bystander) == 200 + assert _models_status(peer, bystander) == 200 + + +@pytest.mark.timeout(OWNED_PROXY_TIMEOUT) +def test_group_delete_after_a_worker_was_killed_still_completes(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process( + gateway, tmp_path, {}, remove_environment=("DATABASE_URL_READ_REPLICA",), workers=2 + ) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + workers, _ = eventually( + lambda: _worker_startups(owned.log), lambda found: len(found[0]) == 2 and found[1] == 2, seconds=120 + ) + first_users: Final = seed_users(scenario, 3) + second_users: Final = seed_users(scenario, 3) + first: Final = created_team(scenario, create_group(candidate, first_users)) + second: Final = created_team(scenario, create_group(candidate, second_users)) + first_key: Final = team_key(candidate, first, first_users[0]) + second_key: Final = team_key(candidate, second, second_users[0]) + + with _fresh_connection(candidate) as survivor: + with owned.replacing_worker(workers[0]) as kill: + kill.freeze() + kill() + deleted_first: Final = delete_group(survivor, first) + assert deleted_first.status_code == 204, deleted_first.text + assert_gone(survivor, first, first_users, first_key) + + deleted_second: Final = delete_group(survivor, second) + assert deleted_second.status_code == 204, deleted_second.text + assert_gone(survivor, second, second_users, second_key) diff --git a/tests/integration/mcp/test_mcp_transports.py b/tests/integration/mcp/test_mcp_transports.py index 928aa9ff121..ef240cba570 100644 --- a/tests/integration/mcp/test_mcp_transports.py +++ b/tests/integration/mcp/test_mcp_transports.py @@ -2,10 +2,12 @@ import builtins import json import uuid from collections.abc import Mapping +from pathlib import Path from typing import Final import httpx import pytest +import yaml from mcp.shared.exceptions import MCPError from integration._support.client import Gateway from integration._support.database import read_rows @@ -17,15 +19,21 @@ from integration._support.mcp import ( McpCaller, Outcome, PeerKind, + ScriptedTool, _outcome_from_rpc, _parse_rpc_body, + call_tool, + listed_tools, mcp_peer, official_client_outcomes, peer_of, register_mcp, + scripted_peer, + text_result, tool_calls, ) from integration._support.mcp_grants import create_toolset +from integration._support.process import owned_proxy ADD: Final = {"http": "add", "sse": "add", "stdio": "add", "openapi": "getpet"} ARGUMENTS: Final = {"add": {"a": 3, "b": 4}, "getpet": {"petId": "7"}} @@ -817,3 +825,67 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session( ), ) as gateway: asyncio.run(exercise(gateway)) + + +def test_tool_result_larger_than_one_mib_in_a_single_sse_event_reaches_the_caller(gateway: Gateway) -> None: + payload: Final = "x" * (3 * 1024 * 1024) + tool: Final = ScriptedTool("dump", lambda _: text_result(payload), event_stream=True) + + with scripted_peer(tool) as peer, gateway.scenario() as scenario: + alias: Final = "bigsse" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + name: Final = next(full for full in listed_tools(gateway, key, identity) if full.endswith("dump")) + response: Final = call_tool(gateway, key, identity, name, {}) + assert response.status_code == 200, response.text[:500] + body: Final = response.json() + assert body["isError"] is False, response.text[:500] + assert body["content"][0]["text"] == payload + + +def test_configured_sse_limit_rejects_large_results_and_allows_the_next_call(tmp_path: Path) -> None: + limit: Final = 4096 + payload: Final = "x" * 512 + small: Final = ScriptedTool("small", lambda _: text_result(payload), event_stream=True) + large: Final = ScriptedTool("large", lambda _: text_result("x" * (2 * limit)), event_stream=True) + key: Final = "sk-local-mcp-sse-limit-regression-key" + with scripted_peer(small, large) as peer, httpx.Client() as client: + config: Final = tmp_path / "sse-limit.yaml" + config.write_text( + yaml.safe_dump( + { + "general_settings": {"master_key": key, "store_model_in_db": False}, + "mcp_servers": {"limited": peer.registration()}, + } + ) + ) + seed: Final = Gateway(client, key, peer.url) + with owned_proxy( + seed, + tmp_path / "proxy", + { + "STORE_MODEL_IN_DB": "False", + "DISABLE_SCHEMA_UPDATE": "true", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES": str(limit), + }, + config=config, + database_setup=(), + remove_environment=( + "DATABASE_URL", + "DATABASE_URL_READ_REPLICA", + "DIRECT_URL", + "INTEGRATION_PROXY_DATABASE_URL", + "INTEGRATION_PROXY_READ_REPLICA_URL", + ), + ) as gateway: + assert set(listed_tools(gateway, key)) == {"small", "large"} + for name in ("small", "large", "small"): + response: Final = call_tool(gateway, key, "limited", name, {}) + assert response.status_code == 200, response.text[:500] + body: Final = response.json() + assert body["isError"] is (name == "large"), response.text[:500] + if name == "large": + assert str(limit) in body["content"][0]["text"], response.text[:500] + else: + assert body["content"][0]["text"] == payload, response.text[:500] diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index b46503703c7..58785c89033 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -312,7 +312,7 @@ def _chat_caller_stream_usage_event(event: dict[str, JsonValue], model: str) -> "object": event["object"], "created": 1, "model": model, - "choices": [{"index": 0, "delta": {}}], + "choices": [], "usage": {**_json_object_value(event["usage"]), "cost": 4.05e-6}, } diff --git a/tests/integration/observability/test_custom_logger_hook_counts.py b/tests/integration/observability/test_custom_logger_hook_counts.py new file mode 100644 index 00000000000..420ac2a5fcb --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts.py @@ -0,0 +1,96 @@ +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.hook_counts import ( + STREAMING_SURFACES, + SURFACE_CALL_TYPES, + HookRecords, + assert_hooks, + one_async_hook, + provider, + send, + write_hook_count_config, +) +from integration._support.process import owned_proxy +from integration._support.wire import Reply, wire_server + +pytestmark = pytest.mark.timeout(240) + + +@pytest.fixture(scope="module", params=["cache", "cache_and_sync_callback"]) +def hook_proxy( + request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[tuple[Gateway, HookRecords]]: + tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-counts") + legacy_callbacks: Final = ( + { + "success_callback": ["integration._support.hook_count_probe.sync_success_function"], + "failure_callback": ["integration._support.hook_count_probe.sync_failure_function"], + } + if request.param == "cache_and_sync_callback" + else {} + ) + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config( + tmp_path / "hook-counts.yaml", + upstream.url, + { + "callbacks": ["integration._support.hook_count_probe.probe"], + **legacy_callbacks, + "cache": True, + "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, + }, + ) + with owned_proxy( + rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as proxy: + yield proxy, HookRecords(sink, function_callbacks=request.param == "cache_and_sync_callback") + + +@pytest.mark.parametrize("outcome", ["success", "failure"]) +@pytest.mark.parametrize("surface", sorted(SURFACE_CALL_TYPES)) +def test_custom_logger_gets_one_async_hook_per_request_with_the_response_cache_on( + hook_proxy: tuple[Gateway, HookRecords], surface: str, outcome: str +) -> None: + proxy, records = hook_proxy + call_id: Final = str(uuid.uuid4()) + response: Final = send(proxy, surface, outcome, call_id) + assert response.status_code == (200 if outcome == "success" else 400), response.text + assert_hooks(records, surface, call_id, one_async_hook(records, surface, outcome)) + + +@pytest.mark.parametrize("outcome", ["success", "failure"]) +@pytest.mark.parametrize("surface", STREAMING_SURFACES) +def test_custom_logger_gets_one_async_hook_per_streamed_request( + hook_proxy: tuple[Gateway, HookRecords], surface: str, outcome: str +) -> None: + proxy, records = hook_proxy + call_id: Final = str(uuid.uuid4()) + response: Final = send(proxy, surface, outcome, call_id, stream=True) + assert response.status_code == (200 if outcome == "success" else 400), response.text + if outcome == "success": + assert "ok" in response.text, response.text + assert_hooks(records, surface, call_id, one_async_hook(records, surface, outcome)) + + +@pytest.mark.parametrize("surface", sorted(SURFACE_CALL_TYPES)) +def test_custom_logger_gets_one_async_hook_per_request_rejected_by_key_auth( + hook_proxy: tuple[Gateway, HookRecords], surface: str +) -> None: + proxy, records = hook_proxy + known_calls: Final = {record["litellm_call_id"] for record in records.drained()} + response: Final = send(proxy, surface, "success", str(uuid.uuid4()), key="sk-not-a-virtual-key") + assert response.status_code == 401, response.text + + def new_calls(found: tuple[dict[str, str], ...]) -> set[str]: + return {record["litellm_call_id"] for record in found if record["litellm_call_id"] not in known_calls} + + (rejected_call,) = eventually(lambda: new_calls(records.drained()), bool, seconds=20) + assert_hooks(records, surface, rejected_call, one_async_hook(records, surface, "failure")) diff --git a/tests/integration/observability/test_custom_logger_hook_counts_chaos.py b/tests/integration/observability/test_custom_logger_hook_counts_chaos.py new file mode 100644 index 00000000000..7b19ef3b1ff --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts_chaos.py @@ -0,0 +1,181 @@ +"""One hook per request under a concurrent burst while the logger's destination stalls and fails, and while a +worker is killed.""" + +import threading +import time +import uuid +from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment, held +from integration._support.hook_counts import ( + STREAMING_SURFACES, + SURFACE_CALL_TYPES, + HookRecords, + assert_hooks, + one_async_hook, + probe_hooks, + provider, + send, + write_hook_count_config, +) +from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +RecordProperty = Callable[[str, object], None] + +_SETTINGS: Final = {"callbacks": ["integration._support.hook_count_probe.probe"]} +_HOLD_SECONDS: Final = 2.0 + + +@dataclass(frozen=True) +class _Call: + surface: str + outcome: str + stream: bool + call_id: str + + +@dataclass(frozen=True) +class _Outcome: + call: _Call + status: int | None + error: str | None + + +def _burst(rounds: int) -> tuple[_Call, ...]: + shapes: Final = ( + *((surface, outcome, False) for surface in sorted(SURFACE_CALL_TYPES) for outcome in ("success", "failure")), + *((surface, "success", True) for surface in STREAMING_SURFACES), + ) + return tuple( + _Call(surface, outcome, stream, str(uuid.uuid4())) for _ in range(rounds) for surface, outcome, stream in shapes + ) + + +def _attempt(proxy: Gateway, call: _Call) -> _Outcome: + try: + response: Final = send(proxy, call.surface, call.outcome, call.call_id, stream=call.stream) + except httpx.HTTPError as error: + return _Outcome(call, None, repr(error)) + return _Outcome(call, response.status_code, None) + + +def _expected_status(call: _Call) -> int: + return 200 if call.outcome == "success" else 400 + + +def _done(futures: tuple[Future[_Outcome], ...]) -> int: + return sum(future.done() for future in futures) + + +class _Destination: + """The probe's sink: answers at once, holds each report until released, or refuses it with a 503.""" + + def __init__(self) -> None: + self.mode = "ok" + self.release = threading.Event() + + def __call__(self, request: Request) -> Reply: + if self.mode == "hold": + self.release.wait(timeout=_HOLD_SECONDS + 1) + if self.mode == "fail": + return Reply(status=503, body=b'{"error":"destination unavailable"}') + return Reply() + + +def test_logger_destination_stalling_then_failing_mid_burst_gets_one_hook_per_request( + tmp_path: Path, record_property: RecordProperty +) -> None: + destination: Final = _Destination() + with ( + wire_server(provider) as upstream, + wire_server(destination) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config(tmp_path / "hook-count-chaos.yaml", upstream.url, _SETTINGS) + with owned_proxy( + rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as proxy: + records: Final = HookRecords(sink, function_callbacks=False) + calls: Final = _burst(2) + destination.mode = "hold" + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + futures: Final = tuple(pool.submit(_attempt, proxy, call) for call in calls) + eventually(lambda: _done(futures), lambda done: done >= 1, seconds=60) + held_liveliness: Final = proxy.client.get("/health/liveliness").status_code + time.sleep(_HOLD_SECONDS) + destination.mode = "fail" + destination.release.set() + time.sleep(_HOLD_SECONDS) + destination.mode = "ok" + outcomes: Final = tuple(future.result(timeout=180) for future in futures) + record_property("burst_size", len(calls)) + assert held_liveliness == 200, held_liveliness + assert [outcome for outcome in outcomes if outcome.status != _expected_status(outcome.call)] == [] + for call in calls: + assert_hooks(records, call.surface, call.call_id, one_async_hook(records, call.surface, call.outcome)) + + +def test_worker_killed_mid_burst_leaves_at_most_one_hook_per_request( + tmp_path: Path, record_property: RecordProperty +) -> None: + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config(tmp_path / "hook-count-kill.yaml", upstream.url, _SETTINGS) + with owned_proxy_process( + rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as owned: + proxy: Final = owned.gateway + records: Final = HookRecords(sink, function_callbacks=False) + victim: Final = min(eventually(owned.worker_pids, lambda found: len(found) == 2, seconds=60)) + calls: Final = _burst(2) + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + futures: Final = tuple(pool.submit(_attempt, proxy, call) for call in calls) + eventually(lambda: _done(futures), lambda done: done >= 1, seconds=60) + with owned.replacing_worker(victim) as kill: + kill() + done_at_kill: Final = _done(futures) + survived: Final = eventually( + lambda: _attempt(proxy, _Call("speech", "success", False, str(uuid.uuid4()))), + lambda outcome: outcome.status == 200, + seconds=30, + ) + outcomes: Final = tuple(future.result(timeout=180) for future in futures) + respawned: Final = _Call("speech", "success", False, str(uuid.uuid4())) + after_respawn: Final = _attempt(proxy, respawned) + completed: Final = tuple(outcome for outcome in outcomes if outcome.status is not None) + interrupted: Final = tuple(outcome for outcome in outcomes if outcome.status is None) + assert after_respawn.status == 200, after_respawn + assert [outcome for outcome in completed if outcome.status != _expected_status(outcome.call)] == [] + for call in (survived.call, respawned): + assert_hooks(records, call.surface, call.call_id, one_async_hook(records, call.surface, call.outcome)) + burst_hooks: Final = held( + lambda: _hooks_of(records, calls), + lambda found: all( + hooks in ([], [f"async_{call.outcome}"]) for call, hooks in zip(calls, found, strict=True) + ), + holding=3, + seconds=20, + ) + record_property("burst_size", len(calls)) + record_property("killed_worker_pid", victim) + record_property("done_at_kill", done_at_kill) + record_property("interrupted_calls", str(tuple(outcome.error for outcome in interrupted))) + record_property("burst_calls_without_a_hook", sum(not hooks for hooks in burst_hooks)) + + +def _hooks_of(records: HookRecords, calls: tuple[_Call, ...]) -> tuple[list[str], ...]: + """A burst call the killed worker served may lose its logging with the worker, whether or not its response + arrived first, so a burst call gets its one async hook or none, never two.""" + return tuple(probe_hooks(records, call.surface, records.for_call(call.call_id)) for call in calls) diff --git a/tests/integration/observability/test_custom_logger_hook_counts_modes.py b/tests/integration/observability/test_custom_logger_hook_counts_modes.py new file mode 100644 index 00000000000..e143559040e --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts_modes.py @@ -0,0 +1,95 @@ +"""A logger the async handler never reaches keeps its sync hook on async calls outside the async flag list. + +A logger with both hooks registered only under `success_callback` and `failure_callback` is never reached by the async +failure handler, so it keeps its sync failure hook. The async success handler does reach it, and gives it its async +success hook instead of the sync one. + +The sync success handler only runs for an async call when a legacy sync callback is configured, so every +registration here carries the legacy function callbacks too.""" + +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.hook_counts import ( + IN_THE_ASYNC_FLAG_LIST, + SURFACE_CALL_TYPES, + HookRecords, + assert_hooks, + provider, + send, + write_hook_count_config, +) +from integration._support.process import owned_proxy +from integration._support.wire import Reply, wire_server + +pytestmark = pytest.mark.timeout(240) + +_PROBES: Final = "integration._support.hook_count_probe" +_REGISTRATIONS: Final[dict[str, dict[str, object]]] = { + "sync_only_logger": { + "callbacks": [f"{_PROBES}.sync_only_probe"], + "success_callback": [f"{_PROBES}.sync_success_function"], + "failure_callback": [f"{_PROBES}.sync_failure_function"], + }, + "sync_event_lists": { + "success_callback": [f"{_PROBES}.probe", f"{_PROBES}.sync_success_function"], + "failure_callback": [f"{_PROBES}.probe", f"{_PROBES}.sync_failure_function"], + }, +} +_OUTSIDE_THE_ASYNC_FLAG_LIST: Final = sorted(SURFACE_CALL_TYPES.keys() - IN_THE_ASYNC_FLAG_LIST) + + +@pytest.fixture(scope="module", params=sorted(_REGISTRATIONS)) +def sync_hook_proxy( + request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[tuple[Gateway, HookRecords, str]]: + tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-count-modes") + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config( + tmp_path / "hook-count-modes.yaml", upstream.url, _REGISTRATIONS[request.param] + ) + with owned_proxy( + rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as proxy: + yield proxy, HookRecords(sink, function_callbacks=True), request.param + + +@pytest.mark.parametrize("outcome", ["success", "failure"]) +@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST) +def test_logger_the_async_handler_skips_gets_one_sync_hook_per_request( + sync_hook_proxy: tuple[Gateway, HookRecords, str], surface: str, outcome: str +) -> None: + proxy, records, registration = sync_hook_proxy + reached_by_the_async_success: Final = registration == "sync_event_lists" and outcome == "success" + call_id: Final = str(uuid.uuid4()) + response: Final = send(proxy, surface, outcome, call_id) + assert response.status_code == (200 if outcome == "success" else 400), response.text + assert_hooks( + records, + surface, + call_id, + sorted([f"function_{outcome}", f"async_{outcome}" if reached_by_the_async_success else f"sync_{outcome}"]), + ) + + +@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST) +def test_logger_the_async_handler_skips_gets_one_sync_hook_per_request_rejected_by_key_auth( + sync_hook_proxy: tuple[Gateway, HookRecords, str], surface: str +) -> None: + proxy, records, _ = sync_hook_proxy + known_calls: Final = {record["litellm_call_id"] for record in records.drained()} + response: Final = send(proxy, surface, "success", str(uuid.uuid4()), key="sk-not-a-virtual-key") + assert response.status_code == 401, response.text + + def new_calls(found: tuple[dict[str, str], ...]) -> set[str]: + return {record["litellm_call_id"] for record in found if record["litellm_call_id"] not in known_calls} + + (rejected_call,) = eventually(lambda: new_calls(records.drained()), bool, seconds=20) + assert_hooks(records, surface, rejected_call, ["function_failure", "sync_failure"]) diff --git a/tests/integration/observability/test_custom_logger_hook_counts_post_call_block.py b/tests/integration/observability/test_custom_logger_hook_counts_post_call_block.py new file mode 100644 index 00000000000..1b708db64d8 --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts_post_call_block.py @@ -0,0 +1,75 @@ +"""A CustomLogger with both hooks gets one failure hook when a post-call guardrail blocks the response. + +The proxy drops the async success logging it parked for the guardrails and logs the block as a failure through both +the async and the sync failure handlers. The logger gets only its async failure hook from them. Speech and +transcription responses never reach post-call guardrails, and image generation logs its success before they run, so +neither parks. + +A legacy sync callback is configured, because that is what makes the sync handler run for an async call.""" + +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway, gateway_from_environment, held +from integration._support.hook_counts import ( + SURFACE_CALL_TYPES, + HookRecords, + probe_hooks, + provider, + send, + write_hook_count_config, +) +from integration._support.process import owned_proxy +from integration._support.wire import Reply, wire_server + +pytestmark = pytest.mark.timeout(240) + +_PROBES: Final = "integration._support.hook_count_probe" +_PARKED_FOR_POST_CALL_GUARDRAILS: Final = sorted( + SURFACE_CALL_TYPES.keys() - {"speech", "transcription", "image_generation"} +) + + +@pytest.fixture(scope="module") +def post_call_block(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, HookRecords]]: + tmp_path: Final = tmp_path_factory.mktemp("hook-count-post-call-block") + guardrail: Final = { + "guardrail_name": "hook-count-post-call-block", + "litellm_params": { + "guardrail": f"{_PROBES}.BlockingPostCallGuardrail", + "mode": "post_call", + "default_on": True, + }, + } + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config( + tmp_path / "hook-count-post-call-block.yaml", + upstream.url, + {"callbacks": [f"{_PROBES}.probe"], "success_callback": [f"{_PROBES}.sync_success_function"]}, + (guardrail,), + ) + overrides: Final = {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"} + with owned_proxy(rig, tmp_path, overrides, config=config, workers=2) as proxy: + yield proxy, HookRecords(sink, function_callbacks=False) + + +@pytest.mark.parametrize("surface", _PARKED_FOR_POST_CALL_GUARDRAILS) +def test_response_a_post_call_guardrail_blocks_gives_the_logger_one_failure_hook( + post_call_block: tuple[Gateway, HookRecords], surface: str +) -> None: + proxy, records = post_call_block + call_id: Final = str(uuid.uuid4()) + response: Final = send(proxy, surface, "success", call_id) + assert response.status_code == 400, response.text + held( + lambda: probe_hooks(records, surface, records.for_call(call_id)), + lambda hooks: hooks == ["async_failure"], + holding=3, + seconds=20, + ) diff --git a/tests/integration/observability/test_custom_logger_hook_counts_vendors.py b/tests/integration/observability/test_custom_logger_hook_counts_vendors.py new file mode 100644 index 00000000000..3c9be9e2a3f --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts_vendors.py @@ -0,0 +1,128 @@ +"""Built-in loggers that override both hooks deliver one record per async request outside the async flag list: +Datadog one log entry and LangSmith one run. OpenTelemetry keeps its own emit-once guard, so it never exports two +spans for a request either way; its speech, OCR and search spans carry no `litellm.call_id`, so only "never two" +is checked for it. + +A legacy sync callback is configured, because that is what makes the sync handler run for an async call.""" + +import gzip +import json +import uuid +import zlib +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import Gateway, gateway_from_environment, held +from integration._support.hook_counts import IN_THE_ASYNC_FLAG_LIST, SURFACE_CALL_TYPES, provider, send +from integration._support.hook_counts import write_hook_count_config +from integration._support.otlp_sink import decode_spans +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server + +pytestmark = pytest.mark.timeout(240) + +_PROBES: Final = "integration._support.hook_count_probe" +_FLUSH_SECONDS: Final = 1 +_OUTSIDE_THE_ASYNC_FLAG_LIST: Final = sorted(SURFACE_CALL_TYPES.keys() - IN_THE_ASYNC_FLAG_LIST) + + +@dataclass(frozen=True) +class _Vendors: + proxy: Gateway + sink: Wire + + +@pytest.fixture(scope="module") +def vendors(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Vendors]: + tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-count-vendors") + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply(status=202)) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config( + tmp_path / "hook-count-vendors.yaml", + upstream.url, + { + "callbacks": ["datadog", "langsmith", "otel"], + "success_callback": [f"{_PROBES}.sync_success_function"], + "failure_callback": [f"{_PROBES}.sync_failure_function"], + }, + ) + overrides: Final = { + "HOOK_COUNT_PROBE_URL": sink.url + "/probe", + "STORE_MODEL_IN_DB": "False", + "DEFAULT_FLUSH_INTERVAL_SECONDS": str(_FLUSH_SECONDS), + "DD_API_KEY": "synthetic-datadog-key", + "DD_SITE": "us5.datadoghq.com", + "DD_BASE_URL": sink.url + "/datadog", + "LANGSMITH_API_KEY": "synthetic-langsmith-key", + "LANGSMITH_PROJECT": "hook-counts", + "LANGSMITH_BASE_URL": sink.url + "/langsmith", + "OTEL_EXPORTER": "otlp_http", + "OTEL_ENDPOINT": sink.url + "/otel/v1/traces", + } + with owned_proxy(rig, tmp_path, overrides, config=config, workers=2) as proxy: + yield _Vendors(proxy, sink) + + +def _unzipped(body: bytes) -> bytes: + if body[:2] == b"\x1f\x8b": + return gzip.decompress(body) + try: + return zlib.decompress(body, zlib.MAX_WBITS | 16) + except zlib.error: + return body + + +def _entries(body: bytes) -> tuple[object, ...]: + decoded: Final = json.loads(_unzipped(body)) + return tuple(decoded) if isinstance(decoded, list) else (decoded,) + + +def _records(request: Request) -> tuple[tuple[str, str], ...]: + """Every record one delivery carries, as (destination, serialized record). A LangSmith batch carries its new + runs under `post`, and an OpenTelemetry export carries every span of the request's trace.""" + if request.target.startswith("/datadog"): + return tuple(("datadog", json.dumps(entry)) for entry in _entries(request.body)) + if request.target.startswith("/langsmith") and request.method == "POST": + runs: Final = tuple( + run + for entry in _entries(request.body) + for run in (entry.get("post", []) if isinstance(entry, dict) and "post" in entry else [entry]) + ) + return tuple(("langsmith", json.dumps(run)) for run in runs) + if request.target.startswith("/otel"): + spans: Final = decode_spans(request.body, request.headers.get("content-type", "")) + return tuple(("otel", json.dumps(span)) for span in spans if "litellm.call_id" in span["attributes"]) + return () + + +@dataclass +class _Deliveries: + sink: Wire + seen: tuple[tuple[str, str], ...] = () + + def per_destination(self, call_id: str) -> dict[str, int]: + self.seen = (*self.seen, *(record for request in self.sink.drain() for record in _records(request))) + destinations: Final = ("datadog", "langsmith", "otel") + return { + destination: sum(1 for name, record in self.seen if name == destination and call_id in record) + for destination in destinations + } + + +@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST) +def test_built_in_loggers_deliver_one_record_per_async_request(vendors: _Vendors, surface: str) -> None: + deliveries: Final = _Deliveries(vendors.sink) + call_id: Final = str(uuid.uuid4()) + response: Final = send(vendors.proxy, surface, "success", call_id) + assert response.status_code == 200, response.text + held( + lambda: deliveries.per_destination(call_id), + lambda counts: counts["datadog"] == 1 and counts["langsmith"] == 1 and counts["otel"] <= 1, + holding=4 * _FLUSH_SECONDS, + seconds=60, + ) diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index 2badf810b6a..d4f9533fcdc 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -1,6 +1,7 @@ import json import re from collections.abc import Callable, Mapping +from itertools import chain from typing import Final from uuid import uuid4 @@ -263,13 +264,17 @@ def _assert_chunk_envelope(chunks: tuple[dict[str, JsonValue], ...], model: str, return next(iter(identities)) +def _content_chunks(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(chunk for chunk in chunks if chunk.get("usage") is None) + + def _finish_reasons(chunks: tuple[dict[str, JsonValue], ...]) -> list[JsonValue]: - reasons: Final = (_only_choice(chunk).get("finish_reason") for chunk in chunks) + reasons: Final = (_only_choice(chunk).get("finish_reason") for chunk in _content_chunks(chunks)) return [reason for reason in reasons if reason is not None] def _deltas(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: - return tuple(object_value(_only_choice(chunk)["delta"]) for chunk in chunks) + return tuple(object_value(_only_choice(chunk)["delta"]) for chunk in _content_chunks(chunks)) def _sdk_base(gateway: Gateway) -> str: @@ -301,12 +306,12 @@ def test_chat_completion_on_a_mantle_claude_id_posts_native_messages_and_spends_ @pytest.mark.parametrize( ("stream_options", "usage_chunks"), [ - pytest.param({"stream_options": {"include_usage": True}}, [_USAGE], id="include_usage"), + pytest.param({"stream_options": {"include_usage": True}}, [(True, *_USAGE)], id="include_usage"), pytest.param({}, [], id="no_usage"), ], ) def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_stream( - gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[int, int, int]] + gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[JsonValue, ...]] ) -> None: marker: Final = uuid4().hex expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) @@ -322,9 +327,9 @@ def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_ identity: Final = _assert_chunk_envelope(chunks, model, response.text) assert "".join(str(delta.get("content") or "") for delta in _deltas(chunks)) == _answer(marker), response.text assert _finish_reasons(chunks) == ["stop"], response.text - assert [_token_usage(chunk) for chunk in chunks if chunk.get("usage") is not None] == usage_chunks, ( - response.text - ) + assert [ + (chunk["choices"] == [], *_token_usage(chunk)) for chunk in chunks if chunk.get("usage") is not None + ] == usage_chunks, response.text _assert_success_row( _spend_row(_BY_REQUEST_ID, identity), model, _HAIKU, "acompletion", _mantle_cost(gateway, _HAIKU) ) @@ -620,15 +625,14 @@ def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: ) identities: Final = {chunk.id for chunk in chunks} assert len(identities) == 1 and {chunk.model for chunk in chunks} == {model}, chunks - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _answer(marker), chunks - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"], ( - chunks - ) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _answer(marker), chunks + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"], chunks assert [ - (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + (len(chunk.choices), chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None - ] == [_USAGE], chunks + ] == [(0, *_USAGE)], chunks _assert_success_row( _spend_row(_BY_REQUEST_ID, next(iter(identities))), model, @@ -688,15 +692,14 @@ async def test_async_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gatew chunks: Final = tuple([chunk async for chunk in stream]) identities: Final = {chunk.id for chunk in chunks} assert len(identities) == 1 and {chunk.model for chunk in chunks} == {model}, chunks - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _answer(marker), chunks - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"], ( - chunks - ) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _answer(marker), chunks + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"], chunks assert [ - (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + (len(chunk.choices), chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None - ] == [_USAGE], chunks + ] == [(0, *_USAGE)], chunks _assert_success_row( _spend_row(_BY_REQUEST_ID, next(iter(identities))), model, diff --git a/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py new file mode 100644 index 00000000000..f995a7ac83f --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py @@ -0,0 +1,531 @@ +import base64 +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import boto3 +import pytest +import yaml +from botocore.config import Config +from botocore.exceptions import ClientError +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180) + +_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_BLOCKED_WORD: Final = "zzzforbidden" +_BLOCKED_TEXT: Final = "Blocked by guardrail" +_VERSIONLESS_GUARDRAIL_MESSAGE: Final = "guardrailIdentifier and guardrailVersion must be provided together." +_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream") +_DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "deployment-guardrail", + "guardrailVersion": "3", + "trace": "enabled", + "streamProcessingMode": "sync", +} +_CALLER_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "caller-guardrail", + "guardrailVersion": "9", + "trace": "disabled", +} +_MANTLE_KEY: Final = "scripted-mantle-key" +_CALLER_INVOKE_HEADERS: Final[Mapping[str, str]] = { + "x-pass-x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-pass-x-amzn-bedrock-guardrailversion": "9", + "x-pass-x-amzn-bedrock-trace": "DISABLED", +} + + +@dataclass(frozen=True, slots=True) +class Applied: + """The guardrail a request asked Bedrock for, read from where Bedrock reads it for that action.""" + + identifier: str + version: str | None + trace: str | None + stream_processing_mode: str | None + + +@dataclass(frozen=True, slots=True) +class Outcome: + """What the AWS SDK caller reads back.""" + + text: str + intervened: bool + traced_guardrail: str | None + + +def _action(request: Request) -> str: + return request.target.partition("?")[0].rsplit("/", 1)[-1] + + +def _applied(request: Request) -> Applied | None: + if _action(request).startswith("converse"): + config: Final = json.loads(request.body).get("guardrailConfig") + if config is None: + return None + return Applied( + config["guardrailIdentifier"], + config.get("guardrailVersion"), + config.get("trace"), + config.get("streamProcessingMode"), + ) + identifier: Final = request.headers.get("x-amzn-bedrock-guardrailidentifier") + if identifier is None: + return None + return Applied( + identifier, + request.headers.get("x-amzn-bedrock-guardrailversion"), + request.headers.get("x-amzn-bedrock-trace"), + None, + ) + + +def _identifier(request: Request) -> str | None: + applied: Final = _applied(request) + return None if applied is None else applied.identifier + + +def _assessment(applied: Applied) -> Mapping[str, JsonValue]: + return { + applied.identifier: { + "wordPolicy": {"customWords": [{"match": _BLOCKED_WORD, "action": "BLOCKED", "detected": True}]}, + "appliedGuardrailDetails": {"guardrailId": applied.identifier, "guardrailVersion": applied.version}, + } + } + + +def _bedrock_with_guardrails(request: Request) -> Reply: + """Bedrock's runtime as the caller sees it: the guardrail the request names blocks a prompt holding the blocked + word, answering with its message and, when the request enabled tracing, a trace naming that guardrail. A guardrail + named without a version is refused with a 400 ValidationException, as real Bedrock in us-east-1 refused it on + every one of these actions on 2026-10-19.""" + applied: Final = _applied(request) + if applied is not None and applied.version is None: + return Reply( + status=400, + body=json.dumps({"message": _VERSIONLESS_GUARDRAIL_MESSAGE}).encode(), + headers={"x-amzn-errortype": "ValidationException"}, + ) + blocked: Final = applied is not None and _BLOCKED_WORD.encode() in request.body + traced: Final = blocked and applied is not None and (applied.trace or "").lower() == "enabled" + text: Final = _BLOCKED_TEXT if blocked else "ok" + stop_reason: Final = "guardrail_intervened" if blocked else "end_turn" + usage: Final = {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6} + converse_trace: Final = {"trace": {"guardrail": {"inputAssessment": _assessment(applied)}}} if traced else {} + action: Final = _action(request) + if action == "converse": + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": stop_reason, + "usage": usage, + "metrics": {"latencyMs": 3}, + **converse_trace, + } + ).encode() + ) + if action == "converse-stream": + return Reply( + chunks=( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": stop_reason}, "sc", "u"), + _aws_event_frame( + "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u" + ), + ), + content_type=_EVENT_STREAM, + ) + message: Final = { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + **({"amazon-bedrock-guardrailAction": "INTERVENED"} if blocked else {}), + **({"amazon-bedrock-trace": {"guardrail": {"input": _assessment(applied)}}} if traced else {}), + } + if action == "invoke": + return Reply(body=json.dumps(message).encode()) + chunk: Final = {"bytes": base64.b64encode(json.dumps(message).encode()).decode()} + return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM) + + +def _bedrock_outage(request: Request) -> Reply: + return Reply( + status=500, + body=json.dumps({"message": "scripted outage"}).encode(), + headers={"x-amzn-errortype": "InternalServerException"}, + ) + + +def _traced_guardrail(assessment: object) -> str | None: + return None if not isinstance(assessment, dict) else next(iter(assessment)) + + +def _send(runtime: object, action: str, model: str, prompt: str, caller_guardrail: bool) -> Outcome: + messages: Final = [{"role": "user", "content": [{"text": prompt}]}] + invoke_body: Final = json.dumps( + {"anthropic_version": "bedrock-2023-05-31", "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]} + ) + invoke_guardrail: Final = ( + {"guardrailIdentifier": "caller-guardrail", "guardrailVersion": "9", "trace": "DISABLED"} + if caller_guardrail + else {} + ) + if action == "converse": + response: Final = getattr(runtime, "converse")( + modelId=model, + messages=messages, + **({"guardrailConfig": dict(_CALLER_GUARDRAIL)} if caller_guardrail else {}), + ) + return Outcome( + text=response["output"]["message"]["content"][0]["text"], + intervened=response["stopReason"] == "guardrail_intervened", + traced_guardrail=_traced_guardrail(response.get("trace", {}).get("guardrail", {}).get("inputAssessment")), + ) + if action == "converse-stream": + events: Final = list( + getattr(runtime, "converse_stream")( + modelId=model, + messages=messages, + **( + {"guardrailConfig": {**_CALLER_GUARDRAIL, "streamProcessingMode": "async"}} + if caller_guardrail + else {} + ), + )["stream"] + ) + (stop,) = tuple(event["messageStop"]["stopReason"] for event in events if "messageStop" in event) + (metadata,) = tuple(event["metadata"] for event in events if "metadata" in event) + return Outcome( + text="".join( + event["contentBlockDelta"]["delta"]["text"] for event in events if "contentBlockDelta" in event + ), + intervened=stop == "guardrail_intervened", + traced_guardrail=_traced_guardrail(metadata.get("trace", {}).get("guardrail", {}).get("inputAssessment")), + ) + if action == "invoke": + message: Final = json.loads( + getattr(runtime, "invoke_model")(modelId=model, body=invoke_body, **invoke_guardrail)["body"].read() + ) + else: + (message,) = tuple( + json.loads(event["chunk"]["bytes"]) + for event in getattr(runtime, "invoke_model_with_response_stream")( + modelId=model, body=invoke_body, **invoke_guardrail + )["body"] + ) + return Outcome( + text=message["content"][0]["text"], + intervened=message.get("amazon-bedrock-guardrailAction") == "INTERVENED", + traced_guardrail=_traced_guardrail(message.get("amazon-bedrock-trace", {}).get("guardrail", {}).get("input")), + ) + + +def _signed_header_names(request: Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _deployment(wire: Wire, guardrail_config: Mapping[str, str] | None = None) -> Mapping[str, JsonValue]: + return { + "model": f"bedrock/{_MODEL_ID}", + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + **({} if guardrail_config is None else {"guardrailConfig": dict(guardrail_config)}), + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + wire: Wire + guarded_model: str + unguarded_model: str + balanced_model: str + balanced_wires: Mapping[str, Wire] + failing_model: str + failing_wire: Wire + fallback_wire: Wire + mantle_model: str + versionless_model: str + key: str + + def runtime( # noqa: ANN201 # boto3 clients are untyped + self, monkeypatch: pytest.MonkeyPatch, extra_headers: Mapping[str, str] = MappingProxyType({}) + ): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", self.key) + client: Final = boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}), + ) + client.meta.events.register( + "before-send.bedrock-runtime.*", + lambda request, **_: request.headers.update(extra_headers), + ) + return client + + def received(self, wire: Wire, marker: str) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if marker.encode() in request.body) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + guarded_model: Final = f"bedrock-guarded-{suffix}" + unguarded_model: Final = f"bedrock-unguarded-{suffix}" + balanced_model: Final = f"bedrock-balanced-{suffix}" + failing_model: Final = f"bedrock-failing-{suffix}" + fallback_model: Final = f"bedrock-fallback-{suffix}" + mantle_model: Final = f"bedrock-mantle-{suffix}" + versionless_model: Final = f"bedrock-versionless-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-guardrail-config") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with ( + wire_server(_bedrock_with_guardrails) as wire, + wire_server(_bedrock_with_guardrails) as wire_a, + wire_server(_bedrock_with_guardrails) as wire_b, + wire_server(_bedrock_outage) as failing_wire, + wire_server(_bedrock_with_guardrails) as fallback_wire, + ): + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + { + "model_name": guarded_model, + "litellm_params": _deployment(wire, _DEPLOYMENT_GUARDRAIL), + }, + {"model_name": unguarded_model, "litellm_params": _deployment(wire)}, + { + "model_name": balanced_model, + "litellm_params": _deployment( + wire_a, {"guardrailIdentifier": "guardrail-a", "guardrailVersion": "1"} + ), + }, + { + "model_name": balanced_model, + "litellm_params": _deployment( + wire_b, {"guardrailIdentifier": "guardrail-b", "guardrailVersion": "2"} + ), + }, + { + "model_name": failing_model, + "litellm_params": _deployment( + failing_wire, + {"guardrailIdentifier": "failing-guardrail", "guardrailVersion": "5", "trace": "enabled"}, + ), + }, + { + "model_name": fallback_model, + "litellm_params": _deployment( + fallback_wire, + {"guardrailIdentifier": "fallback-guardrail", "guardrailVersion": "6", "trace": "enabled"}, + ), + }, + { + "model_name": versionless_model, + "litellm_params": _deployment( + wire, {"guardrailIdentifier": "versionless-guardrail", "trace": "enabled"} + ), + }, + { + "model_name": mantle_model, + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-sonnet-5", + "api_base": wire.url, + "api_key": _MANTLE_KEY, + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + "guardrailConfig": { + "guardrailIdentifier": "mantle-guardrail", + "guardrailVersion": "4", + "trace": "enabled", + }, + }, + }, + ], + "router_settings": { + **base_config.get("router_settings", {}), + "fallbacks": [{failing_model: [fallback_model]}], + }, + } + config_path: Final = directory / "bedrock-guardrail-config.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(session_gateway, directory, {}, config=config_path, workers=2) as proxy, + proxy.scenario() as scenario, + ): + yield Rig( + proxy=proxy, + wire=wire, + guarded_model=guarded_model, + unguarded_model=unguarded_model, + balanced_model=balanced_model, + balanced_wires={"guardrail-a": wire_a, "guardrail-b": wire_b}, + failing_model=failing_model, + failing_wire=failing_wire, + fallback_wire=fallback_wire, + mantle_model=mantle_model, + versionless_model=versionless_model, + key=scenario.key(key_alias=f"guardrail-key-{suffix}"), + ) + + +def _expected_on_the_wire(action: str, request: Request) -> Applied: + if action == "converse": + return Applied("deployment-guardrail", "3", "enabled", None) + if action == "converse-stream": + return Applied("deployment-guardrail", "3", "enabled", "sync") + assert {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} <= ( + _signed_header_names(request) + ), request.headers["authorization"] + return Applied("deployment-guardrail", "3", "ENABLED", None) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_passthrough_on_a_guardrailed_deployment_blocks_with_the_deployment_guardrail( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + blocked: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + allowed: Final = _send(runtime, action, rig.guarded_model, f"{marker} say hello", caller_guardrail=False) + + assert blocked == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), blocked + assert allowed == Outcome("ok", intervened=False, traced_guardrail=None), allowed + received: Final = rig.received(rig.wire, marker) + assert [_action(request) for request in received] == [action, action], received + assert [_applied(request) for request in received] == [ + _expected_on_the_wire(action, request) for request in received + ], [(request.headers, request.body) for request in received] + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_caller_guardrail_is_replaced_by_the_deployment_guardrail( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch, _CALLER_INVOKE_HEADERS) + + outcome: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == _expected_on_the_wire(action, request), (request.headers, request.body) + assert "caller-guardrail" not in json.dumps({**request.headers, "body": request.body.decode()}), request + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_mantle_deployment_passthrough_sends_its_guardrail_with_bearer_auth( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcome: Final = _send(runtime, action, rig.mantle_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="mantle-guardrail"), outcome + (request,) = rig.received(rig.wire, marker) + assert request.headers["authorization"] == f"Bearer {_MANTLE_KEY}", request.headers + assert _applied(request) == Applied( + "mantle-guardrail", "4", "ENABLED" if action.startswith("invoke") else "enabled", None + ), (request.headers, request.body) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_deployment_guardrail_without_a_version_returns_the_bedrock_rejection( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + with pytest.raises(ClientError) as raised: + _send(runtime, action, rig.versionless_model, f"{marker} say hello", caller_guardrail=False) + + error: Final = raised.value.response.get("Error", {}) + assert (error.get("Code"), error.get("Message")) == ("ValidationException", _VERSIONLESS_GUARDRAIL_MESSAGE), error + received: Final = rig.received(rig.wire, marker) + assert {_applied(request) for request in received} == { + Applied("versionless-guardrail", None, "ENABLED" if action.startswith("invoke") else "enabled", None) + }, [(request.headers, request.body) for request in received] + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_deployment_without_a_guardrail_leaves_the_caller_guardrail_as_sent( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + if action.startswith("invoke"): + pytest.skip("BUG: the router-model invoke pass-through drops the caller's X-Amzn-Bedrock-Guardrail* headers") + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcome: Final = _send(runtime, action, rig.unguarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail=None), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == Applied( + "caller-guardrail", "9", "disabled", "async" if action == "converse-stream" else None + ), request.body + + +def test_each_load_balanced_request_carries_the_guardrail_of_the_deployment_it_reached( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcomes: Final = tuple( + _send(runtime, action, rig.balanced_model, f"{marker}-{attempt} say {_BLOCKED_WORD}", caller_guardrail=False) + for attempt, action in enumerate(_ACTIONS * 6) + ) + + assert all(outcome.intervened for outcome in outcomes), outcomes + reached: Final = {name: rig.received(wire, marker) for name, wire in rig.balanced_wires.items()} + assert sum(len(requests) for requests in reached.values()) == len(outcomes), reached + assert all(reached.values()), {name: len(requests) for name, requests in reached.items()} + assert {name: {_identifier(request) for request in requests} for name, requests in reached.items()} == { + "guardrail-a": {"guardrail-a"}, + "guardrail-b": {"guardrail-b"}, + }, reached + + +def test_a_fallback_sends_the_guardrail_of_the_deployment_it_fell_back_to( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcomes: Final = tuple( + _send(runtime, action, rig.failing_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + for action in _ACTIONS + ) + + assert outcomes == (Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="fallback-guardrail"),) * 4, outcomes + failed: Final = rig.received(rig.failing_wire, marker) + served: Final = rig.received(rig.fallback_wire, marker) + assert [_action(request) for request in served] == list(_ACTIONS), served + assert {_identifier(request) for request in failed} == {"failing-guardrail"}, failed + assert {_identifier(request) for request in served} == {"fallback-guardrail"}, served diff --git a/tests/integration/routing/test_router_strategies_count_async_requests_once.py b/tests/integration/routing/test_router_strategies_count_async_requests_once.py new file mode 100644 index 00000000000..f7403ad0437 --- /dev/null +++ b/tests/integration/routing/test_router_strategies_count_async_requests_once.py @@ -0,0 +1,214 @@ +"""Router strategies that keep both a sync and an async hook count a finished async request once: least-busy +keeps counting a request still in flight while others finish around it, and usage-based routing v2 adds a +request's tokens to its deployment once. + +A legacy sync callback is configured, because that is what makes the sync handler run for an async call.""" + +import base64 +import json +import os +import threading +import uuid +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +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, held +from integration._support.hook_counts import PNG_BYTES, WAV_BYTES, HookRecords +from integration._support.process import OwnedState, owned_proxy +from integration._support.wire import Reply, Request, wire_server +from redis import Redis + +pytestmark = pytest.mark.timeout(240) + +_MODEL: Final = "router-strategy-model" +_PROBES: Final = "integration._support.hook_count_probe" +_REJECTION: Final = b'{"error":{"message":"scripted provider rejection","type":"invalid_request_error"}}' +_IMAGE_TOKENS: Final = 12 +_IMAGE_REPLY: Final = json.dumps( + { + "created": 1700000000, + "data": [{"b64_json": base64.b64encode(PNG_BYTES).decode()}], + "usage": { + "input_tokens": 5, + "output_tokens": 7, + "total_tokens": _IMAGE_TOKENS, + "input_tokens_details": {"text_tokens": 5, "image_tokens": 0}, + }, + } +).encode() + + +@dataclass(frozen=True) +class _Rig: + proxy: Gateway + records: HookRecords + release: threading.Event + deployment_id: str + + +def _provider(release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/v1/images/generations": + return Reply(body=_IMAGE_REPLY) + spoken: Final = str(json.loads(request.body or b"{}").get("input", "")) + if spoken.startswith("fail"): + return Reply(status=400, body=_REJECTION) + if spoken.startswith("hold"): + return Reply(chunks=(WAV_BYTES,), content_type="audio/wav", gate_before_first=release) + return Reply(body=WAV_BYTES, content_type="audio/wav") + + return respond + + +def _write_config(path: Path, provider_url: str, strategy: str, model: str, deployment_id: str) -> Path: + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _MODEL, + "litellm_params": { + "model": model, + "api_base": provider_url + "/v1", + "api_key": "synthetic-provider-key", + }, + "model_info": {"id": deployment_id}, + } + ], + "litellm_settings": { + "success_callback": [f"{_PROBES}.sync_success_function"], + "failure_callback": [f"{_PROBES}.sync_failure_function"], + }, + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_model_info_refresh": True, + }, + "router_settings": { + "routing_strategy": strategy, + "routing_strategy_args": {"ttl": 600}, + "num_retries": 0, + "disable_cooldowns": True, + "redis_host": "os.environ/REDIS_HOST", + "redis_port": "os.environ/REDIS_PORT", + }, + } + ) + ) + return path + + +@contextmanager +def _router_proxy(tmp_path: Path, strategy: str, model: str) -> Iterator[_Rig]: + release: Final = threading.Event() + deployment_id: Final = f"{strategy}-{uuid.uuid4().hex}" + with ( + wire_server(_provider(release)) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = _write_config(tmp_path / f"{strategy}.yaml", upstream.url, strategy, model, deployment_id) + try: + with owned_proxy( + rig, + tmp_path, + {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, + config=config, + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as proxy: + yield _Rig(proxy, HookRecords(sink, function_callbacks=True), release, deployment_id) + finally: + release.set() + + +@pytest.fixture(scope="module") +def least_busy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + with _router_proxy(tmp_path_factory.mktemp("least-busy"), "least-busy", "openai/gpt-4o-mini-tts") as rig: + yield rig + + +@pytest.fixture(scope="module") +def usage_based(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + with _router_proxy(tmp_path_factory.mktemp("usage-based"), "usage-based-routing-v2", "openai/gpt-image-1") as rig: + yield rig + + +def _redis() -> Redis: + return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + + +def _in_flight(rig: _Rig) -> int: + with _redis() as cache: + stored: Final = cache.get(f"{_MODEL}_request_count:{rig.deployment_id}") + return 0 if stored is None else int(stored) + + +def _tokens_counted(rig: _Rig) -> int: + with _redis() as cache: + keys: Final = tuple(cache.scan_iter(match=f"{rig.deployment_id}:*:tpm:*")) + return sum(int(cache.get(key) or 0) for key in keys) + + +def _post(rig: _Rig, path: str, body: dict[str, object], call_id: str) -> httpx.Response: + return rig.proxy.client.post( + path, + json={"model": _MODEL, **body}, + headers={"Authorization": f"Bearer {rig.proxy.key}", "x-litellm-call-id": call_id}, + timeout=120, + ) + + +def _speak(rig: _Rig, text: str, call_id: str) -> httpx.Response: + return _post(rig, "/v1/audio/speech", {"input": text, "voice": "alloy", "response_format": "wav"}, call_id) + + +def _logged(rig: _Rig, call_id: str) -> None: + """The request's success or failure logging, the legacy sync callback included, has run.""" + eventually(lambda: rig.records.for_call(call_id), bool, seconds=20) + + +def _finish_a_speech_request(rig: _Rig, outcome: str) -> None: + call_id: Final = str(uuid.uuid4()) + finished: Final = _speak(rig, f"{'ok' if outcome == 'success' else 'fail'} {call_id}", call_id) + assert finished.status_code == (200 if outcome == "success" else 400), finished.text + _logged(rig, call_id) + held(lambda: _in_flight(rig), lambda count: count == 1, holding=3, seconds=20) + + +def _generate_an_image(rig: _Rig) -> None: + call_id: Final = str(uuid.uuid4()) + response: Final = _post(rig, "/v1/images/generations", {"prompt": f"tokens {call_id}", "n": 1}, call_id) + assert response.status_code == 200, response.text + _logged(rig, call_id) + + +@pytest.mark.parametrize("outcome", ["success", "failure"]) +def test_least_busy_keeps_counting_a_request_in_flight_while_async_speech_requests_finish( + least_busy: _Rig, outcome: str +) -> None: + held(lambda: _in_flight(least_busy), lambda count: count == 0, holding=1, seconds=30) + least_busy.release.clear() + with ThreadPoolExecutor(max_workers=1) as pool: + in_flight: Final = pool.submit(_speak, least_busy, f"hold {uuid.uuid4().hex}", str(uuid.uuid4())) + eventually(lambda: _in_flight(least_busy), lambda count: count == 1, seconds=30) + try: + for _ in range(3): + _finish_a_speech_request(least_busy, outcome) + finally: + least_busy.release.set() + assert in_flight.result(timeout=60).status_code == 200 + held(lambda: _in_flight(least_busy), lambda count: count == 0, holding=3, seconds=30) + + +def test_usage_based_routing_counts_an_async_image_request_s_tokens_once(usage_based: _Rig) -> None: + for _ in range(3): + _generate_an_image(usage_based) + held(lambda: _tokens_counted(usage_based), lambda tokens: tokens == 3 * _IMAGE_TOKENS, holding=3, seconds=30) diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index 6d481a98fd0..1b7fc975217 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -1,8 +1,10 @@ import asyncio +import base64 import json import uuid from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor +from itertools import chain from typing import Final import httpx @@ -10,8 +12,12 @@ import openai import pytest from integration._support.client import Gateway, Scenario, eventually, held, list_value, object_value, string_value from integration._support.database import read_rows, write_rows +from integration._support.otlp_sink import decode_spans from integration._support.upstream import delete_scenario, register_scenario +from integration._support.wire import Request, wire_server from integration.caching.response_cache_case import entry_key, redis_store +from integration.security._callback_traffic import langfuse_sink +from integration.security._sinks import Recorder from integration.spend._background_responses import ( ATTRIBUTION_COLUMNS, CONCURRENT_READS, @@ -60,7 +66,7 @@ from integration.spend._background_responses import ( ) from openai._legacy_response import LegacyAPIResponse from openai.types.responses import Response -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN @@ -823,3 +829,78 @@ def test_a_client_cannot_make_a_background_create_free_by_forging_the_managed_ma scenario.cleanups.callback(forget_row, response_id) assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == free_reads("completed", 2) assert_charged_exactly(lambda: key_spend(key), PRICE) + + +LANGFUSE_PUBLIC_KEY: Final = "pk-lf-background" +LANGFUSE_SECRET_KEY: Final = "sk-lf-background" +LANGFUSE_TRACES_PATH: Final = "/api/public/otel/v1/traces" +LANGFUSE_COST: Final = "langfuse.observation.cost_details" +_LANGFUSE_COST_DETAILS: Final = TypeAdapter(dict[str, float]) + + +def _langfuse_callback(host: str) -> dict[str, JsonValue]: + return { + "callback_name": "langfuse", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": LANGFUSE_PUBLIC_KEY, + "langfuse_secret_key": LANGFUSE_SECRET_KEY, + "langfuse_host": host, + }, + } + + +def _key_with_its_own_logging(scenario: Scenario, langfuse_url: str) -> str: + return scenario.key(metadata={"logging": [_langfuse_callback(langfuse_url)]}) + + +def _key_in_a_team_with_logging(scenario: Scenario, langfuse_url: str) -> str: + team_id: Final = scenario.team() + added: Final = scenario.gateway.request("POST", f"/team/{team_id}/callback", _langfuse_callback(langfuse_url)) + assert added.status_code == 200, added.text + return scenario.key(team_id=team_id) + + +LOGGING_ON: Final[Mapping[str, Callable[[Scenario, str], str]]] = { + "the_key": _key_with_its_own_logging, + "the_team": _key_in_a_team_with_logging, +} + + +def _trace_exports(langfuse: Recorder) -> tuple[Request, ...]: + return tuple(request for request in langfuse.requests() if request.target.startswith(LANGFUSE_TRACES_PATH)) + + +def _exported_costs(langfuse: Recorder) -> tuple[float, ...]: + spans: Final = tuple( + chain.from_iterable( + decode_spans(export.body, export.headers.get("content-type", "")) for export in _trace_exports(langfuse) + ) + ) + return tuple( + _LANGFUSE_COST_DETAILS.validate_json(string_value(span["attributes"][LANGFUSE_COST]))["total"] + for span in spans + if LANGFUSE_COST in span["attributes"] + ) + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("logging_on", LOGGING_ON.keys()) +def test_the_billed_cost_reaches_the_langfuse_project_set_on_the_creating_key_or_team( + gateway: Gateway, logging_on: str +) -> None: + with wire_server(langfuse_sink) as wire, gateway.scenario() as scenario: + langfuse: Final = Recorder(wire) + key: Final = LOGGING_ON[logging_on](scenario, wire.url) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, key, upstream.model) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + exported: Final = eventually( + lambda: _exported_costs(langfuse), + lambda costs: pytest.approx(PRICE) in costs, + seconds=SPEND_LANDING_SECONDS, + ) + assert [cost for cost in exported if cost > 0] == [pytest.approx(PRICE)] + credentials: Final = base64.b64encode(f"{LANGFUSE_PUBLIC_KEY}:{LANGFUSE_SECRET_KEY}".encode()).decode() + assert {export.headers.get("authorization") for export in _trace_exports(langfuse)} == {f"Basic {credentials}"} diff --git a/tests/integration/spend/test_spend_log_tool_payload_content.py b/tests/integration/spend/test_spend_log_tool_payload_content.py index 9055f6f06c4..226accdbf54 100644 --- a/tests/integration/spend/test_spend_log_tool_payload_content.py +++ b/tests/integration/spend/test_spend_log_tool_payload_content.py @@ -15,7 +15,7 @@ import openai import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -723,6 +723,33 @@ def test_chat_sdk_stream_include_usage_masks_request_fields(gateway: Gateway, tm assert len(post_requests) == 1 +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_chat_stream_usage_the_proxy_adds_keeps_its_marker_out_of_the_stored_request( + gateway: Gateway, tmp_path: Path +) -> None: + response_id: Final = f"chatcmpl-proxy-usage-{uuid4()}" + with ( + wire_server( + _answering_model_listing(lambda request: _chat_stream(response_id, "sort_key", include_usage=True)) + ) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, + isolated.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") + response: Final = isolated.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "stream marker"}], "stream": True}, + ) + assert response.status_code == 200, response.text + assert {string_value(event["id"]) for event in _sse_events(response.text)} == {response_id}, response.text + stored_request: Final = object_value(_stored_row(response_id)["proxy_server_request"]) + assert stored_request["messages"] == [{"role": "user", "content": "stream marker"}], stored_request + assert "_litellm_stream_usage_injected" not in stored_request, stored_request + observed: Final = _provider_calls(wire.drain()) + assert [_json_object(request.body)["stream_options"] for request in observed] == [{"include_usage": True}] + + def test_chat_history_keeps_string_tool_arguments_and_tool_content(gateway: Gateway, tmp_path: Path) -> None: arguments: Final = '{"sort_key":"created_at"}' tool_content: Final = "tool-result-created_at" diff --git a/tests/integration/streaming/test_stream_contracts.py b/tests/integration/streaming/test_stream_contracts.py index d7537a78db8..1d2a8b25792 100644 --- a/tests/integration/streaming/test_stream_contracts.py +++ b/tests/integration/streaming/test_stream_contracts.py @@ -3,6 +3,7 @@ import json import threading import uuid from collections.abc import Callable +from itertools import chain from pathlib import Path from queue import SimpleQueue from typing import Final @@ -13,9 +14,10 @@ from hypothesis import Phase, example, given, settings from hypothesis import strategies as st from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server -from openai import OpenAI +from openai import OpenAI, omit +from openai.types.chat import ChatCompletionChunk from pydantic import JsonValue, TypeAdapter pytestmark: Final = pytest.mark.shares_worker_state @@ -197,6 +199,7 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat assert {chunk.id for chunk in chunks} == {identity} usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) assert len(usages) == (1 if include else 0) + assert [len(chunk.choices) for chunk in chunks if chunk.usage is not None] == ([0] if include else []) if include: assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4 requests: Final = wire.drain() @@ -214,6 +217,119 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat assert float(rows[0]["spend"]) == pytest.approx(0.019) +_STREAM_USAGE_INJECTED_FLAG: Final = "_litellm_stream_usage_injected" + + +def _registered_model_names(gateway: Gateway) -> tuple[JsonValue, ...]: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list), entries + return tuple(object_value(entry).get("model_name") for entry in entries) + + +def _registered(gateway: Gateway, model: str) -> None: + eventually(lambda: _registered_model_names(gateway), lambda model_names: model in model_names, seconds=60) + + +def _sdk_stream( + gateway: Gateway, + route_prefix: str, + model: str, + identity: str, + *, + include_usage: bool | None, + extra_body: dict[str, JsonValue], +) -> tuple[ChatCompletionChunk, ...]: + with OpenAI( + api_key=gateway.key, + base_url=str(gateway.client.base_url).rstrip("/") + route_prefix, + timeout=30, + max_retries=0, + ) as client: + stream: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": identity}], + stream=True, + stream_options=omit if include_usage is None else {"include_usage": include_usage}, + extra_body=extra_body, + ) + with stream: + return tuple(stream) + + +def _streamed_text(chunks: tuple[ChatCompletionChunk, ...]) -> str: + choices: Final = chain.from_iterable(chunk.choices for chunk in chunks) + return "".join(choice.delta.content or "" for choice in choices) + + +def _usage_chunk_shapes(chunks: tuple[ChatCompletionChunk, ...]) -> list[tuple[int, int]]: + return [(len(chunk.choices), chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None] + + +@pytest.mark.parametrize( + ("route_prefix", "sends", "extra_body"), + [pytest.param("/v1", 2, {}, id="chat_miss_then_cache_hit"), pytest.param("/queue", 1, {"priority": 0}, id="queue")], +) +def test_caller_requested_stream_usage_has_no_choices_even_with_a_client_sent_injection_flag( + gateway: Gateway, route_prefix: str, sends: int, extra_body: dict[str, JsonValue] +) -> None: + identity: Final = "stream-usage-flag-" + uuid.uuid4().hex + with ( + wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=wire.url + "/v1") + _registered(gateway, model) + streams: Final = tuple( + _sdk_stream( + gateway, + route_prefix, + model, + identity, + include_usage=True, + extra_body={**extra_body, _STREAM_USAGE_INJECTED_FLAG: True}, + ) + for _ in range(sends) + ) + assert [_streamed_text(chunks) for chunks in streams] == ["Hello 雪 café"] * sends + assert [_usage_chunk_shapes(chunks) for chunks in streams] == [[(0, 15)]] * sends + requests: Final = wire.drain() + assert len(requests) == 1 + assert _STREAM_USAGE_INJECTED_FLAG not in json.loads(requests[0].body) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_stream_usage_the_deployment_adds_keeps_one_empty_choice(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "deployment-stream-usage-" + uuid.uuid4().hex + with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire: + shared_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **shared_config, + "model_list": [ + { + "model_name": "deployment-usage", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "synthetic-deployment-usage-key", + "api_base": wire.url + "/v1", + "stream_options": {"include_usage": True}, + }, + } + ], + "general_settings": {**shared_config["general_settings"], "always_include_stream_usage": False}, + } + path: Final = tmp_path / "deployment-stream-usage.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(gateway, tmp_path, {}, config=path) as candidate: + chunks: Final = _sdk_stream( + candidate, "/v1", "deployment-usage", identity, include_usage=None, extra_body={} + ) + assert _streamed_text(chunks) == "Hello 雪 café" + assert _usage_chunk_shapes(chunks) == [(1, 15)] + requests: Final = wire.drain() + assert len(requests) == 1 + assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True + + @pytest.mark.covers("other.streaming.messages_bridge.empty_choices_usage_chunk_completes_stream") def test_messages_stream_completes_through_trailing_empty_choices_usage_chunk(gateway: Gateway) -> None: identity: Final = "messages-empty-choices-" + uuid.uuid4().hex diff --git a/tests/test_litellm_rust/AGENTS.md b/tests/test_litellm_rust/AGENTS.md index ea661bc36f1..f688e2253e8 100644 --- a/tests/test_litellm_rust/AGENTS.md +++ b/tests/test_litellm_rust/AGENTS.md @@ -1,27 +1,62 @@ -# Invariants +# Scope -- This directory tests the Rust path and nothing else. Treat the Python implementation of every route as already removed. Never call it on purpose, never run one request through both engines, and never compare Python and Rust outputs, headers, cache keys or logs -- A test that would still pass with Python serving the call does not belong here. Every test asserts the call was served natively with `support.routing.assert_served_natively` -- Select the Rust path with `support.routing.native_route(...)`, which prepends `RUST_REQUIRED` rules to `catalog.RULES` and restores them on exit. Each route folder's `conftest.py` applies it to every test through an autouse fixture. A test outside a route folder wraps its whole body in one `native_route(...)` block, and a child interpreter script assigns `catalog.RULES` once at its top. Never select a backend by reading or setting `LITELLM_RUST`, calling `litellm.rust(...)` or rebinding `catalog.RULES` any other way. The rollout in `litellm/rust_bridge/catalog.py` is a product decision, and a test here must pass no matter what it says, including for a route marked `PYTHON_ONLY` -- `catalog.RULES` is static at runtime and `native_route(...)` swaps a module global, so open it once per test and never per call. Two overlapping blocks, such as one per task inside an `asyncio.gather`, restore out of order and leave the wrong rules behind for later tests -- A test whose subject is the rollout switch itself (env, `litellm.rust(...)`, catalog rules) is a unit test of `litellm/rust_bridge/configuration.py` or `catalog.py` and lives in `tests/unit/rust_bridge/` -- A test lives here only when it cannot be written in Rust. That means it needs the Python public API, Python-owned state such as `litellm.cache` and callbacks, or the process boundary (fork, threads, event loops). Anything checkable from the Rust side goes in the crate's own tests +- This directory tests the Rust path end to end through the Python public API, plus the dispatch tests below, and nothing else +- A test belongs here only when it cannot be written in Rust, meaning it needs + - the Python public API + - Python-owned state such as `litellm.cache` and callbacks + - the process boundary (fork, threads, event loops) +- Anything checkable from the Rust side goes in the crate's own tests - The wire is the local recording server. No test here reaches a real provider -- Never parametrize one test across routes. Each route folder owns its tests, its `server` fixture, its request data and its `call` -- Call `_native` directly only for a case the public API cannot reach + +# Rust only + +- Outside the dispatch tests, treat every route's Python implementation as already removed + - never call it on purpose + - never run one request through both engines + - never compare Python and Rust outputs, headers, cache keys or logs +- A test that would still pass with Python serving the call does not belong here +- Every test except the dispatch tests asserts native serving with `support.routing.assert_served_natively` + +# Dispatch tests + +- They check which path a public call takes under given rules + - each route folder's `test_dispatch.py` + - `test_dispatch.py` at the top, which tests `PublicDispatch` with stub routes + - `test_public_entries.py` at the top, which runs every route's public call forms through one table +- They pass `rules=` to a dispatch directly or wrap the whole test in one `support.routing.routed(...)` block, and they stub the native bindings +- They may call Python, since the choice between the two paths is what they test +- They live here until they can run in `tests/unit/` without the compiled extension + +# Selecting a path + +- Use `support.routing.native_route(...)`, which prepends `RUST_REQUIRED` rules to `catalog.RULES` and restores them on exit + - route folders apply it to every test through the autouse fixture in their `conftest.py` + - a test outside a route folder wraps its whole body in one `native_route(...)` block + - a child interpreter script assigns `catalog.RULES` once at its top +- Dispatch tests use `support.routing.routed(...)` instead, which sets `catalog.RULES` to exactly the rules given +- Open either once per test, never per call. They swap a module global, so overlapping blocks (one per task in an `asyncio.gather`) restore out of order and leak rules into later tests +- Never select a backend through `LITELLM_RUST`, `litellm.rust(...)` or any other rebinding of `catalog.RULES` +- No test here depends on the rollout switch + - only no rule, `PYTHON_ONLY` and `RUST_REQUIRED` pick a path without it, so those are the only rollouts used here + - Rust with Python fallback needs the switch, so it is not reachable from this directory + - what `LITELLM_RUST` and `litellm.rust(...)` mean is tested in `tests/unit/rust_bridge/test_configuration.py`, mostly through `decide`, which takes both as plain values +- Tests pass whatever the shipped rollout in `litellm/rust_bridge/catalog.py` says, including for a `PYTHON_ONLY` route # Layout -Each route folder (`ocr/`, `messages/`, `responses/`, `chat_completions/`) uses the same file names, one concern per file - -- `test_dispatch.py` the call reaches the native route -- `test_requests.py` Python inputs that reach the wire -- `test_failures.py` errors raised to the caller and rejection before provider IO -- `test_settings.py` globals, env, named credentials and their precedence -- `test_streaming.py` stream relay, early close and stream-time logging -- `test_callbacks.py` the logger hook contract -- `test_body_hooks.py` handler body hooks such as interceptors and agentic loops -- `test_lifecycle.py` cancellation, concurrency, context and nested calls -- `test_cache.py` the configured cache seen from this route - -Route data and the route's `call` live in `/data.py`, and `/conftest.py` owns the `server` fixture and the autouse `native` fixture. `support/` holds only route-neutral helpers +- One folder per route (`ocr/`, `messages/`, `responses/`, `chat_completions/`), all using the same file names + - `test_dispatch.py` which path the rules pick for each public call, and that the call reaches the native route + - `test_requests.py` Python inputs that reach the wire + - `test_failures.py` errors raised to the caller and rejection before provider IO + - `test_settings.py` globals, env, named credentials and their precedence + - `test_streaming.py` stream relay, early close and stream-time logging + - `test_callbacks.py` the logger hook contract + - `test_body_hooks.py` handler body hooks such as interceptors and agentic loops + - `test_lifecycle.py` cancellation, concurrency, context and nested calls + - `test_cache.py` the configured cache seen from this route +- `embeddings/` holds only `test_dispatch.py` so far +- `/data.py` holds the route's request data and its `call` +- `/conftest.py` owns the `server` fixture and the autouse `native` fixture +- `support/` holds only route-neutral helpers +- Never parametrize one test across routes, except in `test_public_entries.py`, whose subject is that every route serves its public call forms the same way +- Call `_native` directly only for a case the public API cannot reach diff --git a/tests/test_litellm_rust/chat_completions/test_dispatch.py b/tests/test_litellm_rust/chat_completions/test_dispatch.py index 4ed75ea5f95..69d9e8e12df 100644 --- a/tests/test_litellm_rust/chat_completions/test_dispatch.py +++ b/tests/test_litellm_rust/chat_completions/test_dispatch.py @@ -4,14 +4,15 @@ from typing import Final, cast # noqa: TID251 # narrows legacy callable signat import pytest import litellm +from litellm.caching.caching import Cache from litellm.chat_completions import dispatch from litellm.chat_completions.dispatch import ( _ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch ) -from litellm.rust_bridge import catalog +from litellm.integrations.custom_logger import CustomLogger from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteRule, Rules +from litellm.rust_bridge.catalog import Feature, Route, RouteRule, Rules from litellm.rust_bridge.chat_completions.entrypoints import ( NATIVE_ACOMPLETION, NATIVE_COMPLETION, @@ -147,29 +148,6 @@ def test_native_receives_bound_request_and_original_call_shape() -> None: assert call_kwargs["metadata"] is metadata -def test_internal_async_marker_bypasses_native() -> None: - response: Final = ModelResponse() - called: Final[list[bool]] = [] - - def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: records public call shape - called.append(True) - return response - - def native(request: NativeCall) -> ModelResponse: - pytest.fail("acompletion's inner completion call must stay on Python") - - result: Final = _DISPATCH.run( - ("gpt-4o", MESSAGES), - {"acompletion": True}, - python=python, - binding=completion_binding(native), - native=native_call_hook, - rules=RUST_RULES, - ) - assert result is response - assert called == [True] - - @pytest.mark.parametrize( ("args", "kwargs"), ( @@ -202,56 +180,6 @@ def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Map assert captured == [(args, kwargs)] -def test_public_completion_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: - captured: Final[list[NativeCall]] = [] - expected: Final = ModelResponse() - - def native(request: NativeCall) -> ModelResponse: - captured.append(request) - return expected - - NATIVE_COMPLETION.override(native) - monkeypatch.setattr(catalog, "RULES", RUST_RULES) - public_completion: Final = cast(Callable[..., ModelResponse], litellm.completion) - try: - result: Final = public_completion(model="gpt-4o", messages=MESSAGES) - finally: - NATIVE_COMPLETION.reset() - assert result is expected - assert [request.resolved["model"] for request in captured] == ["gpt-4o"] - - -@pytest.mark.asyncio -async def test_public_acompletion_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: - captured: Final[list[NativeCall]] = [] - expected: Final = ModelResponse() - - async def native(request: NativeCall) -> ModelResponse: - captured.append(request) - return expected - - NATIVE_ACOMPLETION.override(native) - monkeypatch.setattr(catalog, "RULES", RUST_RULES) - public_acompletion: Final = cast(Callable[..., Awaitable[ModelResponse]], litellm.acompletion) - try: - result: Final = await public_acompletion(model="gpt-4o", messages=MESSAGES) - finally: - NATIVE_ACOMPLETION.reset() - assert result is expected - assert [request.resolved["model"] for request in captured] == ["gpt-4o"] - - -@pytest.mark.asyncio -async def test_public_completion_calls_keep_the_python_result() -> None: - sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok") - async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok") - - assert isinstance(sync_response, ModelResponse) - assert isinstance(async_response, ModelResponse) - assert sync_response.choices[0].message.content == "ok" - assert async_response.choices[0].message.content == "ok" - - def test_sync_completion_request_projects_public_arguments() -> None: rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) expected: Final = ModelResponse() @@ -277,64 +205,133 @@ def test_sync_completion_request_projects_public_arguments() -> None: assert response is expected -@pytest.mark.asyncio -async def test_async_completion_falls_back_after_native_declines() -> None: - from litellm.rust_bridge.bindings import native_exception_types - - native_types: Final = native_exception_types() - if native_types is None: - pytest.skip("native bridge is unavailable") - declined, _ = native_types - expected: Final = ModelResponse() - rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),) - - async def native(request: NativeCall) -> ModelResponse: - raise declined("unsupported") - - async def python(*args: object, **kwargs: object) -> ModelResponse: - return expected - - binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None) - binding.override(native) - response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", MESSAGES), - {}, - python=python, - binding=binding, - native=native_call_hook, - rules=rules, - ) - - assert response is expected - - -def test_internal_acompletion_marker_bypasses_native() -> None: - rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) - expected: Final = ModelResponse() - - def python(*args: object, **kwargs: object) -> ModelResponse: - return expected - - def native(request: NativeCall) -> ModelResponse: - pytest.fail("acompletion's inner completion call must stay on Python") - - binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", MESSAGES), - {"custom_llm_provider": "openai", "acompletion": True}, - python=python, - binding=binding, - native=native_call_hook, - rules=rules, - ) - - assert response is expected - - def test_positional_parameters_remain_available_to_native_projection() -> None: request: Final = _DISPATCH.request(("anthropic/test-model", MESSAGES, 12.0, 0.25), {}) assert request is not None assert request.resolved["timeout"] == 12.0 assert request.resolved["temperature"] == 0.25 assert request.resolved["messages"] is MESSAGES + + +@pytest.mark.asyncio +async def test_acompletion_marker_binds_positionals_by_the_sync_signature() -> None: + captured: Final[list[NativeCall]] = [] + + async def native(call: NativeCall) -> ModelResponse: + captured.append(call) + return ModelResponse() + + NATIVE_ACOMPLETION.override(native) + public_completion: Final = cast(Callable[..., Awaitable[ModelResponse]], litellm.completion) + try: + _ = await public_completion("gpt-4o", MESSAGES, 30.0, acompletion=True) + finally: + NATIVE_ACOMPLETION.reset() + assert [(request.resolved["timeout"], request.resolved["functions"]) for request in captured] == [(30.0, None)] + + +class _ChatAgentic(CustomLogger): + async def async_should_run_chat_completion_agentic_loop( + self, *args: object, **kwargs: object + ) -> tuple[bool, dict[str, object]]: # pyright: ignore[reportIncompatibleMethodOverride] # test double only needs the override + return False, {} + + +@pytest.mark.parametrize( + ("feature", "kwargs", "settings"), + ( + (Feature.MOCK_RESPONSE, {"mock_response": "hi"}, {}), + pytest.param(Feature.MOCK_RESPONSE, {"mock_tool_calls": [{"id": "call_1"}]}, {}, id="mock-tool-calls"), + pytest.param(Feature.MOCK_RESPONSE, {"mock_timeout": True}, {}, id="mock-timeout"), + (Feature.CALLER_CLIENT, {"client": object()}, {}), + (Feature.TIMEOUT_VARIANTS, {"request_timeout": 5}, {}), + (Feature.AGENTIC_LOOP_HOOK, {}, {"callbacks": [_ChatAgentic()]}), + (Feature.TLS_POLICY, {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), + (Feature.RESPONSE_CACHE, {}, {"cache": Cache()}), + (Feature.STREAMING, {"stream": True}, {}), + (Feature.DROP_PARAMS, {}, {"drop_params": True}), + (Feature.MODIFY_PARAMS, {}, {"modify_params": True}), + (Feature.EXTRA_BODY, {"extra_body": {"top_k": 1}}, {}), + (Feature.RETRIES_FALLBACKS, {"num_retries": 2}, {}), + pytest.param( + Feature.RETRIES_FALLBACKS, + {"context_window_fallback_dict": {"anthropic/claude-sonnet-4-5": "anthropic/claude-opus-4-5"}}, + {}, + id="context-window-fallbacks", + ), + pytest.param( + Feature.RETRIES_FALLBACKS, + {}, + {"model_fallbacks": [{"anthropic/claude-sonnet-4-5": ["anthropic/claude-opus-4-5"]}]}, + id="model-fallbacks", + ), + (Feature.PROMPT_MANAGEMENT, {"prompt_id": "welcome"}, {}), + pytest.param(Feature.PROMPT_MANAGEMENT, {"model": "litellm_agent/gpt-5"}, {}, id="litellm-agent-model"), + (Feature.CACHE_CONTROL_INJECTION, {"enable_prompt_caching": True}, {}), + (Feature.MODEL_PARAM_GATING, {"logit_bias": {"1": 1}}, {}), + ( + Feature.MODEL_LIST, + {"model_list": [{"model_name": "anthropic/claude-sonnet-4-5", "litellm_params": {"model": "gpt-5"}}]}, + {}, + ), + ( + Feature.MODEL_ALIAS_MAP, + {}, + {"model_alias_map": {"anthropic/claude-sonnet-4-5": "anthropic/claude-opus-4-5"}}, + ), + ), +) +def test_each_detected_feature_reaches_its_rule( + monkeypatch: pytest.MonkeyPatch, feature: Feature, kwargs: Mapping[str, object], settings: Mapping[str, object] +) -> None: + rules: Final = ( + RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY, features=frozenset({feature})), + RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED), + ) + python_result: Final = ModelResponse() + native_result: Final = ModelResponse() + + def served(call_kwargs: Mapping[str, object]) -> object: + return _DISPATCH.run( + (), + {"model": "anthropic/claude-sonnet-4-5", "messages": MESSAGES, **call_kwargs}, + python=lambda *call_args, **python_kwargs: python_result, + binding=completion_binding(lambda request: native_result), + native=native_call_hook, + rules=rules, + ) + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "ssl_security_level", None) + monkeypatch.setattr(litellm, "cache", None) + monkeypatch.setattr(litellm, "model_alias_map", {}) + for name in ("num_retries", "model_fallbacks"): + monkeypatch.setattr(litellm, name, None) + for name in ("drop_params", "modify_params", "enable_anthropic_prompt_caching"): + monkeypatch.setattr(litellm, name, False) + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE"): + monkeypatch.delenv(name, raising=False) + assert served({}) is native_result + for name, value in settings.items(): + monkeypatch.setattr(litellm, name, value) + assert served(kwargs) is python_result + + +def test_acompletion_flag_hands_off_only_when_true() -> None: + expected: Final = ModelResponse() + + def native(call: NativeCall) -> ModelResponse: + return expected + + async def anative(call: NativeCall) -> ModelResponse: + pytest.fail("completion(..., acompletion=1) stays sync, since the Python @client only hands off on is True") + + NATIVE_COMPLETION.override(native) + NATIVE_ACOMPLETION.override(anative) + public_completion: Final = cast(Callable[..., object], litellm.completion) + try: + result: Final = public_completion("gpt-4o", MESSAGES, acompletion=1) + finally: + NATIVE_COMPLETION.reset() + NATIVE_ACOMPLETION.reset() + assert result is expected diff --git a/tests/test_litellm_rust/embeddings/test_dispatch.py b/tests/test_litellm_rust/embeddings/test_dispatch.py index 750c88cd235..fbebf1c64e4 100644 --- a/tests/test_litellm_rust/embeddings/test_dispatch.py +++ b/tests/test_litellm_rust/embeddings/test_dispatch.py @@ -1,12 +1,10 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable +from collections.abc import Callable from typing import Final import pytest -from pydantic import TypeAdapter -import litellm from litellm.embeddings import dispatch from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import Route, RouteRule, Rules @@ -15,19 +13,6 @@ from litellm.rust_bridge.public_call import NativeCall, native_call_hook from litellm.types.utils import EmbeddingResponse -@pytest.mark.asyncio -async def test_public_embedding_calls_keep_the_python_result() -> None: - vector: Final = [0.1, 0.2] - - sync_response: Final = litellm.embedding(model="openai/test-model", input="hello", mock_response=vector) - async_response: Final = await litellm.aembedding(model="openai/test-model", input="hello", mock_response=vector) - - assert isinstance(sync_response, EmbeddingResponse) - rows: Final = TypeAdapter(list[dict[str, object]]) - assert rows.validate_python(sync_response.model_dump()["data"])[0]["embedding"] == vector - assert rows.validate_python(async_response.model_dump()["data"])[0]["embedding"] == vector - - def test_sync_embedding_request_projects_public_arguments() -> None: rules: Final[Rules] = (RouteRule(Route.EMBEDDINGS, Rollout.RUST_REQUIRED),) expected: Final = EmbeddingResponse(model="test-model", data=[]) @@ -52,36 +37,3 @@ def test_sync_embedding_request_projects_public_arguments() -> None: ) assert response is expected - - -@pytest.mark.asyncio -async def test_async_embedding_falls_back_after_native_declines() -> None: - from litellm.rust_bridge.bindings import native_exception_types - - native_types: Final = native_exception_types() - if native_types is None: - pytest.skip("native bridge is unavailable") - declined, _ = native_types - expected: Final = EmbeddingResponse(model="test-model", data=[]) - rules: Final[Rules] = (RouteRule(Route.EMBEDDINGS, Rollout.RUST_OPT_OUT),) - - async def native(request: NativeCall) -> EmbeddingResponse: - raise declined("unsupported") - - async def python(*args: object, **kwargs: object) -> EmbeddingResponse: - return expected - - binding: Final[NativeBinding[Callable[[NativeCall], Awaitable[EmbeddingResponse]]]] = NativeBinding( - "aembedding", validate=lambda _: None - ) - binding.override(native) - response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - ("test-model", "hello"), - {}, - python=python, - binding=binding, - native=native_call_hook, - rules=rules, - ) - - assert response is expected diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index e8d49453065..cd5478dfc3a 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -67,7 +67,7 @@ async def test_success_hook_fires_the_expected_number_of_times( pytest.param("async_log_failure_event", 1), pytest.param("async_post_call_failure_deployment_hook", 1), pytest.param("log_model_group_rate_limit_error", 1), - pytest.param("log_failure_event", 0, marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("log_failure_event", 0), ), ) async def test_failure_hook_fires_the_expected_number_of_times( diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 2a32e31c250..50f71569d73 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -3,9 +3,11 @@ from collections.abc import Awaitable, Callable, Mapping from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect import pytest -from pydantic import TypeAdapter import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.llms.anthropic.pass_through.messages import handler as python_messages from litellm.messages import dispatch from litellm.messages.dispatch import ( @@ -13,17 +15,15 @@ from litellm.messages.dispatch import ( _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch ) from litellm.proxy import proxy_server -from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteRule, Rules +from litellm.rust_bridge.catalog import Feature, Route, RouteRule, Rules, Target from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES, NativeAmessages, NativeMessages from litellm.rust_bridge.public_call import NativeCall +from litellm.types.llms.anthropic import ANTHROPIC_ADVISOR_TOOL_TYPE from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse MESSAGES: Final = [{"role": "user", "content": "hi"}] -DANGEROUS_TOOL_USE: Final = [{"type": "dangerous_tool_use", "classifier_context": {"permission_mode": "auto"}}] -CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "gpt-oss-120b"} PYTHON_RULES: Final[Rules] = () RUST_RULES: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) @@ -194,81 +194,44 @@ def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_tim assert captured[0].resolved["timeout"] == expected_timeout -def test_internal_async_marker_bypasses_native() -> None: - args: Final[tuple[object, ...]] = (16, MESSAGES, "claude-sonnet-4-5") - kwargs: Final[Mapping[str, object]] = {"is_async": True} - captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = [] - expected: Final = response() - - def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape - captured.append((call_args, call_kwargs)) - return expected - - def native( - request: NativeCall, - ) -> AnthropicMessagesResponse: - pytest.fail("The async handler's inner sync call must stay on Python") - - result: Final = _DISPATCH.run( - args, - kwargs, - python=python, - binding=messages_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - assert result is expected - assert captured == [(args, kwargs)] - - @pytest.mark.parametrize( - ("model", "custom_llm_provider"), - [ - ("vertex_ai/gemini-2.5-pro", None), - ("gemini-2.5-pro", "vertex_ai"), - ], + ("model", "kwargs", "passes_through"), + ( + ("vertex_ai/claude-sonnet-4-5@20250929", {}, True), + ("claude-sonnet-4-5@20250929", {"custom_llm_provider": "vertex_ai"}, True), + ("deepseek/deepseek-v4-flash", {}, True), + ("openai/gpt-5", {"model_info": {"supported_endpoints": ["/v1/messages"]}}, True), + ("vertex_ai/gemini-2.5-pro", {}, False), + ("gemini-2.5-pro", {"custom_llm_provider": "vertex_ai"}, False), + ("openai/gpt-5", {}, False), + ("groq/llama-3.3-70b-versatile", {}, False), + ("unknown-provider/unknown-model", {}, False), + ), ) -def test_vertex_models_other_than_claude_stay_on_python(model: str, custom_llm_provider: str | None) -> None: - args: Final[tuple[object, ...]] = (16, MESSAGES, model) - kwargs: Final[Mapping[str, object]] = {"custom_llm_provider": custom_llm_provider} - expected: Final = response(model) +def test_only_requests_that_pass_through_to_a_messages_api_reach_a_messages_target_rule( + model: str, kwargs: Mapping[str, object], passes_through: bool +) -> None: + python_result: Final = response(model) + native_result: Final = response(model) - def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape - return expected - - def native(request: NativeCall) -> AnthropicMessagesResponse: - pytest.fail("Rust serves only Claude on Vertex AI") + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback + return python_result result: Final = _DISPATCH.run( - args, + (16, MESSAGES, model), kwargs, python=python, - binding=messages_binding(native), + binding=messages_binding(lambda request: native_result), native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, + rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, targets=frozenset({Target.MESSAGES})),), ) - assert result is expected + assert result is (native_result if passes_through else python_result) -def test_vertex_claude_reaches_native() -> None: - args: Final[tuple[object, ...]] = (16, MESSAGES, "vertex_ai/claude-sonnet-4-5@20250929") - expected: Final = response("vertex_ai/claude-sonnet-4-5@20250929") - - def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape - pytest.fail("Claude on Vertex AI is served natively") - - result: Final = _DISPATCH.run( - args, - {}, - python=python, - binding=messages_binding(lambda request: expected), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - assert result is expected +DANGEROUS_TOOL_USE: Final = [{"type": "dangerous_tool_use", "classifier_context": {"permission_mode": "auto"}}] +CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "gpt-oss-120b"} -@pytest.mark.asyncio @pytest.mark.parametrize( ("model", "kwargs", "general_settings", "serves_python"), ( @@ -285,7 +248,7 @@ def test_vertex_claude_reaches_native() -> None: ("anthropic/claude-sonnet-4-5", {"safeguards": DANGEROUS_TOOL_USE}, CLASSIFIER_SETTINGS, False), ), ) -async def test_python_keeps_the_requests_it_answers_with_safeguard_results( +def test_proxy_classified_safeguards_reach_the_feature_rule( monkeypatch: pytest.MonkeyPatch, model: str, kwargs: Mapping[str, object], @@ -296,19 +259,19 @@ async def test_python_keeps_the_requests_it_answers_with_safeguard_results( python_result: Final = response(model) native_result: Final = response(model) - async def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback return python_result - async def native(request: NativeCall) -> AnthropicMessagesResponse: - return native_result - - result: Final = await _ADISPATCH.arun( + result: Final = _DISPATCH.run( (16, MESSAGES, model), kwargs, python=python, - binding=amessages_binding(native), + binding=messages_binding(lambda request: native_result), native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, + rules=( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ), ) assert result is (python_result if serves_python else native_result) @@ -345,7 +308,7 @@ def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Map assert captured == [(args, kwargs)] -def test_anthropic_create_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: +def test_anthropic_create_routes_through_dispatch() -> None: captured: Final[list[NativeCall]] = [] expected: Final = response() @@ -356,7 +319,6 @@ def test_anthropic_create_routes_through_dispatch(monkeypatch: pytest.MonkeyPatc return expected NATIVE_MESSAGES.override(native) - monkeypatch.setattr(catalog, "RULES", RUST_RULES) public_create: Final = cast(Callable[..., AnthropicMessagesResponse], litellm.anthropic.create) try: result: Final = public_create(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5") @@ -367,7 +329,7 @@ def test_anthropic_create_routes_through_dispatch(monkeypatch: pytest.MonkeyPatc @pytest.mark.asyncio -async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_anthropic_acreate_routes_through_dispatch() -> None: captured: Final[list[NativeCall]] = [] expected: Final = response() @@ -378,7 +340,6 @@ async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.Mon return expected NATIVE_AMESSAGES.override(native) - monkeypatch.setattr(catalog, "RULES", RUST_RULES) public_acreate: Final = cast(Callable[..., Awaitable[AnthropicMessagesResponse]], litellm.anthropic.acreate) try: result: Final = await public_acreate(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5") @@ -389,14 +350,31 @@ async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.Mon @pytest.mark.asyncio -async def test_public_anthropic_messages_keeps_the_python_result() -> None: - response: Final = await litellm.anthropic_messages( - model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok" - ) +async def test_truthy_is_async_hands_off_to_the_async_entry_like_true() -> None: + captured: Final[list[NativeCall]] = [] + expected: Final = response() - assert isinstance(response, dict) - content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", [])) - assert content[0]["text"] == "ok" + async def native(request: NativeCall) -> AnthropicMessagesResponse: + captured.append(request) + return expected + + def sync_native(request: NativeCall) -> AnthropicMessagesResponse: + pytest.fail("a truthy is_async must hand off to the async entry, as the Python handler does") + + NATIVE_AMESSAGES.override(native) + NATIVE_MESSAGES.override(sync_native) + public_create: Final = cast(Callable[..., object], litellm.anthropic.create) + try: + pending: Final = public_create(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5", is_async=1) + assert inspect.isawaitable(pending) + assert await pending is expected + reference: Final = public_create(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5", is_async=True) + assert inspect.isawaitable(reference) + assert await reference is expected + finally: + NATIVE_AMESSAGES.reset() + NATIVE_MESSAGES.reset() + assert [dict(request.resolved) for request in captured] == [dict(captured[1].resolved)] * 2 def test_sync_messages_request_projects_public_arguments() -> None: @@ -453,62 +431,173 @@ def test_messages_binding_error_delegates_unchanged_to_python() -> None: assert response is expected -@pytest.mark.asyncio -async def test_async_messages_falls_back_after_native_declines() -> None: - from litellm.rust_bridge.bindings import native_exception_types +class _Rewrite(CustomLogger): + async def async_pre_request_hook(self, model: str, messages: list, kwargs: dict) -> dict | None: # pyright: ignore[reportMissingTypeArgument] # mirrors the CustomLogger signature + return kwargs - native_types: Final = native_exception_types() - if native_types is None: - pytest.skip("native bridge is unavailable") - declined, _ = native_types - expected: Final = AnthropicMessagesResponse(model="claude-test") - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),) + +class _Agentic(CustomLogger): + async def async_should_run_agentic_loop(self, *args: object, **kwargs: object) -> tuple[bool, dict]: # pyright: ignore[reportIncompatibleMethodOverride, reportMissingTypeArgument] # test double only needs the override + return False, {} + + +ADVISOR_TOOL: Final = {"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": "claude-opus-5-5"} + + +@pytest.mark.parametrize( + ("feature", "model", "kwargs", "settings"), + ( + (Feature.MOCK_RESPONSE, "anthropic/claude-sonnet-4-5", {"mock_response": "hi"}, {}), + ( + Feature.PROXY_MCP_TOOLS, + "anthropic/claude-sonnet-4-5", + {"tools": [{"type": "mcp", "server_label": "gateway", "server_url": "litellm_proxy"}]}, + {}, + ), + (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), + (Feature.TIMEOUT_VARIANTS, "anthropic/claude-sonnet-4-5", {"request_timeout": 5}, {}), + (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), + (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), + (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), + (Feature.MODIFY_PARAMS, "anthropic/claude-sonnet-4-5", {}, {"modify_params": True}), + ( + Feature.WORKLOAD_IDENTITY, + "anthropic/claude-sonnet-4-5", + {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"}, + {}, + ), + (Feature.CACHE_CONTROL_INJECTION, "anthropic/claude-sonnet-4-5", {"enable_prompt_caching": True}, {}), + ( + Feature.THINKING_SIGNATURE_RETRY, + "anthropic/claude-sonnet-4-5", + {"messages": [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "x"}]}]}, + {}, + ), + ( + Feature.STREAM_COST_INJECTION, + "anthropic/claude-sonnet-4-5", + {"stream": True}, + {"include_cost_in_streaming_usage": True}, + ), + ), +) +def test_each_detected_feature_reaches_its_rule( + monkeypatch: pytest.MonkeyPatch, + feature: Feature, + model: str, + kwargs: Mapping[str, object], + settings: Mapping[str, object], +) -> None: + rules: Final[Rules] = ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({feature})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + python_result: Final = response(model) + native_result: Final = response(model) + + def served(call_kwargs: Mapping[str, object]) -> AnthropicMessagesResponse: + messages: Final = call_kwargs.get("messages", MESSAGES) + return _DISPATCH.run( + (16, messages, model), + {name: value for name, value in call_kwargs.items() if name != "messages"}, + python=lambda *call_args, **python_kwargs: python_result, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, call_args, native_kwargs: hook(request), + rules=rules, + ) + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "cache", None) + monkeypatch.setattr(litellm, "ssl_security_level", None) + for setting in ("modify_params", "enable_anthropic_prompt_caching", "include_cost_in_streaming_usage"): + monkeypatch.setattr(litellm, setting, False) + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE", "ANTHROPIC_FEDERATION_RULE_ID", "ANTHROPIC_ORGANIZATION_ID"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("ANTHROPIC_API_KEY", "") + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) + assert served({}) is native_result + for name, value in settings.items(): + monkeypatch.setattr(litellm, name, value) + assert served(kwargs) is python_result + + +ASYNC_ONLY_FEATURES: Final = ( + pytest.param(Feature.PRE_REQUEST_HOOK, {}, {"callbacks": [_Rewrite()]}, id="pre-request-hook"), + pytest.param( + Feature.WEBSEARCH_INTERCEPTION, {}, {"callbacks": [WebSearchInterceptionLogger()]}, id="websearch-interception" + ), + pytest.param( + Feature.MESSAGES_INTERCEPTOR, + {"model": "deepseek/deepseek-v4-flash", "tools": [ADVISOR_TOOL]}, + {}, + id="messages-interceptor", + ), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("feature", "kwargs", "settings"), ASYNC_ONLY_FEATURES) +async def test_handler_hooks_the_sync_entry_never_runs_keep_only_async_calls_on_python( + monkeypatch: pytest.MonkeyPatch, feature: Feature, kwargs: Mapping[str, object], settings: Mapping[str, object] +) -> None: + rules: Final[Rules] = ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({feature})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + python_result: Final = response() + native_result: Final = response() + model: Final = str(kwargs.get("model", "anthropic/claude-sonnet-4-5")) + call_kwargs: Final = {name: value for name, value in kwargs.items() if name != "model"} + + async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + return python_result async def native(request: NativeCall) -> AnthropicMessagesResponse: - raise declined("unsupported") + return native_result - async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: - return expected - - binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None) - binding.override(native) - response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - {"model": "claude-test", "messages": MESSAGES, "max_tokens": 10}, + monkeypatch.setattr(litellm, "callbacks", []) + for name, value in settings.items(): + monkeypatch.setattr(litellm, name, value) + served_async: Final = await _ADISPATCH.arun( + (16, MESSAGES, model), + call_kwargs, python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request), + binding=amessages_binding(native), + native=lambda hook, request, args, native_kwargs: hook(request), rules=rules, ) - - assert response is expected - - -def test_internal_is_async_marker_bypasses_native() -> None: - rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) - expected: Final = AnthropicMessagesResponse(model="claude-test") - - def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse: - return expected - - def native(request: NativeCall) -> AnthropicMessagesResponse: - pytest.fail("anthropic_messages' inner handler call must stay on Python") - - binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None) - binding.override(native) - response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision - (), - { - "model": "claude-test", - "messages": MESSAGES, - "max_tokens": 10, - "custom_llm_provider": "anthropic", - "is_async": True, - }, - python=python, - binding=binding, - native=lambda hook, request, args, kwargs: hook(request), + served_sync: Final = _DISPATCH.run( + (16, MESSAGES, model), + call_kwargs, + python=lambda *call_args, **python_kwargs: python_result, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, args, native_kwargs: hook(request), rules=rules, ) + assert served_async is python_result + assert served_sync is native_result - assert response is expected + +@pytest.mark.asyncio +async def test_interceptor_feature_follows_the_resolved_provider() -> None: + rules: Final[Rules] = ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.MESSAGES_INTERCEPTOR})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + native_result: Final = response() + + async def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + pytest.fail("anthropic serves the advisor tool itself") + + async def native(request: NativeCall) -> AnthropicMessagesResponse: + return native_result + + result: Final = await _ADISPATCH.arun( + (16, MESSAGES, "claude-sonnet-4-5"), + {"tools": [ADVISOR_TOOL], "custom_llm_provider": "anthropic"}, + python=python, + binding=amessages_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=rules, + ) + assert result is native_result diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 5f961288800..18925599acf 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -237,7 +237,7 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp with pytest.raises(litellm.InternalServerError): await call_aocr(server, callbacks=[Observe()]) - assert [observation[0] for observation in observations] == ["sync", "async"] + assert [observation[0] for observation in observations] == ["async"] assert all(observation[1] == "aocr" for observation in observations) assert all(isinstance(observation[2], litellm.InternalServerError) for observation in observations) assert all(observation[3] is None for observation in observations) @@ -278,7 +278,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal with pytest.raises(litellm.InternalServerError): await call_aocr(server, callbacks=[TrackInFlightRequest()]) - assert [event for event, _ in observed] == ["sync", "async"] + assert [event for event, _ in observed] == ["async"] assert all(observed_token is token for _, observed_token in observed) @@ -445,7 +445,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f await call_aocr(server) await drain_logging() - assert observed == [("sync", token, raised.value), ("async", token, raised.value)] + assert observed == [("async", token, raised.value)] @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py index db60b6aea84..3206cc606a8 100644 --- a/tests/test_litellm_rust/ocr/test_dispatch.py +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -1,8 +1,21 @@ +from collections.abc import Mapping from typing import Final, NoReturn +import httpx import pytest import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.ocr.dispatch import ( + _ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch + _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch +) +from litellm.rust_bridge import runtime +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.catalog import Route, RouteRule, Rules +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.ocr.entrypoints import NativeAocr, NativeOcr +from litellm.rust_bridge.public_call import NativeCall from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.routing import assert_served_natively @@ -110,3 +123,214 @@ async def test_ocr_health_check_uses_the_native_route(recording_server: Recordin assert "error" not in response assert_served_natively(recording_server) assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,") + + +RUST_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),) + + +def ocr_binding(native: NativeOcr | None) -> NativeBinding[NativeOcr]: + binding: Final[NativeBinding[NativeOcr]] = NativeBinding("ocr", validate=lambda _: None) + binding.override(native) + return binding + + +def aocr_binding(native: NativeAocr | None) -> NativeBinding[NativeAocr]: + binding: Final[NativeBinding[NativeAocr]] = NativeBinding("aocr", validate=lambda _: None) + binding.override(native) + return binding + + +def response(model: str = "mistral/mistral-ocr-latest") -> OCRResponse: + return OCRResponse(pages=[], model=model) + + +def test_native_receives_normalized_positional_request_and_original_call_shape() -> None: + document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"} + timeout: Final = httpx.Timeout(30) + extra_headers: Final[dict[str, object]] = {"x-test": "1"} + pages: Final = [0, 2] + args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document) + kwargs: Final[Mapping[str, object]] = { + "api_key": "test-key", + "api_base": "https://example.invalid", + "timeout": timeout, + "custom_llm_provider": "mistral", + "extra_headers": extra_headers, + "pages": pages, + } + captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = [] + expected: Final = response() + + def native( + request: NativeCall, + ) -> OCRResponse: + args: Final = request.args + kwargs: Final = request.kwargs + captured.append((request, args, kwargs)) + return expected + + result: Final = _DISPATCH.run( + args, + kwargs, + python=runtime.NO_PYTHON, + binding=ocr_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + + request, call_args, call_kwargs = captured[0] + assert result is expected + assert request.resolved["model"] == "mistral/mistral-ocr-latest" + assert request.resolved["document"] is document + assert request.resolved["api_key"] == "test-key" + assert request.resolved["api_base"] == "https://example.invalid" + assert request.resolved["timeout"] is timeout + assert request.resolved["custom_llm_provider"] == "mistral" + assert request.resolved["extra_headers"] is extra_headers + assert request.kwargs == kwargs + assert request.kwargs["pages"] is pages + assert call_args is args + assert call_kwargs is kwargs + + +def test_native_preserves_keyword_model_and_document_in_original_call_shape() -> None: + document: Final[Mapping[str, object]] = { + "type": "document_url", + "document_url": "https://example.invalid/document.pdf", + } + pages: Final = [1] + args: Final[tuple[object, ...]] = () + kwargs: Final[Mapping[str, object]] = { + "model": "mistral/mistral-ocr-latest", + "document": document, + "pages": pages, + } + captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = [] + expected: Final = response() + + def native( + request: NativeCall, + ) -> OCRResponse: + args: Final = request.args + kwargs: Final = request.kwargs + captured.append((request, args, kwargs)) + return expected + + result: Final = _DISPATCH.run( + args, + kwargs, + python=runtime.NO_PYTHON, + binding=ocr_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + + request, call_args, call_kwargs = captured[0] + assert result is expected + assert request.resolved["model"] == "mistral/mistral-ocr-latest" + assert request.resolved["document"] is document + assert request.kwargs == kwargs + assert call_args is args + assert call_kwargs is kwargs + assert call_kwargs["model"] == "mistral/mistral-ocr-latest" + assert call_kwargs["document"] is document + + +def test_missing_native_binding_is_a_required_rust_error() -> None: + args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}) + + with pytest.raises(RuntimeError, match="Rust ocr bridge"): + _DISPATCH.run( + args, + {}, + python=runtime.NO_PYTHON, + binding=ocr_binding(None), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + + +def test_non_required_rule_cannot_be_served_without_python() -> None: + args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}) + + def native( + request: NativeCall, + ) -> OCRResponse: + return response() + + with pytest.raises(runtime.NoPythonImplementationError, match="RUST_REQUIRED"): + _DISPATCH.run( + args, + {}, + python=runtime.NO_PYTHON, + binding=ocr_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=(RouteRule(Route.OCR, Rollout.PYTHON_ONLY),), + ) + + +@pytest.mark.parametrize( + ("args", "kwargs", "message"), + ( + ( + ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}), + {"model": "duplicate"}, + r"ocr\(\) got multiple values for argument 'model'", + ), + ( + ("mistral/mistral-ocr-latest",), + {}, + r"ocr\(\) missing 1 required positional argument: 'document'", + ), + ), +) +def test_ocr_parser_errors_before_native(args: tuple[object, ...], kwargs: Mapping[str, object], message: str) -> None: + def native( + request: NativeCall, + ) -> OCRResponse: + pytest.fail("OCR parser failures must not call native") + + with pytest.raises(TypeError, match=message): + _DISPATCH.run( + args, + kwargs, + python=runtime.NO_PYTHON, + binding=ocr_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("args", "kwargs", "message"), + ( + ( + ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}), + {"model": "duplicate"}, + r"aocr\(\) got multiple values for argument 'model'", + ), + ( + ("mistral/mistral-ocr-latest",), + {}, + r"aocr\(\) missing 1 required positional argument: 'document'", + ), + ), +) +async def test_aocr_parser_errors_before_native( + args: tuple[object, ...], kwargs: Mapping[str, object], message: str +) -> None: + async def native( + request: NativeCall, + ) -> OCRResponse: + pytest.fail("OCR parser failures must not call native") + + with pytest.raises(TypeError, match=message): + await _ADISPATCH.arun( + args, + kwargs, + python=runtime.NO_PYTHON, + binding=aocr_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index c35f6836144..611316fb52c 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -144,7 +144,7 @@ async def test_mapped_failure_identity_and_deployment_snapshot(server: Recording with pytest.raises(litellm.InternalServerError) as caught: await call_aocr(server, callbacks=[recorder]) failures: Final = tuple(event for event in recorder.events if "failure" in event.name) - assert [event.name for event in failures] == ["log_failure_event", "async_log_failure_event"] + assert [event.name for event in failures] == ["async_log_failure_event"] assert all(event.kwargs["exception"] is caught.value for event in failures) assert caught.value.status_code == 500 assert snapshots[0] is not caught.value @@ -324,24 +324,26 @@ async def test_failure_callbacks_continue_within_both_families(server: Recording server.enqueue(ResponseSpec(body={"message": "failed"}, status=500)) observed: Final = [] - class Broken(CustomLogger): + class BrokenSync(CustomLogger): def log_failure_event(self, kwargs, response_obj, start_time, end_time): observed.append(("broken-sync", kwargs["exception"])) raise RuntimeError("sync observer") + class FollowingSync(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("following-sync", kwargs["exception"])) + + class BrokenAsync(CustomLogger): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): observed.append(("broken-async", kwargs["exception"])) raise RuntimeError("async observer") - class Following(CustomLogger): - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - observed.append(("following-sync", kwargs["exception"])) - + class FollowingAsync(CustomLogger): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): observed.append(("following-async", kwargs["exception"])) with pytest.raises(litellm.InternalServerError) as caught: - await call_aocr(server, callbacks=[Broken(), Following()]) + await call_aocr(server, callbacks=[BrokenSync(), FollowingSync(), BrokenAsync(), FollowingAsync()]) assert [name for name, _ in observed] == ["broken-sync", "following-sync", "broken-async", "following-async"] assert all(error is caught.value for _, error in observed) diff --git a/tests/test_litellm_rust/ocr/test_public_dispatch.py b/tests/test_litellm_rust/ocr/test_public_dispatch.py deleted file mode 100644 index 34ee8e40552..00000000000 --- a/tests/test_litellm_rust/ocr/test_public_dispatch.py +++ /dev/null @@ -1,300 +0,0 @@ -from collections.abc import Awaitable, Callable, Mapping -from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.ocr.dispatch import ( - _ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch - _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch -) -from litellm.rust_bridge import catalog, runtime -from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteRule, Rules -from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR, NativeAocr, NativeOcr -from litellm.rust_bridge.public_call import NativeCall - -RUST_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),) - - -def ocr_binding(native: NativeOcr | None) -> NativeBinding[NativeOcr]: - binding: Final[NativeBinding[NativeOcr]] = NativeBinding("ocr", validate=lambda _: None) - binding.override(native) - return binding - - -def aocr_binding(native: NativeAocr | None) -> NativeBinding[NativeAocr]: - binding: Final[NativeBinding[NativeAocr]] = NativeBinding("aocr", validate=lambda _: None) - binding.override(native) - return binding - - -def response(model: str = "mistral/mistral-ocr-latest") -> OCRResponse: - return OCRResponse(pages=[], model=model) - - -def test_native_receives_normalized_positional_request_and_original_call_shape() -> None: - document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"} - timeout: Final = httpx.Timeout(30) - extra_headers: Final[dict[str, object]] = {"x-test": "1"} - pages: Final = [0, 2] - args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document) - kwargs: Final[Mapping[str, object]] = { - "api_key": "test-key", - "api_base": "https://example.invalid", - "timeout": timeout, - "custom_llm_provider": "mistral", - "extra_headers": extra_headers, - "pages": pages, - } - captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = [] - expected: Final = response() - - def native( - request: NativeCall, - ) -> OCRResponse: - args: Final = request.args - kwargs: Final = request.kwargs - captured.append((request, args, kwargs)) - return expected - - result: Final = _DISPATCH.run( - args, - kwargs, - python=runtime.NO_PYTHON, - binding=ocr_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - - request, call_args, call_kwargs = captured[0] - assert result is expected - assert request.resolved["model"] == "mistral/mistral-ocr-latest" - assert request.resolved["document"] is document - assert request.resolved["api_key"] == "test-key" - assert request.resolved["api_base"] == "https://example.invalid" - assert request.resolved["timeout"] is timeout - assert request.resolved["custom_llm_provider"] == "mistral" - assert request.resolved["extra_headers"] is extra_headers - assert request.kwargs == kwargs - assert request.kwargs["pages"] is pages - assert call_args is args - assert call_kwargs is kwargs - - -def test_native_preserves_keyword_model_and_document_in_original_call_shape() -> None: - document: Final[Mapping[str, object]] = { - "type": "document_url", - "document_url": "https://example.invalid/document.pdf", - } - pages: Final = [1] - args: Final[tuple[object, ...]] = () - kwargs: Final[Mapping[str, object]] = { - "model": "mistral/mistral-ocr-latest", - "document": document, - "pages": pages, - } - captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = [] - expected: Final = response() - - def native( - request: NativeCall, - ) -> OCRResponse: - args: Final = request.args - kwargs: Final = request.kwargs - captured.append((request, args, kwargs)) - return expected - - result: Final = _DISPATCH.run( - args, - kwargs, - python=runtime.NO_PYTHON, - binding=ocr_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - - request, call_args, call_kwargs = captured[0] - assert result is expected - assert request.resolved["model"] == "mistral/mistral-ocr-latest" - assert request.resolved["document"] is document - assert request.kwargs == kwargs - assert call_args is args - assert call_kwargs is kwargs - assert call_kwargs["model"] == "mistral/mistral-ocr-latest" - assert call_kwargs["document"] is document - - -def test_aocr_marker_cannot_be_served_without_python() -> None: - document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"} - args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document) - kwargs: Final[Mapping[str, object]] = {"aocr": True} - - def native( - request: NativeCall, - ) -> OCRResponse: - pytest.fail("the aocr bypass marker must not reach native") - - with pytest.raises(runtime.NoPythonImplementationError, match="bypass"): - _DISPATCH.run( - args, - kwargs, - python=runtime.NO_PYTHON, - binding=ocr_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - - -def test_missing_native_binding_is_a_required_rust_error() -> None: - args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}) - - with pytest.raises(RuntimeError, match="Rust ocr bridge"): - _DISPATCH.run( - args, - {}, - python=runtime.NO_PYTHON, - binding=ocr_binding(None), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - - -def test_non_required_rule_cannot_be_served_without_python() -> None: - args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}) - - def native( - request: NativeCall, - ) -> OCRResponse: - return response() - - with pytest.raises(runtime.NoPythonImplementationError, match="RUST_REQUIRED"): - _DISPATCH.run( - args, - {}, - python=runtime.NO_PYTHON, - binding=ocr_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=(RouteRule(Route.OCR, Rollout.PYTHON_ONLY),), - ) - - -@pytest.mark.parametrize( - ("args", "kwargs", "message"), - ( - ( - ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}), - {"model": "duplicate"}, - r"ocr\(\) got multiple values for argument 'model'", - ), - ( - ("mistral/mistral-ocr-latest",), - {}, - r"ocr\(\) missing 1 required positional argument: 'document'", - ), - ), -) -def test_ocr_parser_errors_before_native(args: tuple[object, ...], kwargs: Mapping[str, object], message: str) -> None: - def native( - request: NativeCall, - ) -> OCRResponse: - pytest.fail("OCR parser failures must not call native") - - with pytest.raises(TypeError, match=message): - _DISPATCH.run( - args, - kwargs, - python=runtime.NO_PYTHON, - binding=ocr_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("args", "kwargs", "message"), - ( - ( - ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}), - {"model": "duplicate"}, - r"aocr\(\) got multiple values for argument 'model'", - ), - ( - ("mistral/mistral-ocr-latest",), - {}, - r"aocr\(\) missing 1 required positional argument: 'document'", - ), - ), -) -async def test_aocr_parser_errors_before_native( - args: tuple[object, ...], kwargs: Mapping[str, object], message: str -) -> None: - async def native( - request: NativeCall, - ) -> OCRResponse: - pytest.fail("OCR parser failures must not call native") - - with pytest.raises(TypeError, match=message): - await _ADISPATCH.arun( - args, - kwargs, - python=runtime.NO_PYTHON, - binding=aocr_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - - -def test_public_ocr_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: - document: Final[Mapping[str, object]] = { - "type": "document_url", - "document_url": "https://example.invalid/document.pdf", - } - captured: Final[list[NativeCall]] = [] - expected: Final = response() - - def native( - request: NativeCall, - ) -> OCRResponse: - captured.append(request) - return expected - - NATIVE_OCR.override(native) - monkeypatch.setattr(catalog, "RULES", RUST_RULES) - public_ocr: Final = cast(Callable[..., OCRResponse], litellm.ocr) - try: - result: Final = public_ocr(model="mistral/mistral-ocr-latest", document=document) - finally: - NATIVE_OCR.reset() - assert result is expected - assert [request.resolved["model"] for request in captured] == ["mistral/mistral-ocr-latest"] - - -@pytest.mark.asyncio -async def test_public_aocr_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: - document: Final[Mapping[str, object]] = { - "type": "document_url", - "document_url": "https://example.invalid/document.pdf", - } - captured: Final[list[NativeCall]] = [] - expected: Final = response() - - async def native( - request: NativeCall, - ) -> OCRResponse: - captured.append(request) - return expected - - NATIVE_AOCR.override(native) - monkeypatch.setattr(catalog, "RULES", RUST_RULES) - public_aocr: Final = cast(Callable[..., Awaitable[OCRResponse]], litellm.aocr) - try: - result: Final = await public_aocr(model="mistral/mistral-ocr-latest", document=document) - finally: - NATIVE_AOCR.reset() - assert result is expected - assert [request.resolved["model"] for request in captured] == ["mistral/mistral-ocr-latest"] diff --git a/tests/test_litellm_rust/responses/data.py b/tests/test_litellm_rust/responses/data.py index 604a7de7442..0347fd7a57b 100644 --- a/tests/test_litellm_rust/responses/data.py +++ b/tests/test_litellm_rust/responses/data.py @@ -69,5 +69,7 @@ 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 declines to Python", raises=RuntimeError) +STREAMING: Final = gap( + "native Responses streaming is unsupported, the catalog keeps it on Python", raises=litellm.APIConnectionError +) BODY_HOOKS: Final = gap("native Responses does not run the handler body hooks") diff --git a/tests/test_litellm_rust/responses/test_callbacks.py b/tests/test_litellm_rust/responses/test_callbacks.py index a4abdd43529..0e921e9d5db 100644 --- a/tests/test_litellm_rust/responses/test_callbacks.py +++ b/tests/test_litellm_rust/responses/test_callbacks.py @@ -95,7 +95,7 @@ async def test_success_hook_fires_the_expected_number_of_times( "async_post_call_failure_deployment_hook", True, 1, id="async-async_post_call_failure_deployment_hook" ), pytest.param("log_model_group_rate_limit_error", True, 1, id="async-log_model_group_rate_limit_error"), - pytest.param("log_failure_event", True, 0, id="async-log_failure_event", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("log_failure_event", True, 0, id="async-log_failure_event"), pytest.param("log_failure_event", False, 1, id="sync-log_failure_event"), ), ) diff --git a/tests/test_litellm_rust/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py index e6a0193a063..4498dbdf39d 100644 --- a/tests/test_litellm_rust/responses/test_dispatch.py +++ b/tests/test_litellm_rust/responses/test_dispatch.py @@ -1,27 +1,28 @@ import inspect -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Callable, Mapping from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect import pytest import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger from litellm.responses import dispatch as responses_dispatch from litellm.responses import main as python_responses from litellm.responses.dispatch import ( _ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch ) -from litellm.rust_bridge import catalog +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.catalog import Feature, Route, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.public_call import NativeCall from litellm.rust_bridge.responses.entrypoints import ( - NATIVE_ARESPONSES, - NATIVE_RESPONSES, NativeAresponses, NativeResponses, ) +from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_TOOL_TYPES from litellm.types.llms.openai import ResponsesAPIResponse INPUT: Final = [{"role": "user", "content": "hi"}] @@ -181,35 +182,6 @@ def test_native_receives_normalized_request_and_original_call_shape() -> None: assert call_kwargs["litellm_metadata"] is metadata -def test_internal_async_marker_bypasses_native() -> None: - args: Final[tuple[object, ...]] = (INPUT, "gpt-4o") - kwargs: Final[Mapping[str, object]] = {"aresponses": True} - captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = [] - response: Final = _response() - - def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: records call shape - captured.append((call_args, call_kwargs)) - return response - - def native( - request: NativeCall, - ) -> ResponsesAPIResponse: - pytest.fail("aresponses' inner responses call must stay on Python") - - assert ( - _DISPATCH.run( - args, - kwargs, - python=python, - binding=responses_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=RUST_RULES, - ) - is response - ) - assert captured == [(args, kwargs)] - - @pytest.mark.parametrize( ("args", "kwargs"), ( @@ -244,49 +216,6 @@ def test_binding_errors_delegate_unchanged_to_python(args: tuple[object, ...], k assert captured == [(args, kwargs)] -def test_public_responses_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: - captured: Final[list[NativeCall]] = [] - expected: Final = _response() - - def native( - request: NativeCall, - ) -> ResponsesAPIResponse: - captured.append(request) - return expected - - NATIVE_RESPONSES.override(native) - monkeypatch.setattr(catalog, "RULES", RUST_RULES) - public_responses: Final = cast(Callable[..., ResponsesAPIResponse], litellm.responses) - try: - result: Final = public_responses(input=INPUT, model="gpt-4o") - finally: - NATIVE_RESPONSES.reset() - assert result is expected - assert [request.resolved["model"] for request in captured] == ["gpt-4o"] - - -@pytest.mark.asyncio -async def test_public_aresponses_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: - captured: Final[list[NativeCall]] = [] - expected: Final = _response() - - async def native( - request: NativeCall, - ) -> ResponsesAPIResponse: - captured.append(request) - return expected - - NATIVE_ARESPONSES.override(native) - monkeypatch.setattr(catalog, "RULES", RUST_RULES) - public_aresponses: Final = cast(Callable[..., Awaitable[ResponsesAPIResponse]], litellm.aresponses) - try: - result: Final = await public_aresponses(input=INPUT, model="gpt-4o") - finally: - NATIVE_ARESPONSES.reset() - assert result is expected - assert [request.resolved["model"] for request in captured] == ["gpt-4o"] - - def test_responses_with_retries_uses_the_dispatch_entrypoint(monkeypatch: pytest.MonkeyPatch) -> None: calls: Final[list[Mapping[str, object]]] = [] expected: Final = _response() @@ -310,3 +239,167 @@ def test_positional_parameters_remain_available_to_native_projection() -> None: assert request.resolved["include"] is include assert request.resolved["instructions"] == "Be brief" assert request.resolved["max_output_tokens"] == 16 + + +class _Agentic(CustomLogger): + async def async_should_run_agentic_loop(self, *args: object, **kwargs: object) -> tuple[bool, dict]: # pyright: ignore[reportIncompatibleMethodOverride, reportMissingTypeArgument] # test double only needs the override + return False, {} + + +class _Deployment(CustomLogger): + async def async_pre_call_deployment_hook(self, *args: object, **kwargs: object) -> dict[str, object] | None: # pyright: ignore[reportIncompatibleMethodOverride] # test double only needs the override + return None + + +ENCODED_ID: Final = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-1", "rs_1") # pyright: ignore[reportPrivateUsage] # the encoder litellm stamps on earlier turns + + +@pytest.mark.parametrize( + ("feature", "kwargs", "settings"), + ( + (Feature.MOCK_RESPONSE, {"mock_response": "hi"}, {}), + (Feature.PROXY_MCP_TOOLS, {"tools": [{"type": "mcp", "server_label": "g", "server_url": "litellm_proxy"}]}, {}), + (Feature.CALLER_CLIENT, {"client": object()}, {}), + pytest.param(Feature.CALLER_CLIENT, {"shared_session": object()}, {}, id="shared-session"), + pytest.param(Feature.EXTRA_BODY, {"extra_body": {"stream": True}}, {}, id="extra-body-stream"), + pytest.param( + Feature.PROVIDER_AFFINITY, + {"provider_affinity_header": "X-Session-ID", "session_id": "session-1"}, + {}, + id="provider-affinity", + ), + pytest.param( + Feature.INPUT_REWRITE, + {"input": [{"type": "reasoning", "id": "rs_1", "summary": list[object](), "status": None}]}, + {}, + id="input-rewrite", + ), + (Feature.TIMEOUT_VARIANTS, {"request_timeout": 5}, {}), + (Feature.AGENTIC_LOOP_HOOK, {}, {"callbacks": [_Agentic()]}), + (Feature.TLS_POLICY, {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), + (Feature.DROP_PARAMS, {}, {"drop_params": True}), + (Feature.STREAMING, {"stream": True}, {}), + (Feature.PROMPT_MANAGEMENT, {"prompt_id": "welcome"}, {}), + (Feature.TEXT_FORMAT, {"text_format": dict}, {}), + (Feature.CHAT_COMPLETIONS_BRIDGE, {"use_chat_completions_api": True}, {}), + (Feature.CHAT_COMPLETIONS_BRIDGE, {"model": "openai/chat_completions/gpt-5"}, {}), + (Feature.RETRIES_FALLBACKS, {"num_retries": 2}, {}), + (Feature.DEPLOYMENT_RESPONSE_ID, {"litellm_metadata": {"model_info": {"id": "deployment-1"}}}, {}), + (Feature.ENCODED_ITEM_IDS, {"input": [{"type": "reasoning", "id": ENCODED_ID, "summary": []}]}, {}), + (Feature.MODIFY_PARAMS, {}, {"modify_params": True}), + (Feature.RESPONSE_CACHE, {}, {"cache": Cache()}), + ( + Feature.MANAGED_RESOURCE_IDS, + { + "input": [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-file"}]}], + "model_file_id_mapping": {"litellm-file": {"deployment-1": "file-abc"}}, + "model_info": {"id": "deployment-1"}, + }, + {}, + ), + (Feature.CACHE_CONTROL_MARKS, {"tools": [{"type": "function", "name": "f", "cache_control": {}}]}, {}), + ( + Feature.ANTHROPIC_TOOL_SEARCH, + {"tools": [{"type": sorted(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES)[0], "name": "s"}]}, + {}, + ), + (Feature.TOOL_SCHEMA_REPAIR, {"tools": [{"type": "function", "name": "f", "parameters": "{}"}]}, {}), + (Feature.MODEL_PARAM_GATING, {"max_output_tokens": 1}, {}), + (Feature.CLIENT_HEADERS, {"headers": {"x-tenant": "a"}}, {}), + (Feature.DEPLOYMENT_HOOK, {}, {"callbacks": [_Deployment()]}), + ( + Feature.WORKLOAD_IDENTITY, + { + "openai_identity_provider_id": "idp-1", + "openai_service_account_id": "sa-1", + "openai_identity_token_file": "/var/run/token", + }, + {}, + ), + ), +) +def test_each_detected_feature_reaches_its_rule( + monkeypatch: pytest.MonkeyPatch, feature: Feature, kwargs: Mapping[str, object], settings: Mapping[str, object] +) -> None: + rules: Final = ( + RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY, features=frozenset({feature})), + RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED), + ) + python_result: Final = _response() + native_result: Final = _response() + + def served(call_kwargs: Mapping[str, object]) -> object: + return _DISPATCH.run( + (), + {"model": "gpt-5", "input": INPUT, **call_kwargs}, + python=lambda *call_args, **python_kwargs: python_result, + binding=responses_binding(lambda request: native_result), + native=lambda hook, request, call_args, native_kwargs: hook(request), + rules=rules, + ) + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "modify_params", False) + monkeypatch.setattr(litellm, "ssl_security_level", None) + monkeypatch.setattr(litellm, "cache", None) + for setting in ("api_key", "openai_key", "api_base", "num_retries"): + monkeypatch.setattr(litellm, setting, None) + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE", "OPENAI_API_KEY"): + monkeypatch.delenv(name, raising=False) + assert served({}) is native_result + for name, value in settings.items(): + monkeypatch.setattr(litellm, name, value) + assert served(kwargs) is python_result + + +@pytest.mark.asyncio +async def test_deployment_hook_keeps_only_sync_calls_on_python(monkeypatch: pytest.MonkeyPatch) -> None: + rules: Final = ( + RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY, features=frozenset({Feature.DEPLOYMENT_HOOK})), + RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED), + ) + native_result: Final = _response() + + async def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: stub + pytest.fail("the async handler leaves the deployment hook to the shared call wrapper") + + async def native(request: NativeCall) -> ResponsesAPIResponse: + return native_result + + monkeypatch.setattr(litellm, "callbacks", [_Deployment()]) + result: Final = await _ADISPATCH.arun( + (), + {"model": "gpt-5", "input": INPUT}, + python=python, + binding=aresponses_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=rules, + ) + assert result is native_result + + +@pytest.mark.parametrize( + ("kwargs", "native"), + ( + ({"model": "gpt-5"}, True), + ({"model": "openai/gpt-5"}, True), + ({"model": "gpt-5", "custom_llm_provider": "openai"}, True), + ({"model": "anthropic/claude-sonnet-4-5"}, False), + ({"model": "claude-sonnet-4-5"}, False), + ({"model": "gpt-5", "custom_llm_provider": "azure"}, False), + ), +) +def test_provider_rule_sees_the_resolved_provider(kwargs: Mapping[str, object], native: bool) -> None: + python_result: Final = _response() + native_result: Final = _response() + + result: Final = _DISPATCH.run( + (), + {"input": INPUT, **kwargs}, + python=lambda *call_args, **call_kwargs: python_result, + binding=responses_binding(lambda request: native_result), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=(RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED, providers=frozenset({"openai"})),), + ) + assert result is (native_result if native else python_result) diff --git a/tests/test_litellm_rust/responses/test_failures.py b/tests/test_litellm_rust/responses/test_failures.py index 6dad1bc7278..4c4bb96c6ac 100644 --- a/tests/test_litellm_rust/responses/test_failures.py +++ b/tests/test_litellm_rust/responses/test_failures.py @@ -56,28 +56,3 @@ 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( - "options", - ( - {"stream": True}, - {"extra_body": {"provider_option": True}}, - {"mock_response": "mock"}, - {"num_retries": 1}, - {"use_chat_completions_api": True}, - {"model_list": []}, - ), -) -async def test_unsupported_requests_fail_before_callbacks_and_provider_io( - server: RecordingServer, options: dict[str, object] -) -> None: - recorder: Final = RecordingLogger() - server.expected_requests = 0 - - with pytest.raises(RuntimeError, match="declined the request"): - await call(server, callbacks=[recorder], **options) - - assert not server.requests - assert not recorder.events diff --git a/tests/test_litellm_rust/support/routing.py b/tests/test_litellm_rust/support/routing.py index 4250589848d..e3bcd2ac463 100644 --- a/tests/test_litellm_rust/support/routing.py +++ b/tests/test_litellm_rust/support/routing.py @@ -11,7 +11,7 @@ from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer from tests.test_litellm_rust.support.served import assert_served_natively -__all__ = ["CALL_ID", "assert_served_natively", "call_public", "native_route"] +__all__ = ["CALL_ID", "assert_served_natively", "call_public", "native_route", "routed"] CALL_ID: Final = "callback-contract" @@ -22,6 +22,12 @@ def native_route(*routes: Route) -> Generator[None]: yield +@contextmanager +def routed(*rules: RouteRule) -> Generator[None]: + with rebound(catalog, "RULES", rules): + yield + + async def call_public( create: Callable[..., object], acreate: Callable[..., Awaitable[object]], diff --git a/tests/test_litellm_rust/test_dispatch.py b/tests/test_litellm_rust/test_dispatch.py index 9d53db26ba8..193866409a2 100644 --- a/tests/test_litellm_rust/test_dispatch.py +++ b/tests/test_litellm_rust/test_dispatch.py @@ -1,12 +1,11 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping from dataclasses import dataclass -from typing import Final, TypeAlias +from typing import Final, NoReturn, TypeAlias import pytest -from litellm.rust_bridge import configuration from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules +from litellm.rust_bridge.catalog import Feature, Route, RouteContext, RouteRule, Rules, Target from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.dispatch import PublicDispatch from litellm.rust_bridge.runtime import NO_PYTHON, NoPythonImplementationError @@ -70,29 +69,6 @@ def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None assert result is expected -def test_disabled_optional_rust_rule_forwards_before_projection() -> None: - rules: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),) - - def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: - pytest.fail("Disabled optional Rust must not project the request") - - dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: RouteContext(Route.OCR)) - expected: Final = object() - configuration.rust(False) - try: - result: Final = dispatch.run( - ("model",), - {}, - python=lambda *args, **kwargs: expected, - binding=binding(), - native=lambda hook, request, args, kwargs: pytest.fail("Disabled optional Rust must not call native"), - rules=rules, - ) - finally: - configuration.rust(None) - assert result is expected - - def test_native_stream_result_is_not_consumed_or_wrapped() -> None: request: Final = Request(model="streaming-model") stream: Final[Iterator[int]] = iter((1, 2)) @@ -207,36 +183,42 @@ def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() - assert result is expected -@pytest.mark.asyncio -async def test_async_bypass_forwards_to_python_without_native() -> None: - request: Final = Request(model="bypassed-model") - rules: Final[Rules] = (RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED),) - dispatch: Final = PublicDispatch( - route=Route.RESPONSES, - request=lambda args, kwargs: request, - context=lambda value: RouteContext(Route.RESPONSES, model=value.model), - bypass=lambda value: value.model == "bypassed-model", - ) +@pytest.mark.parametrize( + "constrained", + ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, providers=frozenset({"other"})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, models=frozenset({"other"})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, targets=frozenset({Target.CHAT_COMPLETIONS})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), + ), + ids=("providers", "models", "targets", "features"), +) +def test_constrained_python_rule_still_projects_for_a_later_rust_rule(constrained: RouteRule) -> None: expected: Final = object() - - async def python(*args: object, **kwargs: object) -> object: # kwargs-ok: public pass-through shape - return expected - - result: Final = await dispatch.arun( - ("bypassed-model",), - {}, - python=python, - binding=binding(), - native=lambda hook, value, args, kwargs: pytest.fail("Bypassed requests must not call native"), - rules=rules, + native_binding: Final[NativeBinding[Callable[[Request], object]]] = NativeBinding("x", validate=lambda _: None) + native_binding.override(lambda request: expected) + dispatch: Final = PublicDispatch( + route=Route.MESSAGES, + request=lambda args, kwargs: Request(model="model"), + context=lambda value: RouteContext(Route.MESSAGES, model=value.model, target=Target.MESSAGES), ) + + result: Final = dispatch.run( + ("model",), + {}, + python=lambda *args, **kwargs: pytest.fail("a constrained Python rule must not hide the Rust rule after it"), + binding=native_binding, + native=lambda hook, value, args, kwargs: hook(value), + rules=(constrained, RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED)), + ) + assert result is expected NativeRoute: TypeAlias = Callable[[Request, tuple[object, ...], Mapping[str, object]], object] -def native_route(result: object) -> NativeBinding[NativeRoute]: +def no_python_binding(result: object) -> NativeBinding[NativeRoute]: bound: Final[NativeBinding[NativeRoute]] = NativeBinding("no_python", validate=lambda _: None) bound.override(lambda request, args, kwargs: (result, request, args, dict(kwargs))) return bound @@ -263,52 +245,87 @@ async def dispatch_without_python( return await dispatch.arun(("model",), {"page": 1}, python=NO_PYTHON, binding=bound, native=native, rules=rules) -def ocr_dispatch(request: Request | None, *, bypass: bool = False) -> PublicDispatch[Request]: +def failing_context(request: Request) -> RouteContext: + raise RuntimeError("detector exploded") + + +def failing_dispatch() -> PublicDispatch[Request]: + return PublicDispatch( + route=Route.CHAT_COMPLETIONS, request=lambda args, kwargs: Request(model="model"), context=failing_context + ) + + +def reject_native(hook: object, request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NoReturn: + pytest.fail("a failed context build must not reach native") + + +def test_a_detector_error_keeps_the_call_on_python() -> None: + expected: Final = object() + result: Final = failing_dispatch().run( + ("model",), + {}, + python=lambda *args, **kwargs: expected, + binding=binding(), + native=reject_native, + rules=(RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),), + ) + assert result is expected + + +async def test_a_detector_error_keeps_the_async_call_on_python() -> None: + expected: Final = object() + + async def python(*args: object, **kwargs: object) -> object: # kwargs-ok: stub + return expected + + result: Final = await failing_dispatch().arun( + ("model",), + {}, + python=python, + binding=binding(), + native=reject_native, + rules=(RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),), + ) + assert result is expected + + +def ocr_dispatch(request: Request | None) -> PublicDispatch[Request]: return PublicDispatch( route=Route.OCR, request=lambda args, kwargs: request, context=lambda value: RouteContext(Route.OCR, model=value.model), - bypass=lambda _: bypass, ) @pytest.mark.parametrize("asynchronous", (False, True)) -@pytest.mark.parametrize("switch", (None, False)) -async def test_dispatch_without_python_hands_every_call_to_native(asynchronous: bool, switch: bool | None) -> None: +async def test_dispatch_without_python_hands_every_call_to_native(asynchronous: bool) -> None: request: Final = Request(model="model") result: Final = object() - configuration.rust(switch) - try: - dispatched: Final = await dispatch_without_python( - ocr_dispatch(request), - native_route(result), - (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), - asynchronous=asynchronous, - ) - finally: - configuration.rust(None) + dispatched: Final = await dispatch_without_python( + ocr_dispatch(request), + no_python_binding(result), + (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), + asynchronous=asynchronous, + ) assert dispatched == (result, request, ("model",), {"page": 1}) @pytest.mark.parametrize("asynchronous", (False, True)) @pytest.mark.parametrize( - ("request_value", "bypass", "rules", "reason"), + ("request_value", "rules", "reason"), ( - (Request(model="model"), False, (), "must resolve to RUST_REQUIRED"), - (Request(model="model"), False, (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),), "must resolve to RUST_REQUIRED"), - (None, False, (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), "must project to a native request"), - (Request(model="model"), True, (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), "bypass predicate matches"), + (Request(model="model"), (), "must resolve to RUST_REQUIRED"), + (Request(model="model"), (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),), "must resolve to RUST_REQUIRED"), + (None, (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), "must project to a native request"), ), - ids=("no-rule", "opt-out-rule", "unprojectable-call", "bypassed-call"), + ids=("no-rule", "opt-out-rule", "unprojectable-call"), ) async def test_dispatch_without_python_never_falls_back( - asynchronous: bool, request_value: Request | None, bypass: bool, rules: Rules, reason: str + asynchronous: bool, request_value: Request | None, rules: Rules, reason: str ) -> None: bound: Final[NativeBinding[NativeRoute]] = NativeBinding("no_python", validate=lambda _: None) bound.override(lambda request, args, kwargs: pytest.fail("a misdeclared route must not reach native")) with pytest.raises(NoPythonImplementationError, match=f"ocr has no Python implementation, so .*{reason}"): - await dispatch_without_python( - ocr_dispatch(request_value, bypass=bypass), bound, rules, asynchronous=asynchronous - ) + await dispatch_without_python(ocr_dispatch(request_value), bound, rules, asynchronous=asynchronous) diff --git a/tests/test_litellm_rust/test_public_entries.py b/tests/test_litellm_rust/test_public_entries.py new file mode 100644 index 00000000000..2f5c33e3ced --- /dev/null +++ b/tests/test_litellm_rust/test_public_entries.py @@ -0,0 +1,287 @@ +import inspect +from collections.abc import Awaitable, Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from enum import Enum, auto +from typing import Final, TypeAlias, cast # noqa: TID251 # public entries return union results + +import pytest + +import litellm +from litellm.messages import dispatch as messages +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.chat_completions.entrypoints import NATIVE_ACOMPLETION, NATIVE_COMPLETION +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.embeddings.entrypoints import NATIVE_AEMBEDDING, NATIVE_EMBEDDING +from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES +from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR +from litellm.rust_bridge.public_call import NativeCall +from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES, NATIVE_RESPONSES +from litellm.rust_bridge.runtime import NoPythonImplementationError +from litellm.types.utils import ModelResponse +from tests.test_litellm_rust.support.routing import routed + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + +Served: TypeAlias = Callable[[object], bool] + + +@dataclass(frozen=True, slots=True) +class Entry: + """``python`` is ``None`` for a route with no Python implementation, which only RUST_REQUIRED rules serve.""" + + route: Route + sync: Callable[..., object] + asynchronous: Callable[..., Awaitable[object]] + flag: str + sync_binding: NativeBinding[object] + async_binding: NativeBinding[object] + arguments: Mapping[str, object] + positional: int + python: Served | None + + def args(self) -> tuple[object, ...]: + return tuple(self.arguments.values())[: self.positional] + + def kwargs(self, **extra: object) -> Mapping[str, object]: # kwargs-ok: forwards the async flag + return {**dict(tuple(self.arguments.items())[self.positional :]), **extra} + + +def _binding(value: object) -> NativeBinding[object]: + return cast("NativeBinding[object]", value) # cast-ok: the table only overrides and resets bindings + + +def _chat(result: object) -> bool: + return isinstance(result, ModelResponse) and result.choices[0].message.content == "python" # pyright: ignore[reportAttributeAccessIssue] # non-streaming choice + + +def _embedding(result: object) -> bool: + return getattr(result, "data")[0]["embedding"] == [0.5] # noqa: B009 # union result has no common attribute + + +def _responses(result: object) -> bool: + return getattr(result, "output_text") == "python" # noqa: B009 # union result has no common attribute + + +def _messages(result: object) -> bool: + return isinstance(result, Mapping) and result["content"][0]["text"] == "python" + + +ENTRIES: Final = ( + Entry( + route=Route.CHAT_COMPLETIONS, + sync=litellm.completion, + asynchronous=litellm.acompletion, + flag="acompletion", + sync_binding=_binding(NATIVE_COMPLETION), + async_binding=_binding(NATIVE_ACOMPLETION), + arguments={"model": "openai/gpt-4o", "messages": MESSAGES, "mock_response": "python"}, + positional=2, + python=_chat, + ), + Entry( + route=Route.EMBEDDINGS, + sync=litellm.embedding, + asynchronous=litellm.aembedding, + flag="aembedding", + sync_binding=_binding(NATIVE_EMBEDDING), + async_binding=_binding(NATIVE_AEMBEDDING), + arguments={"model": "openai/text-embedding-3-small", "input": ["hi"], "mock_response": [0.5]}, + positional=2, + python=_embedding, + ), + Entry( + route=Route.RESPONSES, + sync=litellm.responses, + asynchronous=litellm.aresponses, + flag="aresponses", + sync_binding=_binding(NATIVE_RESPONSES), + async_binding=_binding(NATIVE_ARESPONSES), + arguments={"input": "hi", "model": "openai/gpt-4o", "mock_response": "python"}, + positional=2, + python=_responses, + ), + Entry( + route=Route.MESSAGES, + sync=messages.anthropic_messages_handler, + asynchronous=messages.anthropic_messages, + flag="is_async", + sync_binding=_binding(NATIVE_MESSAGES), + async_binding=_binding(NATIVE_AMESSAGES), + arguments={ + "max_tokens": 16, + "messages": MESSAGES, + "model": "anthropic/claude-sonnet-4-5", + "mock_response": "python", + }, + positional=0, + python=_messages, + ), + Entry( + route=Route.OCR, + sync=litellm.ocr, + asynchronous=litellm.aocr, + flag="aocr", + sync_binding=_binding(NATIVE_OCR), + async_binding=_binding(NATIVE_AOCR), + arguments={"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": b"pdf"}}, + positional=2, + python=None, + ), +) + + +class Form(Enum): + SYNC = auto() + SYNC_FLAG_FALSE = auto() + ASYNC = auto() + SYNC_FLAG_TRUE = auto() + ASYNC_FLAG_TRUE = auto() + + +async def call(entry: Entry, form: Form) -> object: + match form: + case Form.SYNC: + return entry.sync(*entry.args(), **entry.kwargs()) + case Form.SYNC_FLAG_FALSE: + return entry.sync(*entry.args(), **entry.kwargs(**{entry.flag: False})) + case Form.ASYNC: + return await entry.asynchronous(*entry.args(), **entry.kwargs()) + case Form.SYNC_FLAG_TRUE: + return await cast( + "Awaitable[object]", entry.sync(*entry.args(), **entry.kwargs(**{entry.flag: True})) + ) # cast-ok: a sync entry with its flag set returns a coroutine + case Form.ASYNC_FLAG_TRUE: + return await entry.asynchronous(*entry.args(), **entry.kwargs(**{entry.flag: True})) + + +@dataclass(slots=True) +class Native: + result: ModelResponse = field(default_factory=ModelResponse) + calls: list[tuple[str, NativeCall]] = field(default_factory=list) # mutable-ok: records native calls + + def install(self, entry: Entry) -> None: + def sync(native_call: NativeCall) -> ModelResponse: + self.calls.append(("sync", native_call)) + return self.result + + async def asynchronous(native_call: NativeCall) -> ModelResponse: + self.calls.append(("async", native_call)) + return self.result + + entry.sync_binding.override(sync) + entry.async_binding.override(asynchronous) + + +@pytest.fixture +def native(entry: Entry) -> Iterator[Native]: + recorder: Final = Native() + recorder.install(entry) + yield recorder + entry.sync_binding.reset() + entry.async_binding.reset() + + +class Outcome(Enum): + PYTHON = auto() + NATIVE = auto() + UNAVAILABLE = auto() + NO_PYTHON = auto() + + +@dataclass(frozen=True, slots=True) +class Scenario: + rollout: Rollout | None + binding: bool + outcome: Outcome + + +SCENARIOS: Final = { + "no-rule": Scenario(None, binding=True, outcome=Outcome.PYTHON), + "python-only": Scenario(Rollout.PYTHON_ONLY, binding=True, outcome=Outcome.PYTHON), + "required": Scenario(Rollout.RUST_REQUIRED, binding=True, outcome=Outcome.NATIVE), + "required-without-extension": Scenario(Rollout.RUST_REQUIRED, binding=False, outcome=Outcome.UNAVAILABLE), +} + + +def expected(entry: Entry, scenario: Scenario) -> Outcome: + if entry.python is None and scenario.rollout is not Rollout.RUST_REQUIRED: + return Outcome.NO_PYTHON + return scenario.outcome + + +@contextmanager +def scenario_rules(entry: Entry, scenario: Scenario) -> Iterator[None]: + if not scenario.binding: + entry.sync_binding.override(None) + entry.async_binding.override(None) + with routed(*(() if scenario.rollout is None else (RouteRule(entry.route, scenario.rollout),))): + yield + + +def entry_id(entry: Entry) -> str: + return entry.route.value + + +@pytest.mark.parametrize("form", tuple(Form), ids=lambda form: form.name.lower()) +@pytest.mark.parametrize("scenario", tuple(SCENARIOS.values()), ids=tuple(SCENARIOS)) +@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id) +async def test_every_public_form_takes_the_path_the_rules_pick( + entry: Entry, scenario: Scenario, form: Form, native: Native +) -> None: + outcome: Final = expected(entry, scenario) + with scenario_rules(entry, scenario): + if outcome is Outcome.UNAVAILABLE: + with pytest.raises(RuntimeError, match=f"Rust {entry.route.value} bridge is unavailable"): + await call(entry, form) + return + if outcome is Outcome.NO_PYTHON: + with pytest.raises(NoPythonImplementationError, match=f"{entry.route.value} has no Python implementation"): + await call(entry, form) + assert native.calls == [] + return + result: Final = await call(entry, form) + + if outcome is Outcome.PYTHON: + assert entry.python is not None + assert entry.python(result) + assert native.calls == [] + return + assert result is native.result + [(binding, native_call)] = native.calls + assert binding == ("sync" if form in (Form.SYNC, Form.SYNC_FLAG_FALSE) else "async") + assert {name: native_call.resolved[name] for name in entry.arguments} == entry.arguments + assert entry.flag not in native_call.resolved + + +@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id) +async def test_async_forms_hand_native_the_same_call(entry: Entry, native: Native) -> None: + with scenario_rules(entry, SCENARIOS["required"]): + for form in (Form.ASYNC, Form.SYNC_FLAG_TRUE, Form.ASYNC_FLAG_TRUE): + _ = await call(entry, form) + + assert [binding for binding, _ in native.calls] == ["async", "async", "async"] + resolved: Final = [dict(native_call.resolved) for _, native_call in native.calls] + assert resolved[1:] == [resolved[0], resolved[0]] + + +@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id) +async def test_sync_entry_with_flag_true_does_nothing_until_awaited(entry: Entry, native: Native) -> None: + with scenario_rules(entry, SCENARIOS["required"]): + pending: Final = entry.sync(*entry.args(), **entry.kwargs(**{entry.flag: True})) + + assert inspect.iscoroutine(pending) + assert native.calls == [] + assert await pending is native.result + assert len(native.calls) == 1 + + +@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id) +async def test_sync_entry_with_flag_true_raises_setup_errors_at_await(entry: Entry, native: Native) -> None: + with scenario_rules(entry, SCENARIOS["required-without-extension"]): + pending: Final = entry.sync(*entry.args(), **entry.kwargs(**{entry.flag: True})) + + assert inspect.iscoroutine(pending) + with pytest.raises(RuntimeError, match="is unavailable"): + await pending diff --git a/tests/unit/caching/test_in_memory_cache.py b/tests/unit/caching/test_in_memory_cache.py index 214f4ecf301..75232476ecd 100644 --- a/tests/unit/caching/test_in_memory_cache.py +++ b/tests/unit/caching/test_in_memory_cache.py @@ -49,6 +49,47 @@ async def test_async_increment_delegates_to_locked_sync_path(): assert cache.get_cache("counter") == 5 +class _AnotherReaderExpiresTheEntryFirstCache(InMemoryCache): + other_reader_started = False + other_reader_result: object = "not run" + + def evict_element_if_expired(self, key: str) -> bool: + if not self.other_reader_started: + self.other_reader_started = True + self.other_reader_result = self.get_cache(key) + return super().evict_element_if_expired(key) + + +def test_get_cache_reads_an_entry_another_reader_just_expired_as_a_miss() -> None: + cache: Final = _AnotherReaderExpiresTheEntryFirstCache() + cache.set_cache("tokenizer_failure", "HF unavailable", ttl=-1) + + assert cache.get_cache("tokenizer_failure") is None + assert cache.other_reader_result is None + + +class _ClockThatLetsAnotherReaderRunFirst: + cache: InMemoryCache | None = None + other_reader_result: object = "not run" + + def __call__(self) -> float: + cache: Final = self.cache + if cache is not None: + self.cache = None + self.other_reader_result = cache.get_cache("tokenizer_failure") + return 100.0 + + +def test_get_cache_reads_an_entry_another_reader_expires_during_the_expiry_check_as_a_miss() -> None: + clock: Final = _ClockThatLetsAnotherReaderRunFirst() + cache: Final = InMemoryCache(clock=clock) + cache.set_cache("tokenizer_failure", "HF unavailable", ttl=-1) + clock.cache = cache + + assert cache.get_cache("tokenizer_failure") is None + assert clock.other_reader_result is None + + def test_in_memory_openai_obj_cache(): from openai import OpenAI diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index c96e3e2c031..e64f5cca5b8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -252,6 +252,7 @@ def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> It @pytest.fixture(autouse=True) def isolate_litellm_globals() -> Iterator[None]: original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS} + original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks) original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS} original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)} original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__} @@ -263,6 +264,7 @@ def isolate_litellm_globals() -> Iterator[None]: _flush_client_caches() for name in CALLBACK_LISTS: setattr(litellm, name, []) + litellm_router_module._routers_holding_global_callbacks.clear() for name in RESET_TO_NONE_GLOBALS: setattr(litellm, name, None) for name in MODULE_LEVEL_CLIENTS: @@ -277,6 +279,8 @@ def isolate_litellm_globals() -> Iterator[None]: litellm.__dict__.update(original_clients) for name, value in (original_callbacks | original_reset | original_restored).items(): setattr(litellm, name, value) + litellm_router_module._routers_holding_global_callbacks.clear() + litellm_router_module._routers_holding_global_callbacks.update(original_callback_holders) for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items(): logger.setLevel(level) logger.disabled = disabled diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 6ea76a0a8a7..b4485bcc5f1 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -34,6 +34,7 @@ from mcp.types import ( InitializeResult, JSONRPCError, JSONRPCMessage, + JSONRPCNotification, JSONRPCRequest, JSONRPCResponse, LoggingMessageNotificationParams, @@ -61,8 +62,10 @@ from litellm.experimental_mcp_client.client import ( strip_auth_scheme, ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + ServerListFault, classify_list_exception, list_fault_http_status, + upstream_rate_limit, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _format_byok_openapi_auth_header, @@ -305,6 +308,36 @@ class _MockTransportClient(MCPClient): return streamable_http_client(self.server_url, http_client=http_client), http_client +def _post_status_responder(failure_method: str, status_code: int) -> Callable[[httpx2.Request], httpx2.Response]: + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method != "POST": + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + assert isinstance(payload, (JSONRPCRequest, JSONRPCNotification)) + if payload.method == "initialize": + assert isinstance(payload, JSONRPCRequest) + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_HANDSHAKE_VERSION, + "capabilities": {"tools": {}}, + "serverInfo": {"name": "status", "version": "1"}, + }, + }, + ) + if payload.method == failure_method: + return httpx2.Response(status_code, headers={"Retry-After": "7"}) + if payload.method == "tools/list": + assert isinstance(payload, JSONRPCRequest) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {"tools": []}}) + return httpx2.Response(202) + + return respond + + class _FakeExceptionGroup(Exception): """Duck-typed stand-in for an anyio/builtin ExceptionGroup. @@ -1467,8 +1500,8 @@ def test_mcp_extra_matches_proxy_extra_and_supports_streamable_http(): } specifier: Final = Requirement(mcp_extra["mcp"]).specifier - assert not specifier.contains("1.28.1") - assert specifier.contains("2.2.0") + assert not specifier.contains("2.2.0") + assert specifier.contains("2.3.0") with (pyproject_path.parent / "uv.lock").open("rb") as f: locked = tomllib.load(f) mcp_versions: Final = [package["version"] for package in locked["package"] if package["name"] == "mcp"] @@ -1862,6 +1895,32 @@ async def test_http_response_handler_preserves_success_and_http_errors(status_co assert caught.value.response.status_code == status_code +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure_method,status_code", + [ + ("notifications/initialized", 429), + ("notifications/initialized", 503), + ("tools/list", 429), + ], +) +async def test_list_tools_classifies_upstream_status_for_post_failures(failure_method: str, status_code: int) -> None: + client: Final = _MockTransportClient( + _post_status_responder(failure_method, status_code), + server_url="https://example.com/mcp", + ) + with pytest.raises(httpx2.HTTPStatusError) as caught: + await client.list_tools(raise_on_error=True) + + assert classify_list_exception(caught.value) == ServerListFault(tag="upstream_error", status_code=status_code) + rate_limit: Final = upstream_rate_limit(caught.value) + if status_code == 429: + assert rate_limit is not None + assert rate_limit.retry_after_seconds == 7 + else: + assert rate_limit is None + + @pytest.mark.asyncio async def test_http_status_check_allows_auth_refresh_before_rejecting() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( @@ -4388,6 +4447,35 @@ async def test_persistent_tool_call_preserves_upstream_http_status_error() -> No await asyncio.wait_for(session.wait_closed(), 5) +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [429, 503]) +async def test_persistent_tool_call_classifies_initialized_notification_status(status_code: int) -> None: + client: Final = _MockTransportClient( + _post_status_responder("notifications/initialized", status_code), + server_url="https://example.com/mcp", + ) + session: Final = client.open_persistent_session() + try: + with pytest.raises(httpx2.HTTPStatusError) as caught: + await client.call_tool( + CallToolRequestParams(name="protected_tool", arguments={}), + raise_on_error=True, + persistent_session=session, + ) + + assert classify_list_exception(caught.value) == ServerListFault(tag="upstream_error", status_code=status_code) + rate_limit: Final = upstream_rate_limit(caught.value) + if status_code == 429: + assert rate_limit is not None + assert rate_limit.retry_after_seconds == 7 + else: + assert rate_limit is None + finally: + if not session.closed: + session.close() + await asyncio.wait_for(session.wait_closed(), 5) + + @pytest.mark.asyncio async def test_persistent_session_remaps_stream_error_over_mcp_error() -> None: app: Final = _stateful_upstream() diff --git a/tests/unit/extraction/__init__.py b/tests/unit/extraction/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/extraction/test_main.py b/tests/unit/extraction/test_main.py new file mode 100644 index 00000000000..584ed9cbb91 --- /dev/null +++ b/tests/unit/extraction/test_main.py @@ -0,0 +1,422 @@ +from __future__ import annotations + +import asyncio +import json +import re +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.extraction import ExtractionResponse + +_CLAIRIC_URL: Final = "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/@cf/cloudflare/clairic" +_STATE: Final = "Invoice #4471 from Acme Tools, dated 2026-03-03. 12 hammers. Total $174. Payment due in 30 days." +_FIELDS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType( + { + "vendor": {"type": "text", "instructions": "Who sent the invoice?"}, + "invoice_date": {"type": "date"}, + "total": {"type": "number", "unit": "USD"}, + "items": {"type": "list", "items": "text"}, + "days_to_pay": {"type": "estimate", "min": 0, "max": 90, "provider_field": "kept"}, + } +) +_INPUT_TOKENS: Final = 416 +_RESULT: Final[Mapping[str, object]] = { + "model": "clairic", + "answers": { + "vendor": { + "type": "text", + "p_null": 0.0177, + "span": [19, 29], + "guess": "Acme Tools", + "confidence": 0.9998, + "value": "Acme Tools", + }, + "invoice_date": {"type": "date", "p_null": 6.4e-07, "confidence": 0.99, "value": "2026-03-03"}, + "total": {"type": "number", "p_null": 4.2e-07, "guess": 174, "confidence": 0.99, "value": 174}, + "items": {"type": "list", "value": ["12", "hammers"], "spans": [[51, 53], [54, 61]], "confidence": 0.07}, + "days_to_pay": { + "type": "estimate", + "value": 31, + "mean": 30.76, + "median": 30, + "p10": 29, + "p90": 34, + "p_below": 1.2e-06, + "p_above": 1.2e-06, + }, + }, + "truncated": False, + "usage": { + "state_tokens": 52, + "image_tokens": 0, + "schema_tokens": 312, + "input_tokens": _INPUT_TOKENS, + "output_tokens": 0, + "passes": 2, + }, +} +_ENVELOPE: Final[Mapping[str, object]] = {"result": _RESULT, "success": True, "errors": [], "messages": []} +_PNG_DATA_URL: Final = "data:image/png;base64,AA==" + + +@pytest.fixture(autouse=True) +def _httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.fixture +def cloudflare_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + monkeypatch.setenv("CLOUDFLARE_API_KEY", "cloudflare-key") + monkeypatch.delenv("CLOUDFLARE_API_BASE", raising=False) + + +def _assert_invoice_answers(response: ExtractionResponse) -> None: + assert response.answers["vendor"].value == "Acme Tools" + assert response.answers["vendor"].span == (19, 29) + assert response.answers["items"].spans == ((51, 53), (54, 61)) + assert response.answers["days_to_pay"].median == 30 + assert response.usage.input_tokens == _INPUT_TOKENS + assert response.usage.schema_tokens == 312 + assert response.truncated is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ("cloudflare/clairic", "cloudflare/@cf/cloudflare/clairic")) +@pytest.mark.usefixtures("cloudflare_env") +async def test_clairic_forwards_the_caller_body_and_unwraps_the_envelope( + respx_mock: respx.MockRouter, model: str +) -> None: + route: Final = respx_mock.post(_CLAIRIC_URL).respond(json=_ENVELOPE) + + response: Final = await litellm.aextraction( + model=model, + state=_STATE, + fields=_FIELDS, + images=[_PNG_DATA_URL, {"content_type": "image/jpeg", "base64": "AA=="}], + reference_date="2026-03-04", + options={"include_quantiles": True}, + ) + + assert route.call_count == 1 + request: Final = route.calls[0].request + assert request.headers["authorization"] == "Bearer cloudflare-key" + assert json.loads(request.content) == { + "model": "clairic", + "state": _STATE, + "fields": _FIELDS, + "images": [_PNG_DATA_URL, {"content_type": "image/jpeg", "base64": "AA=="}], + "reference_date": "2026-03-04", + "options": {"include_quantiles": True}, + } + _assert_invoice_answers(response) + assert response.model == model + assert response.hidden_params["model"] == "cloudflare/@cf/cloudflare/clairic" + assert response.model_dump(mode="json", exclude={"model"}) == { + key: value for key, value in _RESULT.items() if key != "model" + } + + +@pytest.mark.usefixtures("cloudflare_env") +def test_sync_extraction_reads_an_unwrapped_result(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_CLAIRIC_URL).respond(json=_RESULT) + + response: Final = litellm.extraction(model="cloudflare/clairic", state=_STATE, fields=_FIELDS) + + assert route.call_count == 1 + assert json.loads(route.calls[0].request.content) == {"model": "clairic", "state": _STATE, "fields": _FIELDS} + _assert_invoice_answers(response) + + +def _clairic_transport() -> httpx.MockTransport: + def respond(request: httpx.Request) -> httpx.Response: + assert str(request.url) == _CLAIRIC_URL + return httpx.Response(200, json=_ENVELOPE) + + return httpx.MockTransport(respond) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("cloudflare_env") +async def test_async_extraction_sends_through_the_callers_client(respx_mock: respx.MockRouter) -> None: + response: Final = await litellm.aextraction( + model="cloudflare/clairic", + state=_STATE, + fields=_FIELDS, + client=AsyncHTTPHandler(transport=_clairic_transport()), + ) + + assert len(respx_mock.calls) == 0 + _assert_invoice_answers(response) + + +@pytest.mark.usefixtures("cloudflare_env") +def test_sync_extraction_sends_through_the_callers_client(respx_mock: respx.MockRouter) -> None: + response: Final = litellm.extraction( + model="cloudflare/clairic", + state=_STATE, + fields=_FIELDS, + client=HTTPHandler(client=httpx.Client(transport=_clairic_transport())), + ) + + assert len(respx_mock.calls) == 0 + _assert_invoice_answers(response) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("cloudflare_env") +async def test_a_cloudflare_bad_input_is_a_bad_request_with_the_cloudflare_message( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_CLAIRIC_URL).respond( + status_code=400, + json={ + "errors": [{"message": "AiError: Bad input: fields.total.decimals must be 0 to 6", "code": 5006}], + "success": False, + "result": {}, + "messages": [], + }, + ) + + with pytest.raises(litellm.BadRequestError, match=re.escape("fields.total.decimals must be 0 to 6")): + await litellm.aextraction( + model="cloudflare/clairic", state=_STATE, fields={"total": {"type": "number", "decimals": 9}} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("fields", "images"), + ( + ({}, None), + ({"total": {"type": "money"}}, None), + ({"bad name": {"type": "text"}}, None), + ({"total": {"type": "number"}}, ["https://example.com/invoice.png"]), + ({"total": {"type": "number"}}, [_PNG_DATA_URL] * 5), + ), + ids=("no_fields", "unknown_type", "invalid_field_name", "remote_image", "too_many_images"), +) +@pytest.mark.usefixtures("cloudflare_env") +async def test_an_invalid_request_is_refused_before_http( + respx_mock: respx.MockRouter, fields: Mapping[str, object], images: list[str] | None +) -> None: + with pytest.raises(litellm.BadRequestError, match="Invalid Extraction request"): + await litellm.aextraction(model="cloudflare/clairic", state=_STATE, fields=fields, images=images) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("unset", "message"), + ( + ("CLOUDFLARE_ACCOUNT_ID", "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID"), + ("CLOUDFLARE_API_KEY", "Missing API key for Extraction provider 'cloudflare'"), + ), +) +@pytest.mark.usefixtures("cloudflare_env") +async def test_missing_cloudflare_credentials_are_refused_before_http( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, unset: str, message: str +) -> None: + monkeypatch.delenv(unset) + + with pytest.raises(litellm.BadRequestError, match=message): + await litellm.aextraction(model="cloudflare/clairic", state=_STATE, fields=_FIELDS) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +async def test_a_provider_without_extraction_is_refused(respx_mock: respx.MockRouter) -> None: + with pytest.raises(litellm.BadRequestError, match=r"Unknown Extraction provider 'openai'\. Supported.*cloudflare"): + await litellm.aextraction(model="openai/gpt-5.5", state=_STATE, fields=_FIELDS, api_key="k") + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ("cloudflare/clairic", "cloudflare/@cf/cloudflare/clairic")) +@pytest.mark.usefixtures("cloudflare_env") +async def test_clairic_cost_bills_input_tokens_at_the_cost_map_price(respx_mock: respx.MockRouter, model: str) -> None: + respx_mock.post(_CLAIRIC_URL).respond(json=_ENVELOPE) + + response: Final = await litellm.aextraction(model=model, state=_STATE, fields=_FIELDS) + + expected_cost: Final = _INPUT_TOKENS * float(litellm.model_cost[model]["input_cost_per_token"]) + assert expected_cost > 0 + assert litellm.completion_cost(completion_response=response) == pytest.approx(expected_cost) + assert response.hidden_params["response_cost"] == pytest.approx(expected_cost) + + +@pytest.mark.asyncio +async def test_router_extraction_sends_the_deployment_model_and_credentials(respx_mock: respx.MockRouter) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "invoice-reader", + "litellm_params": { + "model": "cloudflare/clairic", + "api_key": "deployment-key", + "api_base": "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1", + "additional_drop_params": ["temperature"], + }, + } + ] + ) + route: Final = respx_mock.post(_CLAIRIC_URL).respond(json=_ENVELOPE) + + response: Final = await router.aextraction(model="invoice-reader", state=_STATE, fields=_FIELDS) + + assert route.call_count == 1 + assert route.calls[0].request.headers["authorization"] == "Bearer deployment-key" + assert json.loads(route.calls[0].request.content) == {"model": "clairic", "state": _STATE, "fields": _FIELDS} + assert response.model == "cloudflare/clairic" + _assert_invoice_answers(response) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("cloudflare_env") +async def test_litellm_params_stay_local(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_CLAIRIC_URL).respond(json=_ENVELOPE) + + await litellm.aextraction( + model="cloudflare/clairic", + state=_STATE, + fields=_FIELDS, + metadata={"team": "billing"}, + num_retries=0, + stream=False, + additional_drop_params=["language"], + ) + + assert json.loads(route.calls[0].request.content) == {"model": "clairic", "state": _STATE, "fields": _FIELDS} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("name", ("language", "priority", "session_id", "mock_delay", "extra_body")) +@pytest.mark.usefixtures("cloudflare_env") +async def test_a_key_clairic_does_not_take_is_refused_before_http_without_drop_params( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, name: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape(f"does not support parameters: ['{name}']")): + await litellm.aextraction(model="cloudflare/clairic", state=_STATE, fields=_FIELDS, **{name: "x"}) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("cloudflare_env") +async def test_additional_drop_params_drops_only_the_keys_it_names( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = respx_mock.post(_CLAIRIC_URL).respond(json=_ENVELOPE) + + await litellm.aextraction( + model="cloudflare/clairic", + state=_STATE, + fields=_FIELDS, + temperature=0.2, + language="en", + additional_drop_params=["temperature", "language"], + ) + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("does not support parameters: ['seed']")): + await litellm.aextraction( + model="cloudflare/clairic", state=_STATE, fields=_FIELDS, seed=7, additional_drop_params=["temperature"] + ) + + assert route.call_count == 1 + assert json.loads(route.calls[0].request.content) == {"model": "clairic", "state": _STATE, "fields": _FIELDS} + + +class _RecordingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.standard_logging_object: Mapping[str, object] | None = None + self.stream: object = None + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.stream = kwargs.get("stream") + standard_logging_object: Final = kwargs.get("standard_logging_object") + if isinstance(standard_logging_object, dict): + self.standard_logging_object = standard_logging_object + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("cloudflare_env") +async def test_stream_is_refused_before_http_without_drop_params(respx_mock: respx.MockRouter) -> None: + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape("does not support parameters: ['stream']")): + await litellm.aextraction(model="cloudflare/clairic", state=_STATE, fields=_FIELDS, stream=True) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("name", "value"), + (("temperature", 0.2), ("user", "end-user-1"), ("seed", 7), ("response_format", {"type": "json_object"})), +) +@pytest.mark.usefixtures("cloudflare_env") +async def test_openai_generation_params_are_refused_before_http_without_drop_params( + respx_mock: respx.MockRouter, name: str, value: object +) -> None: + with pytest.raises(litellm.UnsupportedParamsError, match=re.escape(f"does not support parameters: ['{name}']")): + await litellm.aextraction(model="cloudflare/clairic", state=_STATE, fields=_FIELDS, **{name: value}) + + assert len(respx_mock.calls) == 0 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("cloudflare_env") +async def test_openai_generation_params_are_dropped_with_drop_params(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_CLAIRIC_URL).respond(json=_ENVELOPE) + + await litellm.aextraction( + model="cloudflare/clairic", + state=_STATE, + fields=_FIELDS, + temperature=0.2, + user="end-user-1", + seed=7, + drop_params=True, + ) + + assert json.loads(route.calls[0].request.content) == {"model": "clairic", "state": _STATE, "fields": _FIELDS} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("cloudflare_env") +async def test_a_dropped_stream_is_sent_and_logged_as_a_billed_non_stream_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + route: Final = respx_mock.post(_CLAIRIC_URL).respond(json=_ENVELOPE) + recording_logger: Final = _RecordingLogger() + monkeypatch.setattr(litellm, "callbacks", [recording_logger]) + + response: Final = await litellm.aextraction( + model="cloudflare/clairic", state=_STATE, fields=_FIELDS, stream=True, drop_params=True + ) + await asyncio.sleep(0) + GLOBAL_LOGGING_WORKER.start() + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert "stream" not in json.loads(route.calls[0].request.content) + _assert_invoice_answers(response) + logged: Final = recording_logger.standard_logging_object + assert logged is not None + assert recording_logger.stream is False + assert response.hidden_params["response_cost"] > 0 + assert logged["response_cost"] == pytest.approx(response.hidden_params["response_cost"]) diff --git a/tests/unit/fine_tuning/test_main.py b/tests/unit/fine_tuning/test_main.py index 57c845c6e1a..89f51d2e74a 100644 --- a/tests/unit/fine_tuning/test_main.py +++ b/tests/unit/fine_tuning/test_main.py @@ -1,11 +1,12 @@ import json -from typing import Final +from typing import Final, Literal import httpx import pytest import respx import litellm +from litellm.types.utils import LiteLLMFineTuningJob API_BASE: Final = "https://fine-tuning.test" BASE_MODEL: Final = "gpt-4.1-nano-2025-04-14" @@ -52,3 +53,42 @@ async def test_create_sends_the_model_without_the_provider_prefix( ) assert json.loads(sent.calls.last.request.content)["model"] == BASE_MODEL + + +@pytest.mark.asyncio +@respx.mock +@pytest.mark.parametrize( + ("custom_llm_provider", "api_base", "route", "sent_api_version"), + [ + ("openai", f"{API_BASE}/v1", f"{API_BASE}/v1/fine_tuning/jobs", None), + ("azure", API_BASE, f"{API_BASE}/openai/fine_tuning/jobs", "2024-10-21"), + ], +) +@pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"]) +async def test_create_with_api_version_reaches_the_provider_and_only_azure_gets_it( + monkeypatch: pytest.MonkeyPatch, + custom_llm_provider: Literal["openai", "azure"], + api_base: str, + route: str, + sent_api_version: str | None, + is_async: bool, +) -> None: + sent: Final = respx.post(url__startswith=route).mock(return_value=httpx.Response(200, json=PROVIDER_JOB)) + request: Final = { + "model": BASE_MODEL, + "training_file": "file-abc123", + "custom_llm_provider": custom_llm_provider, + "api_key": "sk-test", + "api_base": api_base, + "api_version": "2024-10-21", + } + + async with httpx.AsyncClient() as session: + monkeypatch.setattr(litellm, "aclient_session", session) + job: Final = ( + await litellm.acreate_fine_tuning_job(**request) if is_async else litellm.create_fine_tuning_job(**request) + ) + + assert isinstance(job, LiteLLMFineTuningJob) + assert job.id == PROVIDER_JOB["id"] + assert sent.calls.last.request.url.params.get("api-version") == sent_api_version diff --git a/tests/unit/harness/test_endpoint.py b/tests/unit/harness/test_endpoint.py index 19bd6d0450e..d1447fd40cc 100644 --- a/tests/unit/harness/test_endpoint.py +++ b/tests/unit/harness/test_endpoint.py @@ -272,6 +272,7 @@ async def test_sdk_chat_stream(monkeypatch: pytest.MonkeyPatch) -> None: assert lines[-1] == "data: [DONE]" assert json.loads(lines[0][6:])["choices"][0]["delta"]["content"] == "he" assert calls[0]["stream_options"] == {"include_usage": True} + assert calls[0]["_litellm_stream_usage_injected"] is True assert (ep.usage.input_tokens, ep.usage.output_tokens) == (6, 2) assert ep.usage.cost == pytest.approx(6 * 0.001 + 2 * 0.002) diff --git a/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json index cfdf707c977..6e5461c3b60 100644 --- a/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json +++ b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null, \"kubernetes_pod_routing\": null}", + "metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"billed_retry_attempt\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null, \"kubernetes_pod_routing\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/unit/integrations/test_responses_background_cost.py b/tests/unit/integrations/test_responses_background_cost.py index 1bf95e35c4d..7d266c94700 100644 --- a/tests/unit/integrations/test_responses_background_cost.py +++ b/tests/unit/integrations/test_responses_background_cost.py @@ -282,6 +282,12 @@ class TestCheckResponsesCost: client = MagicMock() client.db = MagicMock() client.db.litellm_managedobjecttable = MagicMock() + client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + client.db.litellm_deprecatedverificationtoken.find_unique = AsyncMock(return_value=None) + client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None) + client.db.litellm_deletedteamtable.find_first = AsyncMock(return_value=None) return client @pytest.fixture diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py index 00b0e2f6972..6d9d8c7997a 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -17,6 +17,7 @@ import pytest import litellm import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod import litellm.proxy.common_request_processing as common_request_processing_mod +from litellm.cost_calculator import select_model_name_for_cost_calc from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ResponseMetadata, @@ -649,3 +650,38 @@ def test_update_response_metadata_prices_per_second_deployment_from_its_stamped_ assert result._response_ms == pytest.approx(2000) assert result._hidden_params["response_cost"] == pytest.approx(0.02 * 2) + + +def test_update_response_metadata_keeps_the_provider_model_after_the_proxy_rewrites_response_model(): + start_time: Final = datetime.datetime(2026, 9, 21, 12, 0, 0) + logging_obj: Final = Logging( + model="azure/ptu-deployment", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=start_time, + litellm_call_id="provider-model-response-metadata", + function_id="f", + ) + logging_obj.update_environment_variables( + model="azure/ptu-deployment", litellm_params={}, optional_params={}, custom_llm_provider="azure" + ) + result: Final = ModelResponse( + model="gpt-5.4-nano-2026-03-17", + usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18), + ) + + update_response_metadata( + result=result, + logging_obj=logging_obj, + model="azure/ptu-deployment", + kwargs={}, + start_time=start_time, + end_time=start_time + datetime.timedelta(seconds=2), + ) + result.model = "model-group-alias" + + assert ( + select_model_name_for_cost_calc(model=None, completion_response=result, custom_llm_provider="azure") + == "azure/gpt-5.4-nano-2026-03-17" + ) 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 251bfa624a5..89173956dc6 100644 --- a/tests/unit/litellm_core_utils/test_get_litellm_params.py +++ b/tests/unit/litellm_core_utils/test_get_litellm_params.py @@ -317,6 +317,10 @@ class TestGetLitellmParamsExplicitFields: result = get_litellm_params(no_log=True) assert result["no-log"] is True + def test_github_copilot_host_is_registered_as_a_litellm_param(self): + result: Final = get_litellm_params(github_copilot_host="octo.ghe.com") + assert result["github_copilot_host"] == "octo.ghe.com" + class TestGetLitellmParamsDataResidency: """Verify that data_residency is inferred from OpenAI regional api_base.""" 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 c23749f7a20..b584a72855a 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -195,9 +195,7 @@ def test_update_model_params_with_health_check_tracking_information(): }, } - result = HealthCheckHelpers.update_model_params_with_health_check_tracking_information( - initial_model_params - ) + result = HealthCheckHelpers.update_model_params_with_health_check_tracking_information(initial_model_params) # Verify that litellm_metadata was added assert "litellm_metadata" in result @@ -753,9 +751,7 @@ async def test_ahealth_check_keeps_the_noul_probe_for_other_decisions_providers( assert "error" not in result, result assert upstream.called sent: Final = json.loads(upstream.calls[0].request.content) - assert sent["questions"] == { - "reachable": {"type": "noul", "instructions": "Is the service reachable?"} - } + assert sent["questions"] == {"reachable": {"type": "noul", "instructions": "Is the service reachable?"}} @pytest.mark.parametrize( @@ -916,6 +912,7 @@ async def test_ahealth_check_without_mode_reports_the_real_failure( assert "Missing `mode`" not in result["error"] assert "raw_request_typed_dict" in result + def test_update_litellm_params_for_health_check(): """ Test if update_litellm_params_for_health_check correctly: @@ -1010,12 +1007,10 @@ def test_update_litellm_params_for_health_check(): "model": f"bedrock/{prefix}anthropic.claude-3-haiku-20240307-v1:0", "api_key": "fake_key", } - updated_params = update_litellm_params_for_health_check( - model_info, litellm_params + updated_params = update_litellm_params_for_health_check(model_info, litellm_params) + assert updated_params["model"] == f"{prefix}anthropic.claude-3-haiku-20240307-v1:0", ( + f"Failed to preserve CRIS prefix: {prefix}" ) - assert ( - updated_params["model"] == f"{prefix}anthropic.claude-3-haiku-20240307-v1:0" - ), f"Failed to preserve CRIS prefix: {prefix}" litellm_params = { "model": "bedrock/us-east-2/us.anthropic.claude-3-haiku-20240307-v1:0", @@ -1036,30 +1031,21 @@ def test_update_litellm_params_for_health_check(): "api_key": "fake_key", } updated_params = update_litellm_params_for_health_check(model_info, litellm_params) - assert ( - updated_params["model"] - == "llama/arn:aws:bedrock:us-east-1:123:imported-model/abc" - ) + assert updated_params["model"] == "llama/arn:aws:bedrock:us-east-1:123:imported-model/abc" litellm_params = { "model": "bedrock/deepseek_r1/arn:aws:bedrock:us-west-2:456:imported-model/xyz", "api_key": "fake_key", } updated_params = update_litellm_params_for_health_check(model_info, litellm_params) - assert ( - updated_params["model"] - == "deepseek_r1/arn:aws:bedrock:us-west-2:456:imported-model/xyz" - ) + assert updated_params["model"] == "deepseek_r1/arn:aws:bedrock:us-west-2:456:imported-model/xyz" litellm_params = { "model": "bedrock/converse/us.anthropic.claude-haiku-4-5-20251001-v1:0", "api_key": "fake_key", } updated_params = update_litellm_params_for_health_check(model_info, litellm_params) - assert ( - updated_params["model"] - == "converse/us.anthropic.claude-haiku-4-5-20251001-v1:0" - ) + assert updated_params["model"] == "converse/us.anthropic.claude-haiku-4-5-20251001-v1:0" litellm_params = { "model": "bedrock/invoke/us-west-2/anthropic.claude-instant-v1", @@ -1073,29 +1059,22 @@ def test_update_litellm_params_for_health_check(): "api_key": "fake_key", } updated_params = update_litellm_params_for_health_check(model_info, litellm_params) - assert ( - updated_params["model"] - == "arn:aws:bedrock:eu-central-1:000:application-inference-profile/abc" - ) + assert updated_params["model"] == "arn:aws:bedrock:eu-central-1:000:application-inference-profile/abc" litellm_params = { "model": "bedrock/us-west-2/llama/arn:aws:bedrock:us-east-1:123:imported-model/abc", "api_key": "fake_key", } updated_params = update_litellm_params_for_health_check(model_info, litellm_params) - assert ( - updated_params["model"] - == "llama/arn:aws:bedrock:us-east-1:123:imported-model/abc" - ) + assert updated_params["model"] == "llama/arn:aws:bedrock:us-east-1:123:imported-model/abc" litellm_params = { "model": "bedrock/converse/us-west-2/eu.anthropic.claude-3-sonnet-20240229-v1:0", "api_key": "fake_key", } updated_params = update_litellm_params_for_health_check(model_info, litellm_params) - assert ( - updated_params["model"] == "converse/eu.anthropic.claude-3-sonnet-20240229-v1:0" - ) + assert updated_params["model"] == "converse/eu.anthropic.claude-3-sonnet-20240229-v1:0" + @pytest.mark.asyncio async def test_perform_health_check_filters_by_model_id(): @@ -1142,6 +1121,7 @@ async def test_perform_health_check_filters_by_model_id(): assert len(healthy_endpoints) == 1 assert healthy_endpoints[0]["api_key"] == "fake-key-2" + @pytest.mark.asyncio async def test_perform_health_check_skip_disabled_background_models(): from litellm.proxy.health_check import perform_health_check @@ -1180,6 +1160,7 @@ async def test_perform_health_check_skip_disabled_background_models(): assert len(captured[0]) == 1 assert captured[0][0]["model_name"] == "a" + @pytest.mark.asyncio async def test_perform_health_check_with_health_check_model(): """ @@ -1206,9 +1187,7 @@ async def test_perform_health_check_with_health_check_model(): return {"status": "healthy"} with patch("litellm.ahealth_check", side_effect=mock_health_check): - healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check( - model_list - ) + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list) print("health check calls: ", health_check_calls) assert health_check_calls[0] == "openai/gpt-5-mini" @@ -1217,6 +1196,7 @@ async def test_perform_health_check_with_health_check_model(): assert len(healthy_endpoints) == 1 assert len(unhealthy_endpoints) == 0 + @pytest.mark.asyncio async def test_image_generation_health_check_prompt(monkeypatch): """Health checks should respect default and environment-configured prompts.""" @@ -1267,9 +1247,7 @@ async def test_image_generation_health_check_prompt(monkeypatch): health_check_calls = await run_health_check(reloaded_health_check) assert len(health_check_calls) == 1 - assert ( - health_check_calls[0]["prompt"] == reloaded_constants.DEFAULT_HEALTH_CHECK_PROMPT - ) + assert health_check_calls[0]["prompt"] == reloaded_constants.DEFAULT_HEALTH_CHECK_PROMPT # Environment override should change the prompt without code changes override_prompt = "environment override prompt" @@ -1293,9 +1271,7 @@ async def test_health_check_with_custom_llm_provider( "object": "chat.completion", "created": 0, "model": "deepseek-r1-distill-qwen-1.5B-q4", - "choices": [ - {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} - ], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, } ) @@ -1503,3 +1479,37 @@ async def test_audio_transcription_health_check_returns_transcribed_text( assert "error" not in response assert upstream.called assert b'name="file"' in upstream.calls.last.request.content + + +@pytest.mark.parametrize("model", ("cloudflare/clairic", "cloudflare/@cf/cloudflare/clairic")) +async def test_ahealth_check_probes_extraction_models_through_the_extraction_api( + model: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct-1") + monkeypatch.delenv("CLOUDFLARE_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post( + "https://api.cloudflare.com/client/v4/accounts/acct-1/ai/run/@cf/cloudflare/clairic" + ).respond( + json={ + "result": { + "model": "clairic", + "answers": {"reachable": {"type": "text", "value": "health"}}, + "truncated": False, + "usage": {"input_tokens": 20, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + } + ) + + result: Final = await ahealth_check({"model": model, "api_key": "sk-test"}, mode=None) + + assert "error" not in result, result + sent: Final = json.loads(upstream.calls[0].request.content) + assert sent["state"] == "health check" + assert sent["fields"]["reachable"]["type"] == "text" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 981c43a9792..43190852c60 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -42,6 +42,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( + BilledRoundsSpend, Logging, StandardLoggingPayloadSetup, _extract_response_obj_and_hidden_params, @@ -402,6 +403,32 @@ def test_get_combined_callback_list_preserves_insertion_order(logging_obj): ) == ["prometheus", "langfuse", "datadog", "otel", "s3", "gcs_bucket", "arize", "logfire"] +class _InputEventRecorder(CustomLogger): + def __init__(self, label: str, events: list[str]) -> None: + self._label = label + self.events = events + + def log_pre_api_call(self, model: object, messages: object, kwargs: object) -> None: + self.events.append(f"{self._label} pre") + + def log_post_api_call(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.events.append(f"{self._label} post") + + +def test_input_events_reach_the_replacement_and_not_the_global_callback_it_replaced( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final[list[str]] = [] + monkeypatch.setattr(litellm, "input_callback", [_InputEventRecorder("global", events)]) + logging_obj.replace_global_callback(_InputEventRecorder("replacement", events)) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + logging_obj.pre_call(input="", api_key="") + logging_obj.post_call(original_response="ok", input="", api_key="") + + assert events == ["replacement pre", "replacement post"] + + def test_get_masked_api_base(logging_obj): api_base = "https://api.openai.com/v1" masked_api_base = logging_obj._get_masked_api_base(api_base) @@ -3088,7 +3115,7 @@ async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_c prefer_async_handlers=True, ) - mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception) + mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception, for_async_request=True) @pytest.mark.asyncio @@ -6811,6 +6838,26 @@ def test_failure_handler_zeroes_spend_without_recovered_usage(logging_obj): assert payload["total_tokens"] == 0 +def test_failure_after_earlier_rounds_bills_only_the_failed_round_when_it_priced_no_breakdown( + logging_obj: LitellmLogging, +) -> None: + earlier_breakdown: Final = CostBreakdown(total_cost=0.3, margin_fixed_amount=0.1) + logging_obj.cost_breakdown = earlier_breakdown + logging_obj.model_call_details["litellm_params"] = {} + + logging_obj.record_failure_after_earlier_rounds( + BilledRoundsSpend(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.3, earlier_breakdown), + Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), + 0.0, + ) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "anthropic", "claude"), "tb") + + row: Final = logging_obj.model_call_details["standard_logging_object"] + assert (row["prompt_tokens"], row["completion_tokens"]) == (7, 0) + assert row["response_cost"] == pytest.approx(0.0) + assert row.get("cost_breakdown") is None or row["cost_breakdown"]["total_cost"] == pytest.approx(0.0) + + def test_set_cost_breakdown_stores_reasoning_cost(): """reasoning_cost is stored only when positive, mirroring the cache-cost fields.""" from datetime import datetime @@ -8188,6 +8235,155 @@ def test_cache_callback_alone_does_not_start_sync_callbacks_for_async_requests( assert logging_obj._should_run_sync_callbacks_for_async_calls() is expected +class _FailureHookRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.hooks: list[str] = [] + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.hooks.append("sync_failure") + + +class _BothFailureHooksRecorder(_FailureHookRecorder): + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.hooks.append("async_failure") + + +def _failure_hook_count_logging_obj( + call_type: str, recorder: _FailureHookRecorder, *, async_failure_listed: bool +) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini-tts", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type=call_type, + start_time=datetime.datetime.now(), + litellm_call_id=f"failure-hook-count-{call_type}", + function_id=f"failure-hook-count-{call_type}", + dynamic_failure_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder] if async_failure_listed else [], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + return logging_obj + + +@pytest.fixture +def no_global_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, name, []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("recorder_type", "async_listed", "expected"), + [ + (_BothFailureHooksRecorder, True, ["async_failure"]), + (_FailureHookRecorder, True, ["sync_failure"]), + (_BothFailureHooksRecorder, False, ["sync_failure"]), + ], +) +@pytest.mark.parametrize("call_type", ["aspeech", "atranscription", "aimage_edit", "asearch"]) +@pytest.mark.usefixtures("no_global_callbacks") +async def test_async_call_outside_the_async_flag_list_gets_one_custom_logger_failure_hook( + call_type: str, recorder_type: type[_FailureHookRecorder], async_listed: bool, expected: list[str] +) -> None: + recorder: Final = recorder_type() + logging_obj: Final = _failure_hook_count_logging_obj(call_type, recorder, async_failure_listed=async_listed) + + logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True) + await logging_obj.async_failure_handler(ValueError("upstream rejected"), "traceback") + + assert recorder.hooks == expected + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_logger_in_the_global_failure_callback_lists_gets_one_failure_hook( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder: Final = _BothFailureHooksRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini-tts", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="aspeech", + start_time=datetime.datetime.now(), + litellm_call_id="global-failure", + function_id="global-failure", + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True) + await logging_obj.async_failure_handler(ValueError("upstream rejected"), "traceback") + + assert recorder.hooks == ["async_failure"] + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_failure_handler_run_without_an_async_failure_run_keeps_the_sync_hook() -> None: + recorder: Final = _BothFailureHooksRecorder() + + _failure_hook_count_logging_obj("aspeech", recorder, async_failure_listed=True).failure_handler( + ValueError("upstream rejected"), "traceback" + ) + + assert recorder.hooks == ["sync_failure"] + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_async_call_outside_the_async_flag_list_still_runs_sync_failure_function_callbacks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm_logging, "customLogger", CustomLogger()) + calls: Final[list[str]] = [] + + def failure_function(kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + calls.append("function_failure") + + recorder: Final = _BothFailureHooksRecorder() + logging_obj: Final = _failure_hook_count_logging_obj("aspeech", recorder, async_failure_listed=True) + logging_obj.dynamic_failure_callbacks = [recorder, failure_function] + + logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True) + + assert calls == ["function_failure"] + assert recorder.hooks == [] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_client_wrapped_failure_gives_a_logger_in_litellm_callbacks_one_hook( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder: Final = _BothFailureHooksRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + + @litellm.utils.client + async def aspeech(model: str, input: str, voice: str, **kwargs: object) -> None: + raise litellm.BadRequestError(message="upstream rejected", model=model, llm_provider="openai") + + @litellm.utils.client + def speech(model: str, input: str, voice: str, **kwargs: object) -> None: + raise litellm.BadRequestError(message="upstream rejected", model=model, llm_provider="openai") + + with pytest.raises(litellm.BadRequestError): + await aspeech(model="gpt-4o-mini-tts", input="hi", voice="alloy") + with pytest.raises(litellm.BadRequestError): + speech(model="gpt-4o-mini-tts", input="hi", voice="alloy") + + assert recorder.hooks == ["async_failure", "sync_failure"] + + @pytest.mark.asyncio @pytest.mark.parametrize("turn_off_message_logging", [False, True]) @pytest.mark.parametrize("make_response", [ModelResponse, _responses_api_response_with_text]) @@ -10019,6 +10215,7 @@ class TestAzurePTUSpilloverCost: ROUTER_MODEL_ID: Final = "ptu-spill-router-model-id" SERVED_MODEL: Final = "azure/spill-served-model-ptu" + PTU_DEPLOYMENT: Final = "azure/spill-ptu-deployment" PTU_MODEL_INFO: Final = { "id": ROUTER_MODEL_ID, "team_id": "team-1", @@ -10054,7 +10251,9 @@ class TestAzurePTUSpilloverCost: litellm.model_cost.pop(TestAzurePTUSpilloverCost.ROUTER_MODEL_ID, None) litellm.model_cost.pop(TestAzurePTUSpilloverCost.SERVED_MODEL, None) - def _logging_obj(self, model_info: dict, *, flag: str, litellm_rate: float, monkeypatch) -> LitellmLogging: + def _logging_obj( + self, model_info: dict, *, flag: str, litellm_rate: float, monkeypatch, deployment: str | None = None + ) -> LitellmLogging: monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", flag) obj = LitellmLogging( model=self.SERVED_MODEL, @@ -10071,7 +10270,7 @@ class TestAzurePTUSpilloverCost: optional_params={}, litellm_params={ "api_base": "", - "metadata": {"model_info": model_info}, + "metadata": {"model_info": model_info, "deployment": deployment}, "input_cost_per_token": litellm_rate, "output_cost_per_token": litellm_rate, }, @@ -10175,6 +10374,50 @@ class TestAzurePTUSpilloverCost: finally: self._unregister_models() + def test_spillover_without_base_model_reports_the_pay_as_you_go_rates_billing_used(self, monkeypatch) -> None: + self._register_models() + litellm.model_cost[self.PTU_DEPLOYMENT] = {"litellm_provider": "azure", "mode": "chat"} + try: + obj = self._logging_obj( + dict(self.PTU_MODEL_INFO), + flag="True", + litellm_rate=0.0, + monkeypatch=monkeypatch, + deployment=self.PTU_DEPLOYMENT, + ) + response = self._response() + response._hidden_params["additional_headers"] = {"llm_provider-x-ms-is-spilled-over": "true"} + response._hidden_params["provider_response_model"] = response.model + response_cost: Final = obj.response_cost_calculator(result=response) + response.model = "spill-model-group-alias" + + payload: Final = get_standard_logging_object_payload( + kwargs={ + "model": self.SERVED_MODEL, + "custom_llm_provider": "azure", + "litellm_params": obj.litellm_params, + "response_cost": response_cost, + }, + init_response_obj=response, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + logging_obj=obj, + status="success", + ) + + assert response_cost == pytest.approx(self.EXPECTED_SPILL_COST) + assert payload is not None + model_map_information: Final = payload["model_map_information"] + assert model_map_information["model_map_key"] != self.PTU_DEPLOYMENT + model_map_value: Final = model_map_information["model_map_value"] + assert model_map_value is not None + assert response_cost == pytest.approx( + 100 * model_map_value["input_cost_per_token"] + 50 * model_map_value["output_cost_per_token"] + ) + finally: + litellm.model_cost.pop(self.PTU_DEPLOYMENT, None) + self._unregister_models() + def test_spillover_header_does_not_touch_non_ptu_custom_pricing(self, monkeypatch) -> None: self._register_models() custom_model_id: Final = "non-ptu-custom-router-model-id" diff --git a/tests/unit/litellm_core_utils/test_logging_callback_manager.py b/tests/unit/litellm_core_utils/test_logging_callback_manager.py index f8260702068..25a52137d65 100644 --- a/tests/unit/litellm_core_utils/test_logging_callback_manager.py +++ b/tests/unit/litellm_core_utils/test_logging_callback_manager.py @@ -145,6 +145,7 @@ def test_remove_callback_from_all_lists(): obj = TestLogger() manager.add_litellm_callback(obj) + manager.add_litellm_input_callback(obj) manager.add_litellm_success_callback(obj) manager.add_litellm_failure_callback(obj) manager.add_litellm_async_success_callback(obj) @@ -153,6 +154,7 @@ def test_remove_callback_from_all_lists(): manager.remove_callback_from_all_lists(obj) assert obj not in litellm.callbacks + assert obj not in litellm.input_callback assert obj not in litellm.success_callback assert obj not in litellm.failure_callback assert obj not in litellm._async_success_callback diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 68aabc3c01e..550ecd3c0f7 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -603,6 +603,91 @@ async def test_streaming_handler_with_usage( assert chunk_has_usage +def _usage_test_logging( + caller_stream_options: dict[str, bool] | None, stream_usage_injected: bool = False +) -> Logging: + logging_obj: Final = Logging( + model="hosted_vllm/m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=1_790_000_000.0, + litellm_call_id="12345", + function_id="1245", + kwargs={"_litellm_stream_usage_injected": True} if stream_usage_injected else {}, + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="hosted_vllm/m", stream_options=caller_stream_options + ) + return logging_obj + + +@pytest.mark.parametrize( + ("caller_stream_options", "stream_usage_injected", "usage_chunk_choice_count"), + [ + ({"include_usage": True}, False, 0), + (None, False, 1), + ({"include_usage": True}, True, 1), + ], + ids=["caller-asked", "provider-handler-injected", "proxy-injected"], +) +@pytest.mark.parametrize("plain_iterator", [False, True]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_usage_chunk_has_empty_choices_only_for_callers_that_asked( + sync_mode: bool, + plain_iterator: bool, + caller_stream_options: dict[str, bool] | None, + stream_usage_injected: bool, + usage_chunk_choice_count: int, +): + usage: Final = Usage(prompt_tokens=5, completion_tokens=1, total_tokens=6) + provider_chunks: Final = [ + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="ok"))]), + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), + ModelResponseStream(id="c1", choices=[], provider_specific_fields={"citations": ["a"]}), + ModelResponseStream(id="c1", choices=[], usage=usage), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=( + iter(provider_chunks) if plain_iterator else ModelResponseListIterator(model_responses=provider_chunks) + ), + model="hosted_vllm/m", + custom_llm_provider="hosted_vllm", + logging_obj=_usage_test_logging(caller_stream_options, stream_usage_injected), + stream_options={"include_usage": True}, + ) + + chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + + assert [len(chunk.choices) for chunk in chunks] == [1, 1, 0, usage_chunk_choice_count] + assert chunks[2].provider_specific_fields == {"citations": ["a"]} + assert [getattr(chunk, "usage", None) is not None for chunk in chunks] == [False, False, False, True] + assert chunks[-1].usage.total_tokens == 6 + + +@pytest.mark.asyncio +async def test_plain_iterator_stream_keeps_usage_on_its_chunk_when_usage_was_not_requested(): + usage: Final = Usage(prompt_tokens=5, completion_tokens=1, total_tokens=6) + provider_chunks: Final = [ + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="ok"))]), + ModelResponseStream( + id="c1", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], usage=usage + ), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(provider_chunks), + model="hosted_vllm/m", + custom_llm_provider="hosted_vllm", + logging_obj=_usage_test_logging(None), + ) + + chunks: Final = [chunk async for chunk in wrapper] + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"] + assert chunks[-1].usage.total_tokens == 6 + + @pytest.mark.parametrize("sync_mode", [False]) @pytest.mark.asyncio @pytest.mark.flaky(reruns=3) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 309cbd33131..194d7995735 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -14,7 +14,7 @@ from concurrent.futures import Future, ThreadPoolExecutor, wait from itertools import accumulate, chain from pathlib import Path from typing import Final -from unittest.mock import MagicMock +from unittest.mock import MagicMock, sentinel import anyio.to_thread import pytest @@ -27,7 +27,7 @@ import litellm from litellm import decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +from litellm.constants import HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, TOKEN_COUNTER_MAX_CONCURRENT_COUNTS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( _encoding_count, @@ -940,10 +940,17 @@ def test_token_counter(): import unittest -from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding +from litellm.utils import ( + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _load_huggingface_tokenizer, + _select_tokenizer_helper, + claude_json_str, + encoding, +) # Clear the cache at module load to ensure clean state _load_huggingface_tokenizer.cache_clear() +_HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() class TestTokenizerSelection(unittest.TestCase): @@ -955,6 +962,7 @@ class TestTokenizerSelection(unittest.TestCase): --dist=loadscope (tests from same file run on same worker). """ _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() @patch("litellm.utils.tokenizer_dispatch.from_pretrained") def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): @@ -1038,6 +1046,69 @@ class TestTokenizerSelection(unittest.TestCase): finally: monkeypatch.undo() + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") + def test_failed_huggingface_tokenizer_download_is_not_retried_on_every_count( + self, mock_from_pretrained: MagicMock + ): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + try: + with patch("litellm.utils.time.monotonic", return_value=0.0): + mock_from_pretrained.side_effect = RuntimeError("HF unavailable") + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + text: Final = "Count the tokens for this message." + expected: Final = len(encoding.encode(text, disallowed_special=())) + + first_count: Final = litellm.token_counter(model=model, text=text) + first_attempts: Final = mock_from_pretrained.call_count + second_count: Final = litellm.token_counter(model=model, text=text) + + assert first_attempts > 0 + assert mock_from_pretrained.call_count == first_attempts + assert first_count == expected + assert second_count == expected + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + + @patch("litellm.utils.time.monotonic", return_value=299.0) + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") + def test_failed_huggingface_tokenizer_download_is_retried_after_retry_window( + self, mock_from_pretrained: MagicMock, mock_monotonic: MagicMock + ): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + try: + def from_pretrained_side_effect(_: str) -> object: + if mock_from_pretrained.call_count == 1: + mock_monotonic.return_value = 322.0 + raise RuntimeError("HF unavailable") + return sentinel.tokenizer + + mock_from_pretrained.side_effect = from_pretrained_side_effect + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + + first_result: Final = _select_tokenizer_helper(model) + assert first_result["type"] == "openai_tokenizer" + mock_monotonic.return_value = 323.0 + second_result: Final = _select_tokenizer_helper(model) + assert second_result["type"] == "openai_tokenizer", "tokenizer was retried at t=323" + assert mock_from_pretrained.call_count == 1 + + mock_monotonic.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS - 1 + third_result: Final = _select_tokenizer_helper(model) + assert third_result["type"] == "openai_tokenizer" + assert mock_from_pretrained.call_count == 1 + + mock_monotonic.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1 + retried_result: Final = _select_tokenizer_helper(model) + assert retried_result["type"] == "huggingface_tokenizer" + assert retried_result["tokenizer"] is sentinel.tokenizer + assert mock_from_pretrained.call_count == 2 + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _CODE_SAMPLE: Final = " for item in items:\n total += item.price * item.quantity\n" * 20 @@ -1104,6 +1175,7 @@ def test_disable_hf_tokenizer_download_keeps_deepseek_v4_on_cl100k(monkeypatch: def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() start: Final = threading.Barrier(8) def first_tokenizer(_: int) -> object: @@ -1116,6 +1188,44 @@ def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): assert len({id(tokenizer) for tokenizer in tokenizers}) == 1 +def test_remembered_tokenizer_failure_warns_once_instead_of_on_every_count(caplog: pytest.LogCaptureFixture): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + try: + with ( + patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable")), + caplog.at_level("WARNING", logger="LiteLLM"), + ): + types: Final = tuple(_select_tokenizer_helper(model)["type"] for _ in range(3)) + fallback_warnings: Final = [r for r in caplog.records if "token counts may be approximate" in r.getMessage()] + assert types == ("openai_tokenizer",) * 3 + assert len(fallback_warnings) == 1 + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + + +def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered(): + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted + + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + forked: Final = ForkedAfterNativeRuntimeStarted("native runtime started before fork") + try: + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=forked): + with pytest.raises(ForkedAfterNativeRuntimeStarted): + _select_tokenizer_helper(model) + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", return_value=sentinel.tokenizer): + loaded: Final = _select_tokenizer_helper(model) + assert loaded["type"] == "huggingface_tokenizer" + assert loaded["tokenizer"] is sentinel.tokenizer + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + + def test_token_counter_with_anthropic_tool_use(): """ Test that _count_anthropic_content() correctly handles tool_use blocks. diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 354c91920fc..61e797d7166 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -406,16 +406,26 @@ def test_huggingface_encoding_exposes_the_tokenizers_lookup_and_mutation_surface @pytest.mark.parametrize("log_level", ["WARNING", "ERROR"]) def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeypatch, log_level): from unittest.mock import patch - from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper + + from litellm.utils import ( + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _load_huggingface_tokenizer, + _select_tokenizer_helper, + ) monkeypatch.setenv("LITELLM_RUST", "false") _load_huggingface_tokenizer.cache_clear() - with caplog.at_level(log_level, logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": None}): - result = _select_tokenizer_helper("llama-2") - assert result["type"] == "openai_tokenizer" - assert result["tokenizer"].encode("hello") - assert ("token counts may be approximate" in caplog.text) is (log_level == "WARNING") - assert ("install tokenizers" in caplog.text) is (log_level == "WARNING") + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + try: + with caplog.at_level(log_level, logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": None}): + result = _select_tokenizer_helper("llama-2") + assert result["type"] == "openai_tokenizer" + assert result["tokenizer"].encode("hello") + assert ("token counts may be approximate" in caplog.text) is (log_level == "WARNING") + assert ("install tokenizers" in caplog.text) is (log_level == "WARNING") + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) @@ -471,10 +481,16 @@ def test_added_token_return_annotation_resolves_without_optional_import(python_i def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): from types import SimpleNamespace from unittest.mock import patch - from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper + + from litellm.utils import ( + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _load_huggingface_tokenizer, + _select_tokenizer_helper, + ) monkeypatch.setenv("LITELLM_RUST", "false") _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() model = "llama-2\r\nforged-model\x1b[31m\u2028\u2029" secret = "sk-" + "x" * 48 failure = OSError("download failed\r\nforged-error\x1b[31m api_key=" + secret) @@ -483,18 +499,26 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): raise failure dependency = SimpleNamespace(Tokenizer=SimpleNamespace(from_pretrained=fail_download)) - with caplog.at_level("WARNING", logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": dependency}): - result = _select_tokenizer_helper(model) - assert result["type"] == "openai_tokenizer" - assert result["tokenizer"].encode("hello") - message = next(record.getMessage() for record in caplog.records if "token counts may be approximate" in record.getMessage()) - assert "llama-2" in message - assert "download failed" in message - assert "forged-model" in message and "forged-error" in message - assert message.isascii() and message.isprintable() - assert secret not in message - assert "REDACTED" in message - assert "install tokenizers and huggingface-hub" in message + try: + with caplog.at_level("WARNING", logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": dependency}): + result = _select_tokenizer_helper(model) + assert result["type"] == "openai_tokenizer" + assert result["tokenizer"].encode("hello") + message = next( + record.getMessage() + for record in caplog.records + if "token counts may be approximate" in record.getMessage() + ) + assert "llama-2" in message + assert "download failed" in message + assert "forged-model" in message and "forged-error" in message + assert message.isascii() and message.isprintable() + assert secret not in message + assert "REDACTED" in message + assert "install tokenizers and huggingface-hub" in message + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 41ad27d6f40..b3bd1b8743a 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -477,6 +477,7 @@ def test_default_max_retries_env_var_reaches_azure_sdk_client(): "asearch", "adecisions", "asystemone", + "aextraction", "avector_store_create", "avector_store_search", "acreate_skill", diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index da5c4699925..44c0e82c006 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -5,8 +5,9 @@ import struct import sys import types from binascii import crc32 -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime +from types import MappingProxyType from typing import Final from unittest.mock import patch @@ -922,7 +923,11 @@ INVOKE_BODY = { async def _sent_passthrough_request( - action: str, body: dict[str, object], identity: dict[str, object], **kwargs: object + action: str, + body: dict[str, object], + identity: dict[str, object], + deployment_params: Mapping[str, object] = MappingProxyType({}), + **kwargs: object, ) -> httpx.Request: requests: Final[list[httpx.Request]] = [] @@ -939,6 +944,7 @@ async def _sent_passthrough_request( "aws_access_key_id": "AKIAROUTERKEY", "aws_secret_access_key": "router-secret", "aws_region_name": "us-west-2", + **deployment_params, }, } ] @@ -1189,3 +1195,176 @@ async def test_passthrough_body_is_sent_as_is_when_params_are_dropped_and_forwar assert json.loads(sent.content) == caller_body assert REQUEST_METADATA_HEADER not in sent.headers + + +DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "deployment-guardrail", + "guardrailVersion": "3", + "trace": "enabled", + "streamProcessingMode": "sync", +} +DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS: Final[Mapping[str, str]] = { + "x-amzn-bedrock-guardrailidentifier": "deployment-guardrail", + "x-amzn-bedrock-guardrailversion": "3", + "x-amzn-bedrock-trace": "ENABLED", +} +CALLER_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "caller-guardrail", + "guardrailVersion": "9", + "trace": "disabled", +} +CALLER_INVOKE_GUARDRAIL_HEADERS: Final[Mapping[str, str]] = { + "x-pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + "x-pass-X-Amzn-Bedrock-GuardrailVersion": "9", + "x-pass-X-Amzn-Bedrock-Trace": "DISABLED", +} + + +def _invoke_guardrail_headers(request: httpx.Request) -> dict[str, str]: + return { + name: value + for name, value in request.headers.items() + if name.lower().startswith("x-amzn-bedrock-guardrail") or name.lower() == "x-amzn-bedrock-trace" + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "sent_guardrail"), + [ + ("converse", {"guardrailIdentifier": "deployment-guardrail", "guardrailVersion": "3", "trace": "enabled"}), + ("converse-stream", DEPLOYMENT_GUARDRAIL), + ], +) +async def test_converse_passthrough_sends_the_deployment_guardrail_in_the_body( + action: str, sent_guardrail: Mapping[str, str], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + request: Final = await _sent_passthrough_request( + action, CONVERSE_BODY, {}, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert json.loads(request.content) == {**CONVERSE_BODY, "guardrailConfig": sent_guardrail} + assert _invoke_guardrail_headers(request) == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_sends_the_deployment_guardrail_as_signed_headers( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + request: Final = await _sent_passthrough_request( + action, INVOKE_BODY, PROXY_IDENTITY, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert _invoke_guardrail_headers(request) == DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS + assert {*DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS, REQUEST_METADATA_HEADER} <= _signed_header_names(request) + assert json.loads(request.headers[REQUEST_METADATA_HEADER]) == EXPECTED_REQUEST_METADATA + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "caller_guardrail"), + [("converse", CALLER_GUARDRAIL), ("converse-stream", {**CALLER_GUARDRAIL, "streamProcessingMode": "async"})], +) +async def test_deployment_guardrail_replaces_the_caller_converse_guardrail( + action: str, caller_guardrail: Mapping[str, str], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + request: Final = await _sent_passthrough_request( + action, + {**CONVERSE_BODY, "guardrailConfig": caller_guardrail}, + {}, + deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + ) + + assert json.loads(request.content) == { + **CONVERSE_BODY, + "guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_deployment_guardrail_replaces_every_invoke_guardrail_header_the_caller_sent( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + {}, + deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS, + ) + + assert _invoke_guardrail_headers(request) == {"x-amzn-bedrock-guardrailidentifier": "deployment-guardrail"} + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +async def test_passthrough_without_a_deployment_guardrail_sends_the_caller_guardrail_as_is( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + caller_body: Final = {**CONVERSE_BODY, "guardrailConfig": CALLER_GUARDRAIL} + + converse: Final = await _sent_passthrough_request("converse", caller_body, {}) + invoke: Final = await _sent_passthrough_request( + "invoke", INVOKE_BODY, {}, request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS + ) + + assert json.loads(converse.content) == caller_body + assert _invoke_guardrail_headers(invoke) == { + "x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-amzn-bedrock-guardrailversion": "9", + "x-amzn-bedrock-trace": "DISABLED", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +@pytest.mark.parametrize("guardrail_config", ["deployment-guardrail", {"guardrailVersion": "3"}]) +async def test_invoke_passthrough_refuses_a_deployment_guardrail_it_cannot_send( + action: str, guardrail_config: object, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + with pytest.raises(BedrockError, match="guardrailConfig") as rejected: + await _sent_passthrough_request( + action, INVOKE_BODY, {}, deployment_params={"guardrailConfig": guardrail_config} + ) + + assert rejected.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_converse_passthrough_forwards_a_non_object_deployment_guardrail_unchanged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + request: Final = await _sent_passthrough_request( + "converse", CONVERSE_BODY, {}, deployment_params={"guardrailConfig": "deployment-guardrail"} + ) + + assert json.loads(request.content) == {**CONVERSE_BODY, "guardrailConfig": "deployment-guardrail"} + + +@pytest.mark.asyncio +async def test_passthrough_leaves_actions_without_a_guardrail_field_as_sent(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + count_tokens_body: Final = {"input": {"converse": CONVERSE_BODY}} + + request: Final = await _sent_passthrough_request( + "count-tokens", count_tokens_body, {}, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert json.loads(request.content) == count_tokens_body + assert _invoke_guardrail_headers(request) == {} diff --git a/tests/unit/llms/cloudflare/test_cloudflare_extraction_guardrail_handler.py b/tests/unit/llms/cloudflare/test_cloudflare_extraction_guardrail_handler.py new file mode 100644 index 00000000000..1db105f4495 --- /dev/null +++ b/tests/unit/llms/cloudflare/test_cloudflare_extraction_guardrail_handler.py @@ -0,0 +1,335 @@ +from typing import Final, Literal + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms import get_guardrail_translation_mapping +from litellm.llms.cloudflare.extraction.guardrail_translation.handler import ExtractionHandler +from litellm.types.extraction import ExtractionResponse +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs + + +class _TaggingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="tagging") + self.seen: tuple[str, tuple[str, ...]] | None = None + self.seen_images: tuple[str, ...] = () + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = tuple(inputs.get("texts", [])) + self.seen = (input_type, texts) + self.seen_images = tuple(inputs.get("images", [])) + return GenericGuardrailAPIInputs(texts=[f"{text} [GUARDRAILED]" for text in texts]) + + +class _MaskingGuardrail(CustomGuardrail): + def __init__(self, secret: str, mask: str) -> None: + super().__init__(guardrail_name="masking") + self.secret = secret + self.mask = mask + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs( + texts=[text.replace(self.secret, self.mask) for text in inputs.get("texts", [])] + ) + + +@pytest.mark.parametrize("call_type", (CallTypes.extraction, CallTypes.aextraction)) +def test_extraction_call_types_resolve_to_the_extraction_handler(call_type: CallTypes) -> None: + assert get_guardrail_translation_mapping(call_type) is ExtractionHandler + + +@pytest.mark.asyncio +async def test_request_state_and_field_instructions_are_guardrailed_on_a_copy() -> None: + guardrail: Final = _TaggingGuardrail() + data: Final = { + "model": "invoice-reader", + "state": "Invoice from Acme Tools, contact jane@acme.example", + "fields": { + "vendor": {"type": "text", "instructions": "Who sent the invoice?"}, + "total": {"type": "number", "unit": "USD"}, + }, + "images": ["data:image/png;base64,AA=="], + } + + result: Final = await ExtractionHandler().process_input_messages(data, guardrail) + + assert guardrail.seen == ( + "request", + ("Invoice from Acme Tools, contact jane@acme.example", "Who sent the invoice?", "USD"), + ) + assert {key: value for key, value in result.items() if key != "metadata"} == { + **{key: value for key, value in data.items() if key != "metadata"}, + "state": "Invoice from Acme Tools, contact jane@acme.example [GUARDRAILED]", + "fields": { + "vendor": {"type": "text", "instructions": "Who sent the invoice? [GUARDRAILED]"}, + "total": {"type": "number", "unit": "USD [GUARDRAILED]"}, + }, + } + assert data["state"] == "Invoice from Acme Tools, contact jane@acme.example" + + +@pytest.mark.asyncio +async def test_a_request_with_no_text_or_images_skips_the_guardrail() -> None: + guardrail: Final = _TaggingGuardrail() + data: Final = {"model": "invoice-reader", "fields": {"total": {"type": "number"}}} + + assert await ExtractionHandler().process_input_messages(data, guardrail) is data + assert guardrail.seen is None + + +@pytest.mark.asyncio +async def test_request_images_reach_the_guardrail_as_data_urls() -> None: + guardrail: Final = _TaggingGuardrail() + data: Final = { + "model": "invoice-reader", + "fields": {"total": {"type": "number"}}, + "images": [{"content_type": "image/png", "base64": "iVBORw0KGgo="}, "data:image/jpeg;base64,/9j/"], + } + + result: Final = await ExtractionHandler().process_input_messages(data, guardrail) + + assert guardrail.seen == ("request", ()) + assert guardrail.seen_images == ("data:image/png;base64,iVBORw0KGgo=", "data:image/jpeg;base64,/9j/") + assert result["images"] == data["images"] + + +@pytest.mark.asyncio +async def test_response_answers_and_state_text_are_guardrailed_on_a_copy_without_stale_spans() -> None: + guardrail: Final = _TaggingGuardrail() + response: Final = ExtractionResponse.model_validate( + { + "model": "invoice-reader", + "answers": { + "vendor": {"type": "text", "value": "Acme Tools", "guess": "Acme Tools", "span": [19, 29]}, + "items": {"type": "list", "value": ["12", "hammers"], "spans": [[51, 53], [54, 61]]}, + "total": {"type": "number", "value": 174, "guess": 174}, + "contact": {"type": "text", "value": None, "p_null": 0.98}, + }, + "truncated": False, + "state_text": "Invoice from Acme Tools", + "usage": {"input_tokens": 52, "output_tokens": 0}, + } + ) + response.set_hidden_params({"response_cost": 0.0000052}) + + result: Final = await ExtractionHandler().process_output_response(response, guardrail) + + assert guardrail.seen == ( + "response", + ("Acme Tools", "12", "hammers", "174", "Invoice from Acme Tools"), + ) + assert isinstance(result, ExtractionResponse) + assert result.model_dump(mode="json") == { + **response.model_dump(mode="json"), + "answers": { + "vendor": {"type": "text", "value": "Acme Tools [GUARDRAILED]", "guess": "Acme Tools [GUARDRAILED]"}, + "items": {"type": "list", "value": ["12 [GUARDRAILED]", "hammers [GUARDRAILED]"]}, + "total": {"type": "number", "value": None}, + "contact": {"type": "text", "value": None, "p_null": 0.98}, + }, + "state_text": "Invoice from Acme Tools [GUARDRAILED]", + } + assert response.answers["vendor"].value == "Acme Tools" + assert result.hidden_params["response_cost"] == 0.0000052 + + +@pytest.mark.asyncio +async def test_spans_into_the_callers_own_state_survive_a_response_guardrail() -> None: + response: Final = ExtractionResponse.model_validate( + { + "model": "invoice-reader", + "answers": {"vendor": {"type": "text", "value": "Acme Tools", "span": [19, 29]}}, + "truncated": False, + "usage": {"input_tokens": 52, "output_tokens": 0}, + } + ) + + result: Final = await ExtractionHandler().process_output_response(response, _TaggingGuardrail()) + + assert isinstance(result, ExtractionResponse) + assert result.answers["vendor"].model_dump(mode="json") == { + "type": "text", + "value": "Acme Tools [GUARDRAILED]", + "span": [19, 29], + } + + +@pytest.mark.asyncio +async def test_a_response_guardrail_change_to_a_typed_answer_nulls_it_instead_of_breaking_its_type() -> None: + guardrail: Final = _MaskingGuardrail("2026-03-03", "[DATE]") + request_data: Final = { + "fields": { + "invoice_date": {"type": "date"}, + "due_dates": {"type": "list", "items": "date"}, + "notes": {"type": "list", "items": "text"}, + } + } + response: Final = ExtractionResponse.model_validate( + { + "model": "invoice-reader", + "answers": { + "invoice_date": {"type": "date", "value": "2026-03-03", "guess": "2026-03-03", "span": [30, 40]}, + "due_dates": {"type": "list", "value": ["2026-03-03", "2026-04-02"], "spans": [[30, 40], [50, 60]]}, + "notes": {"type": "list", "value": ["paid 2026-03-03"], "spans": [[70, 85]]}, + "received": {"type": "date", "value": "2026-02-27", "span": [90, 100]}, + }, + "truncated": False, + "usage": {"input_tokens": 52, "output_tokens": 0}, + } + ) + + result: Final = await ExtractionHandler().process_output_response(response, guardrail, request_data=request_data) + + assert isinstance(result, ExtractionResponse) + assert {name: answer.model_dump(mode="json") for name, answer in result.answers.items()} == { + "invoice_date": {"type": "date", "value": None}, + "due_dates": {"type": "list", "value": None}, + "notes": {"type": "list", "value": ["paid [DATE]"], "spans": [[70, 85]]}, + "received": {"type": "date", "value": "2026-02-27", "span": [90, 100]}, + } + + +@pytest.mark.asyncio +async def test_a_response_guardrail_that_masks_a_number_nulls_every_answer_carrying_it() -> None: + request_data: Final = { + "fields": { + "phone_text": {"type": "text"}, + "phone_number": {"type": "number"}, + "phone_list": {"type": "list", "items": "number"}, + "phone_mean": {"type": "estimate"}, + "phone_median": {"type": "estimate"}, + "phone_p10": {"type": "estimate"}, + "phone_p90": {"type": "estimate"}, + "phone_quantile": {"type": "estimate"}, + "phone_forecast": {"type": "forecast"}, + "total": {"type": "number"}, + } + } + response: Final = ExtractionResponse.model_validate( + { + "model": "invoice-reader", + "answers": { + "phone_text": {"type": "text", "value": "4155550123", "span": [10, 20]}, + "phone_number": {"type": "number", "value": 4155550123, "guess": 4155550123, "span": [10, 20]}, + "phone_list": {"type": "list", "value": [4155550123], "spans": [[10, 20]]}, + "phone_mean": { + "type": "estimate", + "value": None, + "mean": 4155550123.4, + "p10": 4155550000, + "p_null": 0.6, + }, + "phone_median": { + "type": "estimate", + "value": None, + "median": 4155550123, + "p90": 4155550200, + "p_null": 0.6, + }, + "phone_p10": {"type": "estimate", "value": None, "p10": 4155550123, "p_null": 0.6}, + "phone_p90": {"type": "estimate", "value": None, "p90": 4155550123.2, "p_null": 0.6}, + "phone_quantile": { + "type": "estimate", + "value": None, + "quantiles": [4155550000, 4155550123, 4155550200], + "p_null": 0.6, + }, + "phone_forecast": { + "type": "forecast", + "value": None, + "steps": [{"value": 12, "p10": 10}, {"value": 13, "p90": 4155550123}], + "p_null": 0.6, + }, + "total": {"type": "number", "value": 174, "guess": 174}, + }, + "truncated": False, + "usage": {"input_tokens": 52, "output_tokens": 0}, + } + ) + + result: Final = await ExtractionHandler().process_output_response( + response, _MaskingGuardrail("4155550123", "[PHONE]"), request_data=request_data + ) + + assert isinstance(result, ExtractionResponse) + assert {name: answer.model_dump(mode="json") for name, answer in result.answers.items()} == { + "phone_text": {"type": "text", "value": "[PHONE]", "span": [10, 20]}, + "phone_number": {"type": "number", "value": None}, + "phone_list": {"type": "list", "value": None}, + "phone_mean": {"type": "estimate", "value": None, "p_null": 0.6}, + "phone_median": {"type": "estimate", "value": None, "p_null": 0.6}, + "phone_p10": {"type": "estimate", "value": None, "p_null": 0.6}, + "phone_p90": {"type": "estimate", "value": None, "p_null": 0.6}, + "phone_quantile": {"type": "estimate", "value": None, "p_null": 0.6}, + "phone_forecast": {"type": "forecast", "value": None, "p_null": 0.6}, + "total": {"type": "number", "value": 174, "guess": 174}, + } + + +@pytest.mark.asyncio +async def test_a_response_guardrail_does_not_scan_the_fraction_digits_of_an_estimate() -> None: + days_late: Final = { + "type": "estimate", + "value": None, + "mean": 3.4525830151341586, + "median": 2.4525830151341586, + "p_null": 0.2, + } + response: Final = ExtractionResponse.model_validate( + { + "model": "invoice-reader", + "answers": {"days_late": days_late}, + "truncated": False, + "usage": {"input_tokens": 52, "output_tokens": 0}, + } + ) + + result: Final = await ExtractionHandler().process_output_response( + response, + _MaskingGuardrail("4525830151341586", "[CARD]"), + request_data={"fields": {"days_late": {"type": "estimate"}}}, + ) + + assert isinstance(result, ExtractionResponse) + assert result.answers["days_late"].model_dump(mode="json") == days_late + + +@pytest.mark.asyncio +async def test_a_response_guardrail_sees_the_decimal_text_of_a_fractional_number_answer() -> None: + response: Final = ExtractionResponse.model_validate( + { + "model": "invoice-reader", + "answers": { + "unit_price": {"type": "number", "value": 14.5, "guess": 14.5, "span": [30, 36]}, + "prices": {"type": "list", "value": [14.5, 2.25], "spans": [[30, 36], [40, 45]]}, + }, + "truncated": False, + "usage": {"input_tokens": 52, "output_tokens": 0}, + } + ) + + result: Final = await ExtractionHandler().process_output_response( + response, + _MaskingGuardrail("14.5", "[PRICE]"), + request_data={"fields": {"unit_price": {"type": "number"}, "prices": {"type": "list", "items": "number"}}}, + ) + + assert isinstance(result, ExtractionResponse) + assert {name: answer.model_dump(mode="json") for name, answer in result.answers.items()} == { + "unit_price": {"type": "number", "value": None}, + "prices": {"type": "list", "value": None}, + } diff --git a/tests/unit/llms/cloudflare/test_cloudflare_extraction_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_extraction_transformation.py new file mode 100644 index 00000000000..214fec07284 --- /dev/null +++ b/tests/unit/llms/cloudflare/test_cloudflare_extraction_transformation.py @@ -0,0 +1,221 @@ +import json +from typing import Final + +import httpx +import pytest +from pydantic import ValidationError + +import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig +from litellm.types.extraction import ExtractionRequestBody +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +INVOICE_STATE: Final = "Invoice #4471 from Acme Tools, dated March 3 2026. Total $174.00. Payment due in 30 days." +INVOICE_FIELDS: Final = { + "vendor": {"type": "text", "instructions": "Who sent the invoice"}, + "invoice_date": {"type": "date"}, + "total": {"type": "number", "decimals": 2, "unit": "USD"}, + "items": {"type": "list", "items": "text"}, + "days_to_pay": {"type": "estimate", "min": 0, "max": 90, "step": 1}, +} +CLAIRIC_RESULT: Final = { + "model": "clairic", + "answers": { + "vendor": { + "type": "text", + "p_null": 0.01, + "span": [19, 29], + "guess": "Acme Tools", + "confidence": 0.99, + "value": "Acme Tools", + }, + "total": {"type": "number", "p_null": 0.0, "guess": 174, "confidence": 0.99, "value": 174}, + "items": {"type": "list", "value": ["hammers"], "spans": [[54, 61]], "confidence": 0.4}, + "days_to_pay": { + "type": "estimate", + "value": 31, + "mean": 30.7, + "median": 30, + "p10": 29, + "p90": 34, + "p_below": 0.0, + "p_above": 0.0, + }, + }, + "truncated": False, + "usage": { + "state_tokens": 52, + "image_tokens": 0, + "schema_tokens": 312, + "input_tokens": 416, + "output_tokens": 0, + "passes": 2, + }, +} + + +def _cloudflare_config() -> BaseExtractionConfig: + config: Final = ProviderConfigManager.get_provider_extraction_config("clairic", LlmProviders.CLOUDFLARE) + assert config is not None + return config + + +def _response(payload: object) -> httpx.Response: + return httpx.Response(status_code=200, content=json.dumps(payload).encode()) + + +def test_only_cloudflare_serves_extraction() -> None: + assert ProviderConfigManager.get_provider_extraction_config("gpt-5", LlmProviders.OPENAI) is None + assert isinstance(_cloudflare_config(), litellm.CloudflareExtractionConfig) + + +@pytest.mark.parametrize("model", ("clairic", "@cf/cloudflare/clairic")) +def test_clairic_resolves_to_its_workers_ai_run_url(monkeypatch: pytest.MonkeyPatch, model: str) -> None: + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + config: Final = _cloudflare_config() + canonical: Final = config.canonical_model(model) + api_base: Final = config.resolve_api_base(None) + + assert canonical == "@cf/cloudflare/clairic" + assert api_base is not None + assert ( + config.get_complete_url(api_base, canonical) + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/@cf/cloudflare/clairic" + ) + + +@pytest.mark.parametrize( + "api_base", + ( + "https://gateway.example/client/v4/accounts/acct/ai/v1", + "https://gateway.example/client/v4/accounts/acct/ai/run/", + "https://gateway.example/client/v4/accounts/acct", + ), +) +def test_custom_api_base_keeps_the_run_path(api_base: str) -> None: + assert ( + _cloudflare_config().get_complete_url(api_base, "@cf/cloudflare/clairic") + == "https://gateway.example/client/v4/accounts/acct/ai/run/@cf/cloudflare/clairic" + ) + + +def test_missing_account_id_leaves_no_default_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("CLOUDFLARE_ACCOUNT_ID", raising=False) + monkeypatch.delenv("CLOUDFLARE_API_BASE", raising=False) + + assert _cloudflare_config().resolve_api_base(None) is None + + +def test_request_forwards_the_callers_fields_unchanged() -> None: + request: Final = ExtractionRequestBody.model_validate( + { + "state": INVOICE_STATE, + "fields": INVOICE_FIELDS, + "reference_date": "2026-03-01", + "options": {"return_state_text": True}, + "images": [{"content_type": "image/png", "base64": "iVBORw0KGgo="}], + } + ) + + assert _cloudflare_config().transform_extraction_request("@cf/cloudflare/clairic", request) == { + "model": "clairic", + "state": INVOICE_STATE, + "fields": INVOICE_FIELDS, + "reference_date": "2026-03-01", + "options": {"return_state_text": True}, + "images": [{"content_type": "image/png", "base64": "iVBORw0KGgo="}], + } + + +@pytest.mark.parametrize("wrapped", (True, False)) +def test_response_parses_typed_answers_and_usage(wrapped: bool) -> None: + payload: Final = ( + {"result": CLAIRIC_RESULT, "success": True, "errors": [], "messages": []} if wrapped else CLAIRIC_RESULT + ) + + response: Final = _cloudflare_config().transform_extraction_response( + "@cf/cloudflare/clairic", "cloudflare", _response(payload) + ) + + assert response.model_dump(mode="json") == CLAIRIC_RESULT + assert response.answers["vendor"].span == (19, 29) + assert response.usage.input_tokens == 416 + + +def test_unexpected_response_raises_a_provider_error() -> None: + with pytest.raises(BaseLLMException) as error: + _cloudflare_config().transform_extraction_response( + "@cf/cloudflare/clairic", "cloudflare", _response({"result": {"answers": {}}}) + ) + + assert error.value.status_code == 500 + + +def test_unknown_field_type_is_rejected() -> None: + with pytest.raises(ValidationError, match=r"fields\.vendor\.type"): + ExtractionRequestBody.model_validate({"state": INVOICE_STATE, "fields": {"vendor": {"type": "color"}}}) + + +def test_a_list_field_without_an_item_type_is_rejected_locally() -> None: + # Cloudflare answered a list field without items with 422 code 3030 on @cf/cloudflare/clairic, probed 2026-10-20 + with pytest.raises(ValidationError, match=r"fields\.line_items\n.*list fields need items"): + ExtractionRequestBody.model_validate({"state": INVOICE_STATE, "fields": {"line_items": {"type": "list"}}}) + + +@pytest.mark.parametrize( + "spec", + ( + {"type": "number", "decimals": "2"}, + {"type": "number", "integer": "yes"}, + {"type": "number", "integer": 1}, + {"type": "estimate", "min": True, "max": 90}, + ), + ids=("string_decimals", "string_integer", "number_integer", "bool_min"), +) +def test_a_field_value_of_the_wrong_json_type_is_rejected_instead_of_coerced(spec: dict[str, object]) -> None: + # Cloudflare answered each of these specs with a 400 or 422 on @cf/cloudflare/clairic, probed 2026-10-20 + with pytest.raises(ValidationError, match=r"fields\.total\."): + ExtractionRequestBody.model_validate({"state": INVOICE_STATE, "fields": {"total": spec}}) + + +def test_numbers_in_a_field_spec_go_to_cloudflare_exactly_as_sent() -> None: + # Cloudflare answered "decimals": 2.0 and "step": 1.0 with a 200 on @cf/cloudflare/clairic, probed 2026-10-20 + fields: Final = { + "total": {"type": "number", "decimals": 2.0, "null_threshold": 1}, + "days_to_pay": {"type": "estimate", "min": 0, "max": 90, "step": 1.0}, + } + request: Final = ExtractionRequestBody.model_validate({"state": INVOICE_STATE, "fields": fields}) + + body: Final = _cloudflare_config().transform_extraction_request("@cf/cloudflare/clairic", request) + + assert json.dumps(body["fields"]) == json.dumps(fields) + + +@pytest.mark.parametrize("name", ("due date", "1st_item", "x" * 80), ids=("space", "digit_first", "80_chars")) +def test_a_field_name_clairic_refuses_is_rejected_locally(name: str) -> None: + # Cloudflare answered each of these names with 422 code 3030 on @cf/cloudflare/clairic, probed 2026-10-19 + with pytest.raises(ValidationError, match=r"fields\."): + ExtractionRequestBody.model_validate({"state": INVOICE_STATE, "fields": {name: {"type": "text"}}}) + + +@pytest.mark.parametrize("name", ("_vendor", "line.items", "days-to-pay", "x" * 64)) +def test_field_names_with_dots_dashes_and_underscores_up_to_64_chars_are_accepted(name: str) -> None: + # Cloudflare answered each of these names with a 200 on @cf/cloudflare/clairic, probed 2026-10-19 + request: Final = ExtractionRequestBody.model_validate({"state": INVOICE_STATE, "fields": {name: {"type": "text"}}}) + + assert tuple(request.fields) == (name,) + + +def test_decisions_config_still_resolves_clef_after_the_shared_helpers_moved(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + config: Final = litellm.CloudflareDecisionsConfig() + api_base: Final = config.get_default_api_base() + + assert api_base is not None + assert ( + config.get_complete_url(api_base, config.canonical_model("clef")) + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/@cf/cloudflare/clef" + ) + assert config.unwrap_response({"result": {"answers": {}}}) == {"answers": {}} diff --git a/tests/unit/llms/github_copilot/test_github_host.py b/tests/unit/llms/github_copilot/test_github_host.py new file mode 100644 index 00000000000..df2efe4669b --- /dev/null +++ b/tests/unit/llms/github_copilot/test_github_host.py @@ -0,0 +1,59 @@ +from typing import Final + +import pytest + +from litellm.llms.github_copilot.github_host import ( + github_host_endpoints, + normalized_github_host, +) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, "github.com"), + ("", "github.com"), + ("GitHub.com", "github.com"), + ("OCTO.ghe.com", "octo.ghe.com"), + ("octo-1.ghe.com", "octo-1.ghe.com"), + ("https://octo.ghe.com", None), + ("ghes.example.com", None), + ("octo.ghe.com:443", None), + ("a.b.ghe.com", None), + ("octo.ghe.com.", None), + (" api.octo.ghe.com", None), + ("octo.ghe.com/path", None), + ], +) +def test_normalized_github_host_accepts_only_bare_supported_hosts(raw: object, expected: str | None) -> None: + assert normalized_github_host(raw) == expected + + +def test_github_com_endpoints_read_url_overrides_at_call_time(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://device.example/code") + monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", "https://device.example/token") + monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", "https://api.example/copilot") + monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", "custom-client") + + endpoints: Final = github_host_endpoints("github.com") + + assert endpoints.device_code_url == "https://device.example/code" + assert endpoints.access_token_url == "https://device.example/token" + assert endpoints.api_key_url == "https://api.example/copilot" + assert endpoints.client_id == "custom-client" + + +def test_ghe_endpoints_ignore_github_com_url_overrides(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://device.example/code") + monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", "https://device.example/token") + monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", "https://api.example/copilot") + monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", "custom-client") + + endpoints: Final = github_host_endpoints("octo.ghe.com") + + assert endpoints.device_code_url == "https://octo.ghe.com/login/device/code" + assert endpoints.access_token_url == "https://octo.ghe.com/login/oauth/access_token" + assert endpoints.api_key_url == "https://api.octo.ghe.com/copilot_internal/v2/token" + assert endpoints.user_url == "https://api.octo.ghe.com/user" + assert endpoints.default_copilot_api_base == "https://copilot-api.octo.ghe.com" + assert endpoints.client_id == "custom-client" diff --git a/tests/unit/llms/github_copilot/test_per_user_auth.py b/tests/unit/llms/github_copilot/test_per_user_auth.py index 9b16d3c0aa7..cf6fb9fa71f 100644 --- a/tests/unit/llms/github_copilot/test_per_user_auth.py +++ b/tests/unit/llms/github_copilot/test_per_user_auth.py @@ -12,6 +12,7 @@ import httpx import pytest import litellm +from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST from litellm.exceptions import ( AuthenticationError, BadRequestError, @@ -32,20 +33,33 @@ from litellm.llms.github_copilot.per_user_auth import ( attach_github_copilot_user_session, evict_copilot_user_session, exchange_github_token, + github_copilot_host_endpoints, github_copilot_auth_mode, github_copilot_user_session_from, validated_copilot_api_base, ) +from litellm.llms.github_copilot.github_host import github_host_endpoints from litellm.types.utils import CredentialItem, ModelResponse EXCHANGE_URL = "https://api.github.com/copilot_internal/v2/token" DEVICE_CODE_URL = "https://github.com/login/device/code" ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token" +_GITHUB_ENDPOINTS: Final = github_host_endpoints("github.com") -def _credential(auth_type="per_user_oauth"): - values = {} if auth_type is None else {"github_copilot_auth_type": auth_type} - return CredentialItem(credential_name="copilot-cred", credential_values=values, credential_info={}) +def _credential( + auth_type: str | None = "per_user_oauth", + github_host: str | None = None, +) -> CredentialItem: + values: Final = {} if auth_type is None else {"github_copilot_auth_type": auth_type} + return CredentialItem( + credential_name="copilot-cred", + credential_values={ + **values, + **({"github_copilot_host": github_host} if github_host is not None else {}), + }, + credential_info={}, + ) def _exchange_payload(token="copilot-token", expires_at=None, api_base="https://api.githubcopilot.com"): @@ -143,7 +157,61 @@ def test_unknown_named_credential_falls_back_to_kwargs(): ], ) def test_validated_copilot_api_base(raw, expected): - assert validated_copilot_api_base(raw) == expected + assert validated_copilot_api_base(raw, github_host_endpoints("github.com")) == expected + + +@pytest.mark.parametrize( + "raw", + [ + "https://copilot-api.evil.ghe.com", + "https://api.githubcopilot.com", + "http://copilot-api.octo.ghe.com", + "https://user@copilot-api.octo.ghe.com", + "https://copilot-api.octo.ghe.com:8443", + ], +) +def test_ghe_exchange_response_must_stay_on_its_tenant_host(raw: object) -> None: + endpoints: Final = github_host_endpoints("octo.ghe.com") + assert validated_copilot_api_base(raw, endpoints) == "https://copilot-api.octo.ghe.com" + + +def test_github_com_exchange_rejects_a_ghe_copilot_host() -> None: + endpoints: Final = github_host_endpoints("github.com") + assert validated_copilot_api_base("https://copilot-api.octo.ghe.com", endpoints) == "https://api.githubcopilot.com" + + +def test_ghe_exchange_uses_stored_host_not_request_kwargs(): + credential: Final = CredentialItem( + credential_name="copilot-cred", + credential_values={"github_copilot_auth_type": "per_user_oauth", "github_copilot_host": "octo.ghe.com"}, + credential_info={}, + ) + + def respond(request: httpx.Request) -> httpx.Response: + assert str(request.url) == "https://api.octo.ghe.com/copilot_internal/v2/token" + return httpx.Response(200, json=_exchange_payload(api_base="https://copilot-api.octo.ghe.com"), request=request) + + with patch.object(litellm, "credential_list", [credential]), github_http(respond) as requests: + kwargs: Final = _kwargs_with_connection(github_host="octo.ghe.com") + kwargs["github_copilot_host"] = "evil.ghe.com" + attach_github_copilot_user_session(kwargs) + + session: Final = github_copilot_user_session_from(kwargs) + assert session is not None and session.api_base == "https://copilot-api.octo.ghe.com" + assert len(requests) == 1 + + +def test_invalid_stored_host_fails_before_exchange(): + credential: Final = CredentialItem( + credential_name="copilot-cred", + credential_values={"github_copilot_auth_type": "per_user_oauth", "github_copilot_host": "https://octo.ghe.com"}, + credential_info={}, + ) + with patch.object(litellm, "credential_list", [credential]): + with github_http(lambda request: pytest.fail(f"unexpected request to {request.url}")) as requests: + with pytest.raises(BadRequestError): + attach_github_copilot_user_session(_kwargs_with_connection()) + assert not requests def test_session_repr_and_str_hide_the_token(caplog): @@ -159,49 +227,59 @@ def test_session_repr_and_str_hide_the_token(caplog): def test_exchange_returns_session_and_caches_per_user(): with github_http(_exchange_responder()) as requests: - s1 = exchange_github_token("user-a", "gh-token", "copilot-cred") + s1 = exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 1 assert requests[0].headers["Authorization"] == "token gh-token" - s2 = exchange_github_token("user-a", "gh-token", "copilot-cred") + s2 = exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 1 assert s2.token == s1.token == "copilot-token" # different user -> separate exchange - exchange_github_token("user-b", "gh-token", "copilot-cred") + exchange_github_token("user-b", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 2 # different github token -> separate exchange - exchange_github_token("user-a", "gh-token-2", "copilot-cred") + exchange_github_token("user-a", "gh-token-2", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 3 +def test_cached_session_is_not_reused_after_github_token_changes(): + with github_http(_exchange_responder()) as requests: + exchange_github_token("user-a", "gh-token-a", "copilot-cred", _GITHUB_ENDPOINTS) + exchange_github_token("user-a", "gh-token-b", "copilot-cred", _GITHUB_ENDPOINTS) + + assert [request.headers["Authorization"] for request in requests] == ["token gh-token-a", "token gh-token-b"] + + @pytest.mark.parametrize("status", [401, 403, 404]) def test_exchange_auth_failures_raise_caller_auth_error_and_evict(status): with github_http(_exchange_responder(status=status)): with pytest.raises(CallerCredentialAuthenticationError) as exc: - exchange_github_token("user-a", "gh-token", "copilot-cred") + exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert "Reconnect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials)" in str( exc.value ) - assert _SESSION_CACHE.get_cache(_session_cache_key("user-a", "gh-token")) is None + assert ( + _SESSION_CACHE.get_cache(_session_cache_key("user-a", "copilot-cred", GITHUB_COPILOT_DEFAULT_HOST)) is None + ) def test_exchange_429_raises_caller_rate_limit(): with github_http(_exchange_responder(status=429)): with pytest.raises(CallerCredentialRateLimitError): - exchange_github_token("user-a", "gh-token", "copilot-cred") + exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) def test_exchange_other_failure_raises_service_unavailable_without_body(): with github_http(_exchange_responder(status=500, payload={"secret": "body"})): with pytest.raises(ServiceUnavailableError) as exc: - exchange_github_token("user-a", "gh-token", "copilot-cred") + exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert "body" not in str(exc.value) def test_exchange_expiry_margin_refetches_near_expiry(): with github_http(_exchange_responder(_exchange_payload(expires_at=int(time.time()) + 30))) as requests: - exchange_github_token("user-a", "gh-token", "copilot-cred") + exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) # ttl under the 60s safety margin -> not cached, second call refetches - exchange_github_token("user-a", "gh-token", "copilot-cred") + exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 2 @@ -212,8 +290,8 @@ async def test_aexchange_single_flight_one_http_call(): with github_http(respond) as requests: s1, s2 = await asyncio.gather( - aexchange_github_token("user-a", "gh-token", "copilot-cred"), - aexchange_github_token("user-a", "gh-token", "copilot-cred"), + aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS), + aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS), ) assert len(requests) == 1 assert s1.token == s2.token == "copilot-token" @@ -223,16 +301,26 @@ async def test_aexchange_single_flight_one_http_call(): async def test_aexchange_401_raises_and_evicts(): with github_http(_exchange_responder(status=401)): with pytest.raises(CallerCredentialAuthenticationError): - await aexchange_github_token("user-a", "gh-token", "copilot-cred") - assert _SESSION_CACHE.get_cache(_session_cache_key("user-a", "gh-token")) is None + await aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + assert ( + _SESSION_CACHE.get_cache(_session_cache_key("user-a", "copilot-cred", GITHUB_COPILOT_DEFAULT_HOST)) is None + ) -def _kwargs_with_connection(token="gho_user_token"): +def _kwargs_with_connection( + token: str = "gho_user_token", + github_host: str | None = "github.com", +) -> dict[str, object]: return { "litellm_credential_name": "copilot-cred", "secret_fields": { "user_provider_credentials_user_id": "user-a", - "user_provider_credentials": {"copilot-cred": token}, + "user_provider_credentials": { + "copilot-cred": { + "access_token": token, + **({"github_host": github_host} if github_host is not None else {}), + } + }, }, } @@ -261,6 +349,31 @@ def test_attach_per_user_without_connection_raises_connect_401(per_user_credenti assert "Connect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials)" in str(exc.value) +@pytest.mark.parametrize( + ("credential_host", "connection_host"), + [("github.com", "octo.ghe.com"), ("octo.ghe.com", "github.com")], +) +def test_attach_rejects_connection_from_a_different_github_host( + credential_host: str, + connection_host: str, +): + with patch.object(litellm, "credential_list", [_credential(github_host=credential_host)]): + kwargs: Final = _kwargs_with_connection(github_host=connection_host) + with github_http(lambda request: pytest.fail(f"unexpected token exchange at {request.url}")) as requests: + with pytest.raises(CallerCredentialAuthenticationError) as exc: + attach_github_copilot_user_session(kwargs) + assert "Connect GitHub Copilot for credential 'copilot-cred'" in str(exc.value) + assert not requests + + +def test_legacy_connection_without_host_uses_github_com(per_user_credential): + with github_http(_exchange_responder()) as requests: + kwargs: Final = _kwargs_with_connection(github_host=None) + attach_github_copilot_user_session(kwargs) + assert len(requests) == 1 + assert str(requests[0].url) == EXCHANGE_URL + + def test_attach_shared_mode_does_nothing_without_connection(): with patch.object(litellm, "credential_list", [_credential(auth_type=None)]): kwargs = {"litellm_credential_name": "copilot-cred"} @@ -270,9 +383,9 @@ def test_attach_shared_mode_does_nothing_without_connection(): def test_evict_drops_cached_session(): with github_http(_exchange_responder()) as requests: - exchange_github_token("user-a", "gh-token", "copilot-cred") - evict_copilot_user_session("user-a", "gh-token") - exchange_github_token("user-a", "gh-token", "copilot-cred") + exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + evict_copilot_user_session("user-a", "copilot-cred", _GITHUB_ENDPOINTS.host) + exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 2 @@ -293,7 +406,7 @@ async def test_astart_device_flow_shape(): ) with github_http(respond): - start = await astart_device_flow() + start = await astart_device_flow(_GITHUB_ENDPOINTS) assert (start.device_code, start.user_code, start.verification_uri, start.expires_in, start.interval) == ( "dc", "UC-123", @@ -319,7 +432,7 @@ async def test_apoll_device_flow_error_mapping(github_error, expected_status): return httpx.Response(200, json={"error": github_error}, request=request) with github_http(respond): - poll = await apoll_device_flow("dc") + poll = await apoll_device_flow("dc", _GITHUB_ENDPOINTS) assert poll.status == expected_status @@ -330,7 +443,7 @@ async def test_apoll_device_flow_connected(): return httpx.Response(200, json={"access_token": "gho_x"}, request=request) with github_http(respond): - poll = await apoll_device_flow("dc") + poll = await apoll_device_flow("dc", _GITHUB_ENDPOINTS) assert poll.status == "connected" and poll.access_token == "gho_x" @@ -363,7 +476,12 @@ def _chat_completion_response(): def _secret_fields(user_id: str, github_token: str): return { "user_provider_credentials_user_id": user_id, - "user_provider_credentials": {"copilot-cred": github_token}, + "user_provider_credentials": { + "copilot-cred": { + "access_token": github_token, + "github_host": GITHUB_COPILOT_DEFAULT_HOST, + } + }, } @@ -842,7 +960,9 @@ def test_sync_exchange_single_flight_one_http_call(): with github_http(slow_respond) as requests: threads = [ threading.Thread( - target=lambda: sessions.append(exchange_github_token("user-a", "gh-token", "copilot-cred")) + target=lambda: sessions.append( + exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + ) ) for _ in range(6) ] @@ -856,12 +976,34 @@ def test_sync_exchange_single_flight_one_http_call(): def test_two_users_same_github_token_do_not_share_session(): with github_http(_exchange_responder()): - s1 = exchange_github_token("user-a", "gh-same", "copilot-cred") - s2 = exchange_github_token("user-b", "gh-same", "copilot-cred") - assert _session_cache_key("user-a", "gh-same") != _session_cache_key("user-b", "gh-same") + s1: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) + s2: Final = exchange_github_token("user-b", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) + user_a_key: Final = _session_cache_key("user-a", "copilot-cred", _GITHUB_ENDPOINTS.host) + user_b_key: Final = _session_cache_key("user-b", "copilot-cred", _GITHUB_ENDPOINTS.host) + other_credential_key: Final = _session_cache_key("user-a", "another-credential", _GITHUB_ENDPOINTS.host) + other_host_key: Final = _session_cache_key("user-a", "copilot-cred", "octo.ghe.com") + assert user_a_key != user_b_key + assert user_a_key != other_credential_key + assert user_a_key != other_host_key assert s1.token == s2.token == "copilot-token" +def test_same_user_and_token_do_not_share_sessions_across_hosts(): + from litellm.llms.github_copilot.github_host import github_host_endpoints + + ghe_endpoints: Final = github_host_endpoints("octo.ghe.com") + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=_exchange_payload(), request=request) + + with github_http(respond) as requests: + github_session: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) + ghe_session: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", ghe_endpoints) + assert len(requests) == 2 + assert github_session.api_base == "https://api.githubcopilot.com" + assert ghe_session.api_base == "https://copilot-api.octo.ghe.com" + + def test_router_constructs_with_per_user_credential_deployment(per_user_credential): from litellm.llms.github_copilot.authenticator import Authenticator @@ -928,9 +1070,9 @@ async def test_aexchange_cancelled_waiter_does_not_cancel_the_shared_exchange(pe with patch( # test-quality-ok: the gate holds the HTTP edge open so two waiters overlap "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=_GatedClient() ): - first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred")) + first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS)) await asyncio.sleep(0) - second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred")) + second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS)) await asyncio.sleep(0) second.cancel() gate.set() @@ -940,6 +1082,33 @@ async def test_aexchange_cancelled_waiter_does_not_cancel_the_shared_exchange(pe await second +@pytest.mark.asyncio +async def test_aexchange_does_not_share_an_in_flight_session_after_token_changes(per_user_credential): + gate = asyncio.Event() + auth_headers: asyncio.Queue[str] = asyncio.Queue() + + class _GatedClient: + async def get(self, url, headers=None): + await auth_headers.put(headers["authorization"]) + await gate.wait() + token = "copilot-token-a" if headers["authorization"] == "token gh-a" else "copilot-token-b" + return httpx.Response( + 200, + json=_exchange_payload(token=token), + request=httpx.Request("GET", url), + ) + + with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=_GatedClient()): + first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-a", "copilot-cred", _GITHUB_ENDPOINTS)) + assert await asyncio.wait_for(auth_headers.get(), timeout=1) == "token gh-a" + second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-b", "copilot-cred", _GITHUB_ENDPOINTS)) + assert await asyncio.wait_for(auth_headers.get(), timeout=1) == "token gh-b" + gate.set() + sessions = await asyncio.gather(first, second) + + assert [session.token for session in sessions] == ["copilot-token-a", "copilot-token-b"] + + @pytest.mark.asyncio async def test_aexchange_failure_with_a_single_caller_emits_no_unretrieved_warning(per_user_credential): """The leader removes the failed future before anyone awaits it; the stored @@ -955,7 +1124,7 @@ async def test_aexchange_failure_with_a_single_caller_emits_no_unretrieved_warni github_http(_exchange_responder(status=401)), pytest.raises(CallerCredentialAuthenticationError), ): - await aexchange_github_token("user-a", "gh-fail", "copilot-cred") + await aexchange_github_token("user-a", "gh-fail", "copilot-cred", _GITHUB_ENDPOINTS) _IN_FLIGHT_KEYS = list(pua._IN_FLIGHT) import gc diff --git a/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py b/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py index 2f1b3764610..04c2ce153c6 100644 --- a/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py +++ b/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py @@ -10,6 +10,7 @@ import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.huggingface.common_utils import _fetch_inference_provider_mapping +from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, _load_huggingface_tokenizer HUGGINGFACE_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct" HUGGINGFACE_ROUTER_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct-Turbo" @@ -83,9 +84,13 @@ def isolate_huggingface_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator monkeypatch.delenv("HF_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_KEY", raising=False) + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() _fetch_inference_provider_mapping.cache_clear() yield _fetch_inference_provider_mapping.cache_clear() + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() PROVIDER_MAPPING_RESPONSE = { diff --git a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py index 7ecc0b47d9f..dbb5cde6060 100644 --- a/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py +++ b/tests/unit/llms/nvidia_riva/audio_transcription/test_handler.py @@ -12,11 +12,12 @@ import io from types import SimpleNamespace from unittest.mock import MagicMock +import httpx import numpy as np import pytest import soundfile as sf - +import litellm from litellm.llms.nvidia_riva.audio_transcription import handler as handler_mod from litellm.llms.nvidia_riva.audio_transcription.handler import ( NvidiaRivaAudioTranscription, @@ -93,6 +94,24 @@ def mock_riva(monkeypatch): LINEAR_PCM = "LINEAR_PCM" streaming_responses_holder = {"value": []} + offline_response_holder = {"value": None} + offline_error_holder = {"value": None} + offline_future_holder = {"value": None} + streaming_calls = [] + offline_calls = {} + + class FakeRecognizeFuture: + def result(self, timeout: float | None = None) -> object: + offline_calls["timeout"] = timeout + if offline_error_holder["value"] is not None: + raise offline_error_holder["value"] + return offline_response_holder["value"] + + def done(self) -> bool: + return True + + def cancel(self) -> bool: + return False class FakeASRService: def __init__(self, auth): @@ -100,9 +119,15 @@ def mock_riva(monkeypatch): def streaming_response_generator(self, audio_chunks, streaming_config): # Drain audio_chunks generator so we exercise the chunking path. - list(audio_chunks) + streaming_calls.append(list(audio_chunks)) yield from streaming_responses_holder["value"] + def offline_recognize(self, audio_bytes: bytes, config: object, future: bool = False) -> object: + offline_calls["audio_bytes"] = audio_bytes + offline_calls["config"] = config + pending = offline_future_holder["value"] or FakeRecognizeFuture() + return pending if future else pending.result() + fake_riva_client = SimpleNamespace( Auth=FakeAuth, ASRService=FakeASRService, @@ -120,6 +145,11 @@ def mock_riva(monkeypatch): return SimpleNamespace( auth_calls=auth_calls, responses=streaming_responses_holder, + offline_response=offline_response_holder, + offline_error=offline_error_holder, + offline_future=offline_future_holder, + streaming_calls=streaming_calls, + offline_calls=offline_calls, client=fake_riva_client, ) @@ -203,6 +233,339 @@ def test_sync_path_aggregates_only_final_results(mock_riva, logging_obj): ) +def test_riva_offline_uses_unary_recognize_and_keeps_results_without_is_final( + mock_riva: SimpleNamespace, logging_obj: MagicMock +) -> None: + """ + Offline-only NIMs (Parakeet TDT, Whisper) reject StreamingRecognize. With + ``riva_offline`` the handler must call ``offline_recognize`` with the whole + resampled buffer and a bare ``RecognitionConfig``, and must keep every + result even though ``RecognizeResponse`` results carry no ``is_final``. + """ + mock_riva.offline_response["value"] = _fake_response( + results=[ + SimpleNamespace(alternatives=[_fake_alternative("Hello,", words=[_fake_word("Hello,", 0, 320)])]), + SimpleNamespace(alternatives=[_fake_alternative(" world.", words=[_fake_word("world.", 480, 870)])]), + ] + ) + + impl = NvidiaRivaAudioTranscription() + response: TranscriptionResponse = impl.audio_transcriptions( + model="nvidia/parakeet-0.6b-tdt", + audio_file=_make_wav_bytes(), + optional_params={ + "language_code": "multi", + "riva_offline": True, + "enable_word_time_offsets": True, + "response_format": "verbose_json", + "timestamp_granularities": ["word"], + }, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert response.text == "Hello, world." + assert response["words"][1]["end"] == pytest.approx(0.87) + assert response._hidden_params["audio_transcription_duration"] == pytest.approx(1.0, abs=0.05) + assert mock_riva.streaming_calls == [] + assert len(mock_riva.offline_calls["audio_bytes"]) == pytest.approx(32000, abs=1600) + assert mock_riva.offline_calls["config"]._kwargs["language_code"] == "multi" + assert logging_obj.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"]["riva_offline"] is True + + +def test_default_mode_streams_and_never_calls_offline_recognize( + mock_riva: SimpleNamespace, logging_obj: MagicMock +) -> None: + mock_riva.responses["value"] = [ + _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + ] + + impl = NvidiaRivaAudioTranscription() + response = impl.audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US"}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + assert response.text == "streamed" + assert mock_riva.offline_calls == {} + assert len(mock_riva.streaming_calls) == 1 + assert logging_obj.pre_call.call_args.kwargs["additional_args"]["complete_input_dict"]["riva_offline"] is False + + + +@pytest.mark.parametrize(("riva_offline", "expect_offline"), [("false", False), ("False", False), ("true", True)]) +def test_riva_offline_string_from_multipart_form_is_parsed( + mock_riva: SimpleNamespace, logging_obj: MagicMock, riva_offline: str, expect_offline: bool +) -> None: + mock_riva.responses["value"] = [ + _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + ] + mock_riva.offline_response["value"] = _fake_response( + results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] + ) + + response = NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": riva_offline}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert response.text == ("offline" if expect_offline else "streamed") + assert (mock_riva.offline_calls != {}) is expect_offline + + +@pytest.mark.parametrize("riva_offline", ["yes", "", 1]) +def test_riva_offline_malformed_value_is_rejected_before_any_rpc( + mock_riva: SimpleNamespace, logging_obj: MagicMock, riva_offline: object +) -> None: + with pytest.raises(NvidiaRivaException) as excinfo: + NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": riva_offline}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert excinfo.value.status_code == 400 + assert "riva_offline" in excinfo.value.message + assert mock_riva.offline_calls == {} + assert mock_riva.streaming_calls == [] + + +def test_riva_offline_malformed_value_is_dropped_under_drop_params( + mock_riva: SimpleNamespace, logging_obj: MagicMock +) -> None: + mock_riva.responses["value"] = [ + _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + ] + + response = NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": "yes"}, + litellm_params={"drop_params": True}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert response.text == "streamed" + assert mock_riva.offline_calls == {} + + +def test_deployment_riva_offline_reaches_the_handler_through_transcription(mock_riva: SimpleNamespace) -> None: + mock_riva.offline_response["value"] = _fake_response( + results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] + ) + + response = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + riva_offline=True, + ) + + assert response.text == "offline" + assert mock_riva.streaming_calls == [] + assert mock_riva.offline_calls != {} + + +@pytest.mark.parametrize( + ("timeout", "waited_for"), + [("7.5", 7.5), (httpx.Timeout(30.0, read=7.5), 7.5), (httpx.Timeout(5.0, read=None), 600.0), (None, None)], +) +def test_offline_recognition_waits_the_timeout_in_seconds_whatever_its_form( + mock_riva: SimpleNamespace, timeout: str | httpx.Timeout | None, waited_for: float | None +) -> None: + mock_riva.offline_response["value"] = _fake_response( + results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] + ) + + response = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + riva_offline=True, + timeout=timeout, + ) + + assert response.text == "offline" + assert mock_riva.offline_calls["timeout"] == waited_for + + +@pytest.mark.parametrize("timeout", ["abc", "30s"]) +def test_malformed_timeout_on_offline_call_is_a_400_before_any_rpc( + mock_riva: SimpleNamespace, timeout: str +) -> None: + with pytest.raises(NvidiaRivaException) as excinfo: + litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + riva_offline=True, + timeout=timeout, + ) + + assert excinfo.value.status_code == 400 + assert "timeout" in excinfo.value.message + assert "audio_bytes" not in mock_riva.offline_calls + + +def test_malformed_timeout_on_offline_call_waits_the_default_under_drop_params( + mock_riva: SimpleNamespace, +) -> None: + mock_riva.offline_response["value"] = _fake_response( + results=[SimpleNamespace(alternatives=[_fake_alternative("offline")])] + ) + + response = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-0.6b-tdt", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + riva_offline=True, + timeout="abc", + drop_params=True, + ) + + assert response.text == "offline" + assert mock_riva.offline_calls["timeout"] == 600.0 + + +def test_malformed_timeout_still_streams_when_the_client_takes_no_timeout(mock_riva: SimpleNamespace) -> None: + mock_riva.responses["value"] = [ + _fake_response(results=[_fake_result(is_final=True, alternatives=[_fake_alternative("streamed")])]) + ] + + response = litellm.transcription( + model="nvidia_riva/nvidia/parakeet-ctc-1.1b", + file=("clip.wav", _make_wav_bytes(), "audio/wav"), + api_base="localhost:50051", + timeout="abc", + ) + + assert response.text == "streamed" + assert len(mock_riva.streaming_calls) == 1 + + +def test_oversized_offline_audio_is_a_non_retryable_413(mock_riva: SimpleNamespace, logging_obj: MagicMock) -> None: + class FakeGrpcError(Exception): + def code(self) -> SimpleNamespace: + return SimpleNamespace(name="RESOURCE_EXHAUSTED") + + def details(self) -> str: + return "SERVER: Received message larger than max (4348825 vs. 4194304)" + + mock_riva.offline_error["value"] = FakeGrpcError("rpc fail") + + with pytest.raises(NvidiaRivaException) as excinfo: + NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": True}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert excinfo.value.status_code == 413 + assert litellm.utils.should_retry(excinfo.value.status_code) is False + + +def test_stalled_offline_recognition_is_cancelled_and_times_out_as_408( + mock_riva: SimpleNamespace, logging_obj: MagicMock +) -> None: + class FutureTimeoutError(Exception): + pass + + class StalledRecognizeFuture: + def __init__(self) -> None: + self.waited_for: float | None = None + self.cancelled = False + + def result(self, timeout: float | None = None) -> object: + self.waited_for = timeout + raise FutureTimeoutError() + + def done(self) -> bool: + return self.cancelled + + def cancel(self) -> bool: + self.cancelled = True + return True + + stalled = StalledRecognizeFuture() + mock_riva.offline_future["value"] = stalled + + with pytest.raises(NvidiaRivaException) as excinfo: + NvidiaRivaAudioTranscription().audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "en-US", "riva_offline": True}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=7.5, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + + assert excinfo.value.status_code == 408 + assert stalled.waited_for == 7.5 + assert stalled.cancelled is True + +def test_riva_offline_grpc_error_is_wrapped(mock_riva: SimpleNamespace, logging_obj: MagicMock) -> None: + class FakeGrpcError(Exception): + def code(self) -> SimpleNamespace: + return SimpleNamespace(name="INVALID_ARGUMENT") + + def details(self) -> str: + return "Unavailable model requested" + + mock_riva.offline_error["value"] = FakeGrpcError("rpc fail") + + impl = NvidiaRivaAudioTranscription() + with pytest.raises(NvidiaRivaException) as excinfo: + impl.audio_transcriptions( + model="m", + audio_file=_make_wav_bytes(), + optional_params={"language_code": "multi", "riva_offline": True}, + litellm_params={}, + model_response=TranscriptionResponse(), + timeout=60, + logging_obj=logging_obj, + api_key=None, + api_base="localhost:50051", + ) + assert excinfo.value.status_code == 400 + assert "INVALID_ARGUMENT" in excinfo.value.message + + def test_auth_nvcf_defaults_use_ssl_and_attaches_function_id(mock_riva, logging_obj): mock_riva.responses["value"] = [ _fake_response( @@ -343,9 +706,10 @@ def test_async_path_uses_to_thread(mock_riva, logging_obj): ) +@pytest.mark.parametrize("timeout", [12.5, "12.5", httpx.Timeout(60.0, read=12.5)]) def test_timeout_is_forwarded_to_streaming_generator_when_supported( - mock_riva, logging_obj -): + mock_riva: SimpleNamespace, logging_obj: MagicMock, timeout: float | str | httpx.Timeout +) -> None: """ Without a deadline the gRPC stream can block forever on a stalled Riva server. The handler must forward the call-level ``timeout`` to @@ -374,7 +738,7 @@ def test_timeout_is_forwarded_to_streaming_generator_when_supported( optional_params={"language_code": "en-US"}, litellm_params={}, model_response=TranscriptionResponse(), - timeout=12.5, + timeout=timeout, logging_obj=logging_obj, api_key=None, api_base="localhost:50051", diff --git a/tests/unit/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py index f465bbea663..9e2754eef68 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -364,7 +364,8 @@ def test_text_completion_stream_forwards_stream_options_and_usage(monkeypatch): ) ) - assert tuple(chunk.choices[0].text for chunk in chunks) == ("hello", None, None) + assert tuple(chunk.choices[0].text for chunk in chunks[:-1]) == ("hello", None) + assert chunks[-1].choices == [] assert tuple(chunk.usage for chunk in chunks[:-1]) == (None, None) assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == ( 2, diff --git a/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py index 12690d5822b..56c0aee3efd 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py @@ -543,8 +543,9 @@ def test_nova_streaming_sends_stream_flag_and_parses_chunks_and_usage(respx_mock "stream": True, "stream_options": {"include_usage": True}, } - assert tuple(chunk.choices[0].delta.content for chunk in chunks) == ("Nova ", "stream", None, None) - assert tuple(chunk.choices[0].finish_reason for chunk in chunks) == (None, None, "stop", None) + assert tuple(chunk.choices[0].delta.content for chunk in chunks[:-1]) == ("Nova ", "stream", None) + assert chunks[-1].choices == [] + assert tuple(chunk.choices[0].finish_reason for chunk in chunks[:-1]) == (None, None, "stop") assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == ( 2, 1, diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index dc1e762e752..cf4a7c4cbc8 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -12,8 +12,10 @@ from urllib.parse import parse_qs, urlparse import httpx import pytest +from pydantic import TypeAdapter from unittest.mock import MagicMock +import litellm from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, _get_litellm_batch_custom_id_from_labels, @@ -1988,3 +1990,58 @@ class TestVertexEmbeddingsBatchOutputTranslation: content = json.dumps(legacy_row).encode("utf-8") assert config._try_transform_vertex_batch_output_to_openai(content) == content + + +_SAMPLING_KNOBS: Final = MappingProxyType({"temperature": 0, "top_p": 0.5, "top_k": 3}) + + +def _chat_entry(model: str, knobs: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType( + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5, **knobs}, + } + ) + + +def _generation_config(row: Mapping[str, object]) -> object: + request: Final = TypeAdapter(dict[str, object]).validate_python(row["request"]) + return request["generationConfig"] + + +def _flag_model_as_no_sampling(monkeypatch: pytest.MonkeyPatch, model: str) -> None: + monkeypatch.setitem(litellm.model_cost, model, {"supports_sampling_params": False}) + + +def test_batch_row_for_a_flagged_gemini_model_raises_a_400_naming_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: + """A batch file row gets the same contract as a real-time call: the knobs Google ignores are refused + with a 400 that names drop_params instead of silently vanishing from the uploaded row""" + _flag_model_as_no_sampling(monkeypatch, "gemini-zeta-9") + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError, match=r"temperature, top_p, top_k.*drop_params"): + _wrap_entries([_chat_entry("gemini-zeta-9", _SAMPLING_KNOBS)]) + + +def test_batch_row_with_null_knobs_for_a_flagged_gemini_model_is_not_refused(monkeypatch: pytest.MonkeyPatch) -> None: + """A null knob means no override, the same as on a real-time call, so it must not fail the upload""" + _flag_model_as_no_sampling(monkeypatch, "gemini-zeta-9") + monkeypatch.setattr(litellm, "drop_params", False) + (row,) = _wrap_entries([_chat_entry("gemini-zeta-9", {"temperature": None, "top_p": None, "top_k": None})]) + assert _generation_config(row) == {"max_output_tokens": 5} + + +def test_batch_row_for_a_flagged_gemini_model_drops_the_knobs_under_global_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _flag_model_as_no_sampling(monkeypatch, "gemini-zeta-9") + monkeypatch.setattr(litellm, "drop_params", True) + (row,) = _wrap_entries([_chat_entry("gemini-zeta-9", _SAMPLING_KNOBS)]) + assert _generation_config(row) == {"max_output_tokens": 5} + + +def test_batch_row_for_an_unflagged_gemini_model_keeps_the_knobs(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + (row,) = _wrap_entries([_chat_entry("gemini-2.5-flash", _SAMPLING_KNOBS)]) + assert _generation_config(row) == {"max_output_tokens": 5, **_SAMPLING_KNOBS} diff --git a/tests/unit/llms/vertex_ai/gemini/test_long_decode.py b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py index 4465ad7f3df..3fe9a5c7a92 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_long_decode.py +++ b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py @@ -516,9 +516,16 @@ async def test_unary_answer_cut_at_the_callers_cap_is_returned_with_its_token_in assert response.choices[0].message.provider_specific_fields["continuation_token"] == "dG9rLTE=" +@pytest.mark.parametrize( + "choice_count_request", + [{"max_tokens": 100_000, "n": 2}, {"extra_body": {"generationConfig": {"candidateCount": 2}}}], +) @pytest.mark.parametrize("is_async", [True, False]) @pytest.mark.asyncio -async def test_unary_request_for_more_than_one_choice_is_never_continued(is_async: bool): +async def test_unary_request_for_more_than_one_choice_is_never_continued( + monkeypatch: pytest.MonkeyPatch, is_async: bool, choice_count_request: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "100000") sent: Final[list[dict[str, object]]] = [] # Vertex AI answered 400 "Candidate count greater than 1 is not supported when continuation_token is set" # on 2026-10-13 @@ -526,9 +533,8 @@ async def test_unary_request_for_more_than_one_choice_is_never_continued(is_asyn kwargs: Final = { "model": MODEL, "messages": list(MESSAGES), - "max_tokens": 100_000, - "n": 2, "api_key": "fake-gemini-key", + **choice_count_request, } response: Final = ( @@ -548,6 +554,42 @@ async def test_unary_request_for_more_than_one_choice_is_never_continued(is_asyn ] +@pytest.mark.parametrize( + "one_choice_request", + [ + {"max_tokens": 100_000, "n": 0}, + {"extra_body": {"generationConfig": {"candidateCount": None}}}, + {"extra_body": {"generationConfig": {"candidateCount": 0}}}, + ], +) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_request_whose_candidate_count_google_reads_as_one_is_continued( + monkeypatch: pytest.MonkeyPatch, is_async: bool, one_choice_request: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "100000") + sent: Final[list[dict[str, object]]] = [] + hops: Final = (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2)) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "api_key": "fake-gemini-key", + **one_choice_request, + } + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + if is_async + else litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + ) + + assert sent[1] == {**sent[0], "continuationToken": "dG9rLTE="} + assert response.choices[0].message.content == "hello world" + assert response.choices[0].finish_reason == "stop" + + @pytest.mark.asyncio async def test_unary_failed_continuation_hop_returns_the_answer_so_far_with_its_token(): sent: Final[list[dict[str, object]]] = [] @@ -1295,6 +1337,468 @@ async def test_vertex_long_decode_keeps_the_callers_own_authorization_header_on_ assert seen == ["Bearer mine", "Bearer mine"] +BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" +ARGON: Final = "gemini/gemini-4-argon" +UNMAPPED_GEMINI: Final = "gemini/gemini-99-unreleased" +OUTPUT_LIMIT_KEYS: Final = ("max_output_tokens", "maxOutputTokens") + + +def _output_limits(request: Mapping[str, object]) -> dict[str, object]: + generation_config: Final = request.get("generationConfig", {}) + assert isinstance(generation_config, dict) + return {key: value for key, value in generation_config.items() if key in OUTPUT_LIMIT_KEYS} + + +def _cost_map_max_output_tokens(model: str) -> int: + provider, name = model.split("/", 1) + max_output_tokens: Final = litellm.get_model_info(model=name, custom_llm_provider=provider)["max_output_tokens"] + assert isinstance(max_output_tokens, int) + return max_output_tokens + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_request_with_no_output_limit_keeps_continuing_until_the_env_budget_is_spent( + monkeypatch: pytest.MonkeyPatch, is_async: bool, stream: bool +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + (_sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END), _sse(HOP_2_TRUNCATED), _sse(HOP_2)) + if stream + else ( + httpx.Response(200, json=HOP_1), + httpx.Response(200, json=HOP_2_TRUNCATED), + httpx.Response(200, json=HOP_2), + ) + ) + kwargs: Final = {"model": ARGON, "messages": list(MESSAGES), "stream": stream, "api_key": "fake-gemini-key"} + transport: Final = httpx.MockTransport(_replay(hops, sent)) + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport)) + if is_async + else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport))) + ) + chunks: Final = ([chunk async for chunk in response] if is_async else list(response)) if stream else [] + + assert len(sent) == 2 + assert _output_limits(sent[0]) == {"max_output_tokens": 50} + if stream: + assert _stream_text(chunks) == "hello world" + assert _stream_finish_reasons(chunks) == ["length"] + assert {"continuation_token": "dG9rLTI="} in _delta_fields(chunks) + else: + assert response.choices[0].message.content == "hello world" + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message.provider_specific_fields["continuation_token"] == "dG9rLTI=" + + +@pytest.mark.parametrize( + "budget, model, request_params, expected_limits", + [ + ("50", ARGON, {"max_tokens": 30}, {"max_output_tokens": 30}), + ("50", ARGON, {"max_completion_tokens": 30}, {"max_output_tokens": 30}), + ( + "50", + ARGON, + {"extra_body": {"generationConfig": {"maxOutputTokens": 30}}}, + {"maxOutputTokens": 30}, + ), + (None, ARGON, {}, {}), + ("", ARGON, {}, {}), + ("a lot", ARGON, {}, {}), + ("0", ARGON, {}, {}), + ("-5", ARGON, {}, {}), + ("50", UNMAPPED_GEMINI, {}, {}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_request_that_sets_no_output_limit_of_its_own( + monkeypatch: pytest.MonkeyPatch, + budget: str | None, + model: str, + request_params: Mapping[str, object], + expected_limits: Mapping[str, object], +): + if budget is None: + monkeypatch.delenv(BUDGET_ENV, raising=False) + else: + monkeypatch.setenv(BUDGET_ENV, budget) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + **request_params, + ) + + assert _output_limits(sent[0]) == expected_limits + + +@pytest.mark.parametrize( + "modalities, expected_limits", + [ + (["audio"], {}), + (["text", "audio"], {}), + (["text", "image"], {}), + (["text"], {"max_output_tokens": 50}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_request_asking_for_text_output( + monkeypatch: pytest.MonkeyPatch, modalities: list[str], expected_limits: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + modalities=modalities, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == expected_limits + + +def _ai_studio_model_in_mode(mode: str) -> str: + return next( + name + for name, row in litellm.model_cost.items() + if name.startswith("gemini/") and row.get("mode") == mode and row.get("max_output_tokens") + ) + + +@pytest.mark.parametrize("mode", ["audio_speech", "image_generation"]) +@pytest.mark.asyncio +async def test_env_budget_skips_a_model_the_cost_map_does_not_list_as_a_chat_model( + monkeypatch: pytest.MonkeyPatch, mode: str +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=_ai_studio_model_in_mode(mode), + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {} + + +@pytest.mark.parametrize( + "output_modalities, expected_limits", + [ + ({}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["text"]}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["text", "audio"]}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["audio"]}, {}), + ({"supported_output_modalities": "audio"}, {}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_chat_model_the_cost_map_lists_as_answering_text( + monkeypatch: pytest.MonkeyPatch, output_modalities: Mapping[str, object], expected_limits: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + litellm.register_model( + { + UNMAPPED_GEMINI: { + "litellm_provider": "gemini", + "mode": "chat", + "max_output_tokens": 1000, + **output_modalities, + } + } + ) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=UNMAPPED_GEMINI, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == expected_limits + + +@pytest.mark.parametrize( + "extra_body", + [ + {"generation_config": {"temperature": 0.1}}, + {"generation_config": {"max_output_tokens": 30}}, + {"generationConfig": {"response_modalities": ["AUDIO"]}}, + ], +) +@pytest.mark.asyncio +async def test_env_budget_leaves_a_snake_case_generation_config_from_extra_body_alone( + monkeypatch: pytest.MonkeyPatch, extra_body: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + extra_body=dict(extra_body), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {} + + +def _gemini_chat_models_answering_no_text() -> tuple[str, ...]: + rows: Final = ( + (name, row) + for name, row in litellm.model_cost.items() + if row.get("mode") == "chat" and "text" not in row.get("supported_output_modalities", ("text",)) + ) + return tuple( + name if name.startswith("gemini/") else f"vertex_ai/{name}" + for name, row in rows + if row.get("litellm_provider") in ("gemini", "vertex_ai-language-models") + ) + + +@pytest.mark.asyncio +async def test_env_budget_skips_every_chat_model_the_cost_map_lists_as_answering_no_text( + monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint +): + monkeypatch.setenv(BUDGET_ENV, "50") + models: Final = _gemini_chat_models_answering_no_text() + sent: Final[list[dict[str, object]]] = [] + replies: Final = tuple(httpx.Response(200, json=HOP_2) for _ in models) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model in models: + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + vertex_project="p", + vertex_location="us-central1", + vertex_credentials=token_endpoint.credentials, + client=client, + ) + + assert models + assert {model: _output_limits(request) for model, request in zip(models, sent, strict=True)} == dict.fromkeys( + models, {} + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_env_budget_leaves_a_continuation_token_resume_with_no_limit_alone( + monkeypatch: pytest.MonkeyPatch, stream: bool +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + reply: Final = _sse(HOP_2) if stream else httpx.Response(200, json=HOP_2) + + response: Final = await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + continuation_token="dG9rLTE=", + stream=stream, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((reply,), sent))), + ) + if stream: + _ = [chunk async for chunk in response] + + assert sent[0]["continuationToken"] == "dG9rLTE=" + assert _output_limits(sent[0]) == {} + + +@pytest.mark.parametrize("model", ["gemini/gemini-2.5-flash", "vertex_ai/gemini-2.5-flash", ARGON]) +@pytest.mark.asyncio +async def test_env_budget_above_a_models_output_limit_sends_that_limit_instead( + monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint, model: str +): + budget: Final = 2 * _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(budget)) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + vertex_project="p", + vertex_location="us-central1", + vertex_credentials=token_endpoint.credentials, + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": _cost_map_max_output_tokens(model)} + + +@pytest.mark.asyncio +async def test_env_budget_never_goes_above_the_cost_map_limit_when_a_deployment_overrides_it( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + model_limit: Final = _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(2 * model_limit)) + router: Final = litellm.Router( + model_list=[ + {"model_name": "plain", "litellm_params": {"model": model, "api_key": "fake-gemini-key"}}, + { + "model_name": "raised", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": 4 * model_limit}, + }, + ] + ) + sent: Final[list[dict[str, object]]] = [] + replies: Final = (httpx.Response(200, json=HOP_2), httpx.Response(200, json=HOP_2)) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model_name in ("plain", "raised"): + await router.acompletion(model=model_name, messages=list(MESSAGES), client=client) + + sent_limits: Final = [_output_limits(request).get("max_output_tokens", 0) for request in sent] + assert len(sent_limits) == 2 + assert all(isinstance(limit, int) and limit <= model_limit for limit in sent_limits) + + +@pytest.mark.asyncio +async def test_env_budget_caps_each_deployment_at_its_own_model_info_limit_and_never_at_a_siblings( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + budget: Final = _cost_map_max_output_tokens(model) // 2 + low_limit: Final = budget // 4 + monkeypatch.setenv(BUDGET_ENV, str(budget)) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "low", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": low_limit}, + }, + {"model_name": "plain", "litellm_params": {"model": model, "api_key": "fake-gemini-key"}}, + ] + ) + sent: Final[list[dict[str, object]]] = [] + replies: Final = (httpx.Response(200, json=HOP_2), httpx.Response(200, json=HOP_2)) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model_name in ("plain", "low"): + await router.acompletion(model=model_name, messages=list(MESSAGES), client=client) + + assert [_output_limits(request) for request in sent] == [ + {"max_output_tokens": budget}, + {"max_output_tokens": low_limit}, + ] + + +@pytest.mark.parametrize("deployment_limit", [0, -1]) +@pytest.mark.asyncio +async def test_env_budget_ignores_a_deployment_model_info_limit_that_is_not_positive( + monkeypatch: pytest.MonkeyPatch, deployment_limit: int +): + model: Final = "gemini/gemini-2.5-flash" + budget: Final = _cost_map_max_output_tokens(model) // 2 + monkeypatch.setenv(BUDGET_ENV, str(budget)) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "odd", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": deployment_limit}, + } + ] + ) + sent: Final[list[dict[str, object]]] = [] + + await router.acompletion( + model="odd", + messages=list(MESSAGES), + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": budget} + + +@pytest.mark.asyncio +async def test_env_budget_on_a_direct_call_never_goes_above_the_cost_map_limit_after_a_router_raised_it( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + model_limit: Final = _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(2 * model_limit)) + litellm.Router( + model_list=[ + { + "model_name": "raised", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": 4 * model_limit}, + } + ] + ) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": model_limit} + + +@pytest.mark.asyncio +async def test_env_budget_on_a_direct_call_follows_a_lower_limit_set_with_register_model( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + lowered_limit: Final = _cost_map_max_output_tokens(model) // 4 + monkeypatch.setenv(BUDGET_ENV, str(2 * lowered_limit)) + litellm.register_model({model: {"litellm_provider": "gemini", "max_output_tokens": lowered_limit}}) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": lowered_limit} + + +@pytest.mark.parametrize( + "generation_config", + [None, {"responseModalities": None}, {"responseModalities": {"TEXT": 1}}], +) +@pytest.mark.asyncio +async def test_env_budget_sends_a_generation_config_of_an_unexpected_shape_as_it_came( + monkeypatch: pytest.MonkeyPatch, generation_config: object +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + extra_body={"generationConfig": generation_config}, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert sent[0]["generationConfig"] == generation_config + + def test_initial_prompt_tokens_is_absent_on_a_single_hop_and_the_first_hops_count_otherwise(): assert initial_prompt_tokens_of((HOP_1,)) is None assert initial_prompt_tokens_of((HOP_1, HOP_2)) == 10 diff --git a/tests/unit/llms/vertex_ai/gemini/test_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_transformation.py index 2ea02b776ab..ac2d43f64a1 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_transformation.py @@ -1,12 +1,15 @@ import json +import logging import uuid +from typing import Final from unittest.mock import Mock import httpx import pytest import litellm +from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.gemini import transformation from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -469,3 +472,16 @@ async def test_vertex_ai_async_transform_inlines_only_the_urls_gemini_cannot_fet {"file_data": {"mime_type": "application/pdf", "file_uri": files_api_pdf}}, ] assert sorted(async_only_image_fetch.fetched) == sorted([plain_http_png, extensionless_https]) + + +def test_image_only_user_content_logs_the_blank_text_fixup_at_debug(caplog: pytest.LogCaptureFixture) -> None: + messages: Final = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}]} + ] + + with caplog.at_level(logging.DEBUG, logger=verbose_logger.name): + contents: Final = transformation.gemini_convert_messages_with_history(messages) + + fixups: Final = tuple(r for r in caplog.records if "No text in user content" in r.getMessage()) + assert tuple(r.levelno for r in fixups) == (logging.DEBUG,) + assert contents[0]["parts"][-1] == {"text": " "} diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 636d4034ccc..6ea3039edd8 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 @@ -1,7 +1,9 @@ import asyncio, importlib, os, uuid import json +import logging import re -from collections.abc import Awaitable, Callable, Mapping +import threading +from collections.abc import Awaitable, Callable, Iterator, Mapping from copy import deepcopy from typing import Final, List, Literal, cast, get_args from unittest.mock import MagicMock, Mock, patch @@ -13,7 +15,9 @@ from pydantic import BaseModel import litellm from litellm import ModelResponse, completion, embedding, image_generation +from litellm._logging import verbose_logger from litellm.litellm_core_utils.exception_mapping_utils import exception_type +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.llms.anthropic.pass_through.messages import handler as anthropic_messages_handler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig @@ -21,8 +25,9 @@ from litellm.llms.vertex_ai.common_utils import GEMINI_MAX_ENUM_CHARACTERS, Vert from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, + _log_gemini_sampling_params_deprecated, ) -from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata +from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -1889,7 +1894,9 @@ def test_gemini_tool_use(respx_mock: respx.MockRouter) -> None: assert [(call.function.name, json.loads(call.function.arguments)) for call in tool_calls] == [ ("get_weather", {"location": "Lima, Peru"}) ] - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["tool_calls"] + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [ + "tool_calls" + ] @pytest.mark.respx(assert_all_called=True) @@ -3734,6 +3741,169 @@ def test_get_optional_params_gemini_3_has_no_default_temperature( assert optional_params["thinkingConfig"]["thinkingLevel"] == "medium" +_SAMPLING_PARAMS: Final = {"temperature": 0.3, "top_p": 0.9, "top_k": 20} + + +def _gemini_config(config_cls: Literal["vertex", "studio"]) -> VertexGeminiConfig: + return VertexGeminiConfig() if config_cls == "vertex" else GoogleAIStudioGeminiConfig() + + +def _gemini_request_body( + custom_llm_provider: Literal["vertex_ai", "gemini"], + optional_params: dict[str, object], + litellm_params: dict[str, object], +) -> RequestBody: + from litellm.llms.vertex_ai.gemini.transformation import transform_request_body + + return transform_request_body( + messages=[{"role": "user", "content": "hi"}], + model="gemini-zeta-9", + optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + cached_content=None, + ) + + +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "gemini"]) +@pytest.mark.parametrize("param", ["temperature", "top_p"]) +def test_get_optional_params_raises_a_400_naming_drop_params_for_a_flagged_gemini_model( + custom_llm_provider: Literal["vertex_ai", "gemini"], param: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Google ignores sampling params on models flagged ``supports_sampling_params: false``, so without + drop_params the caller gets the 400 that names the escape hatch instead of a 200 that lied (#47015)""" + from litellm.utils import get_optional_params + + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setitem(litellm.model_cost, "gemini-zeta-9", {"supports_sampling_params": False}) + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params") as excinfo: + get_optional_params( + model="gemini-zeta-9", custom_llm_provider=custom_llm_provider, **{param: _SAMPLING_PARAMS[param]} + ) + + assert excinfo.value.status_code == 400 + assert param in str(excinfo.value) + + +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "gemini"]) +def test_allowed_openai_params_still_forwards_a_sampling_param_the_map_flags_off( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + """The 400 tells the caller to send allowed_openai_params, so that escape hatch has to work""" + from litellm.utils import get_optional_params + + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setitem(litellm.model_cost, "gemini-zeta-9", {"supports_sampling_params": False}) + + optional_params: Final = get_optional_params( + model="gemini-zeta-9", + custom_llm_provider=custom_llm_provider, + temperature=0.3, + allowed_openai_params=["temperature"], + ) + + assert optional_params["temperature"] == 0.3 + + +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "gemini"]) +def test_top_k_is_stripped_from_the_request_body_of_a_flagged_gemini_model_under_drop_params( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + """top_k is a provider-specific kwarg that skips the shared unsupported-param check, so it is gated + where the request body is built""" + monkeypatch.setitem(litellm.model_cost, "gemini-zeta-9", {"supports_sampling_params": False}) + + body: Final = _gemini_request_body( + custom_llm_provider, {"top_k": 20, "max_output_tokens": 10}, {"drop_params": True} + ) + + assert body["generationConfig"] == {"max_output_tokens": 10} + + +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "gemini"]) +def test_top_k_raises_a_400_naming_drop_params_for_a_flagged_gemini_model( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setitem(litellm.model_cost, "gemini-zeta-9", {"supports_sampling_params": False}) + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params") as excinfo: + _gemini_request_body(custom_llm_provider, {"top_k": 20}, {}) + + assert excinfo.value.status_code == 400 + assert "top_k" in str(excinfo.value) + + +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "gemini"]) +@pytest.mark.parametrize("row", [{"supports_sampling_params": True}, {}], ids=["flag_true", "flag_absent"]) +def test_gemini_keeps_forwarding_sampling_params_unless_the_model_map_flags_them_off( + custom_llm_provider: Literal["vertex_ai", "gemini"], row: dict[str, bool], monkeypatch: pytest.MonkeyPatch +) -> None: + """An explicit true and a model the map says nothing about both keep today's pass-through, top_k included""" + from litellm.utils import get_optional_params + + monkeypatch.setitem(litellm.model_cost, "gemini-zeta-9", row) + + optional_params: Final = get_optional_params( + model="gemini-zeta-9", custom_llm_provider=custom_llm_provider, drop_params=True, **_SAMPLING_PARAMS + ) + body: Final = _gemini_request_body(custom_llm_provider, dict(optional_params), {"drop_params": True}) + + assert body["generationConfig"] == _SAMPLING_PARAMS + + +@pytest.mark.parametrize("model_key", ["gemini-zeta-9", "vertex_ai/gemini-zeta-9", "gemini/gemini-zeta-9"]) +def test_gemini_sampling_flag_is_read_from_whichever_provider_row_carries_it( + model_key: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem(litellm.model_cost, model_key, {"supports_sampling_params": False}) + + result: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"temperature": 0.3}, + optional_params={}, + model="gemini-zeta-9", + drop_params=True, + ) + + assert "temperature" not in result + + +@pytest.mark.parametrize("config_cls", ["vertex", "studio"]) +def test_gemini_supported_openai_params_omit_sampling_params_when_flagged_off( + config_cls: Literal["vertex", "studio"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem(litellm.model_cost, "gemini-zeta-9", {"supports_sampling_params": False}) + + supported: Final = set(_gemini_config(config_cls).get_supported_openai_params("gemini-zeta-9")) + + assert not supported & set(_SAMPLING_PARAMS) + assert "max_tokens" in supported + + +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "gemini"]) +def test_get_optional_params_drops_sampling_params_for_a_flagged_gemini_model_under_drop_params( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + """The path every SDK and proxy call takes, on both the Vertex and the AI Studio route""" + from litellm.utils import get_optional_params + + monkeypatch.setitem(litellm.model_cost, "gemini-zeta-9", {"supports_sampling_params": False}) + + optional_params: Final = get_optional_params( + model="gemini-zeta-9", + custom_llm_provider=custom_llm_provider, + temperature=0.3, + top_p=0.9, + max_tokens=10, + drop_params=True, + ) + + assert "temperature" not in optional_params + assert "top_p" not in optional_params + assert optional_params["max_output_tokens"] == 10 + + def test_media_resolution_from_detail_parameter(): """Test that OpenAI's detail parameter is correctly mapped to media_resolution""" from litellm.llms.vertex_ai.gemini.transformation import ( @@ -6690,6 +6860,155 @@ def test_supported_reasoning_efforts_still_map(model): assert "thinkingConfig" in result +@pytest.fixture +def fresh_sampling_deprecation_cache() -> Iterator[None]: + _log_gemini_sampling_params_deprecated.cache_clear() + yield + _log_gemini_sampling_params_deprecated.cache_clear() + + +def _mapped_gemini_3_models() -> tuple[str, ...]: + return tuple(sorted(name for name in GetModelCostMap.loaded_model_cost_map() if name.startswith("gemini-3"))) + + +def _deprecation_messages(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple(r.getMessage() for r in caplog.records if "DeprecationWarning" in r.getMessage()) + + +@pytest.mark.parametrize(("param", "value"), (("temperature", 0.2), ("top_p", 0.9), ("top_k", 5))) +@pytest.mark.usefixtures("fresh_sampling_deprecation_cache") +def test_gemini_3_sampling_deprecation_is_logged_once_per_model( + param: str, value: float, caplog: pytest.LogCaptureFixture +) -> None: + first, second = _mapped_gemini_3_models()[:2] + config: Final = VertexGeminiConfig() + + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + mapped: Final = tuple( + config.map_openai_params( + non_default_params={param: value}, optional_params={}, model=model, drop_params=True + ) + for model in (first, first, first, second) + ) + + deprecations: Final = _deprecation_messages(caplog) + assert len(deprecations) == 2 + assert f"({first})" in deprecations[0] + assert f"({second})" in deprecations[1] + assert all(m[param] == value for m in mapped) + + +@pytest.mark.usefixtures("fresh_sampling_deprecation_cache") +def test_gemini_3_sampling_deprecation_is_logged_once_when_first_calls_race(caplog: pytest.LogCaptureFixture) -> None: + model: Final = _mapped_gemini_3_models()[0] + + def map_sampling_params() -> None: + VertexGeminiConfig().map_openai_params( + non_default_params={"temperature": 0.2}, optional_params={}, model=model, drop_params=True + ) + + racing_call: Final = threading.Thread(target=map_sampling_params) + + def start_racing_call_while_first_warning_logs(record: logging.LogRecord) -> bool: + if "DeprecationWarning" in record.getMessage() and racing_call.ident is None: + racing_call.start() + racing_call.join(timeout=0.5) + return True + + verbose_logger.addFilter(start_racing_call_while_first_warning_logs) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + map_sampling_params() + racing_call.join() + finally: + verbose_logger.removeFilter(start_racing_call_while_first_warning_logs) + + assert len(_deprecation_messages(caplog)) == 1 + + +@pytest.mark.usefixtures("fresh_sampling_deprecation_cache") +def test_gemini_3_sampling_deprecation_for_unmapped_names_does_not_hide_mapped_models( + caplog: pytest.LogCaptureFixture, +) -> None: + unmapped: Final = tuple(f"gemini-3.9-flash-{uuid.uuid4().hex}" for _ in range(100)) + mapped: Final = _mapped_gemini_3_models()[0] + config: Final = VertexGeminiConfig() + + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + for model in (*unmapped, mapped) * 2: + config.map_openai_params( + non_default_params={"temperature": 0.2}, optional_params={}, model=model, drop_params=True + ) + + deprecations: Final = _deprecation_messages(caplog) + assert len(deprecations) == 2 + assert all(name not in deprecations[0] for name in unmapped) + assert f"({mapped})" in deprecations[1] + + +@pytest.mark.usefixtures("fresh_sampling_deprecation_cache") +def test_gemini_3_sampling_deprecation_names_models_cataloged_only_under_vertex_ai( + caplog: pytest.LogCaptureFixture, +) -> None: + catalog: Final = GetModelCostMap.loaded_model_cost_map() + vertex_only: Final = next( + name.removeprefix("vertex_ai/") + for name in catalog + if name.startswith("vertex_ai/gemini-3") + and name.removeprefix("vertex_ai/") not in catalog + and f"gemini/{name.removeprefix('vertex_ai/')}" not in catalog + ) + + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + VertexGeminiConfig().map_openai_params( + non_default_params={"temperature": 0.2}, optional_params={}, model=vertex_only, drop_params=True + ) + + deprecations: Final = _deprecation_messages(caplog) + assert len(deprecations) == 1 + assert f"({vertex_only})" in deprecations[0] + + +@pytest.mark.usefixtures("fresh_sampling_deprecation_cache") +def test_gemini_3_sampling_deprecation_for_names_priced_at_runtime_shares_one_notice( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + priced: Final = tuple(f"gemini-3.9-flash-{uuid.uuid4().hex}" for _ in range(20)) + for name in priced: + monkeypatch.setitem( + litellm.model_cost, # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed + f"gemini/{name}", + {"input_cost_per_token": 1e-7, "output_cost_per_token": 1e-7, "litellm_provider": "gemini"}, + ) + config: Final = VertexGeminiConfig() + + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + for model in priced * 2: + config.map_openai_params( + non_default_params={"temperature": 0.2}, optional_params={}, model=model, drop_params=True + ) + + deprecations: Final = _deprecation_messages(caplog) + assert len(deprecations) == 1 + assert all(name not in deprecations[0] for name in priced) + + +@pytest.mark.parametrize(("param", "value"), (("temperature", 0.5), ("top_p", 0.9), ("top_k", 5))) +def test_gemini_2_5_sampling_params_log_no_deprecation( + param: str, value: float, caplog: pytest.LogCaptureFixture +) -> None: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + mapped: Final = VertexGeminiConfig().map_openai_params( + non_default_params={param: value}, + optional_params={}, + model=f"gemini-2.5-flash-{uuid.uuid4().hex}", + drop_params=True, + ) + + assert mapped[param] == value + assert not any("DeprecationWarning" in r.getMessage() for r in caplog.records) + + def _generate_content_body() -> dict: return { "candidates": [ diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index fb1cf4ced2d..0726ebb5ee1 100644 --- a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -632,7 +632,7 @@ class TestVertexGemmaCompletion: assert chunks[0].choices[0].delta.content == "READY" assert chunks[1].choices[0].finish_reason == "stop" if stream_options and stream_options["include_usage"]: - assert chunks[-1].choices[0].delta.content is None + assert chunks[-1].choices == [] assert chunks[-1].usage.total_tokens == 114 assert span.get_metric("_dd.llmobs.total_tokens") == 114 else: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 4896c2f1706..7315315a619 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -317,6 +317,204 @@ def test_validated_admission_jwt_is_not_forwarded_as_an_extra_header(auth_type: assert _openapi_forwarded_extra_headers(server, raw, user) == expected +@pytest.mark.parametrize( + "auth_type", + [MCPAuth.none, MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.true_passthrough], +) +def test_admission_authorization_is_not_forwarded_for_any_auth_type( + auth_type: MCPAuth, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_forwarded_extra_headers + from litellm.proxy._experimental.mcp_server.operations import _prepare_mcp_server_headers + + monkeypatch.setattr(proxy_server, "master_key", "test-master-key") + server: Final = MCPServer( + server_id="admission-header-boundary", + name="admission-header-boundary", + transport=MCPTransport.http, + auth_type=auth_type, + extra_headers=["Authorization", "x-tenant"], + ) + user: Final = UserAPIKeyAuth(api_key="hashed-key") + admitted_raw: Final = { + "Authorization": "Bearer sk-" + "1234", + "x-tenant": "tenant-control", + } + expected_without_admission_key: Final = {"x-tenant": "tenant-control"} + admission_upstream_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer sk-" + "1234"}, + raw_headers=admitted_raw, + user_api_key_auth=user, + )[1] + assert admission_upstream_headers == expected_without_admission_key + assert _openapi_forwarded_extra_headers(server, admitted_raw, user) == expected_without_admission_key + + separate_upstream_raw: Final = { + "x-litellm-api-key": "Bearer another-sk", + "Authorization": "Bearer upstream", + "x-tenant": "tenant-control", + } + expected_with_upstream_authorization: Final = { + "x-tenant": "tenant-control", + "Authorization": "Bearer upstream", + } + distinct_upstream_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer upstream"}, + raw_headers=separate_upstream_raw, + user_api_key_auth=user, + )[1] + assert distinct_upstream_headers == expected_with_upstream_authorization + assert _openapi_forwarded_extra_headers(server, separate_upstream_raw, user) == expected_with_upstream_authorization + + +@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.true_passthrough]) +def test_no_master_key_preserves_caller_authorization_unless_auth_is_enabled( + auth_type: MCPAuth, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_forwarded_extra_headers + from litellm.proxy._experimental.mcp_server.operations import _prepare_mcp_server_headers + + monkeypatch.setattr(proxy_server, "master_key", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + server: Final = MCPServer( + server_id="no-master-key", + name="no-master-key", + transport=MCPTransport.http, + auth_type=auth_type, + extra_headers=["Authorization", "x-tenant"], + ) + raw_headers: Final = {"Authorization": "Bearer upstream", "x-tenant": "t"} + user: Final = UserAPIKeyAuth(api_key="upstream", authenticated_by_custom_auth=False) + expected_forwarded: Final = {"x-tenant": "t", "Authorization": "Bearer upstream"} + forwarded_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer upstream"}, + raw_headers=raw_headers, + user_api_key_auth=user, + )[1] + assert forwarded_headers == expected_forwarded + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected_forwarded + + monkeypatch.setattr(proxy_server, "general_settings", {"enable_oauth2_auth": True}) + auth_enabled_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer upstream"}, + raw_headers=raw_headers, + user_api_key_auth=user, + )[1] + assert auth_enabled_headers == {"x-tenant": "t"} + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == {"x-tenant": "t"} + + +@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.true_passthrough]) +def test_custom_auth_is_stripped_in_no_auth_dev_mode( + auth_type: MCPAuth, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_forwarded_extra_headers + from litellm.proxy._experimental.mcp_server.operations import _prepare_mcp_server_headers + + monkeypatch.setattr(proxy_server, "master_key", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + server: Final = MCPServer( + server_id="custom-auth-no-master-key", + name="custom-auth-no-master-key", + transport=MCPTransport.http, + auth_type=auth_type, + extra_headers=["Authorization", "x-tenant"], + ) + raw_headers: Final = {"Authorization": "Bearer custom-tok", "x-tenant": "t"} + user: Final = UserAPIKeyAuth(api_key="custom-tok") + user.authenticated_by_custom_auth = True + expected_headers: Final = {"x-tenant": "t"} + forwarded_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer custom-tok"}, + raw_headers=raw_headers, + user_api_key_auth=user, + )[1] + + assert forwarded_headers == expected_headers + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected_headers + + +@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.true_passthrough]) +@pytest.mark.parametrize( + "raw_headers,oauth2_headers,expected", + [ + ( + { + "X-Company-Key": "sk-" + "1234", + "Authorization": "Bearer upstream", + "x-tenant": "t", + }, + {"Authorization": "Bearer upstream"}, + {"x-tenant": "t", "Authorization": "Bearer upstream"}, + ), + ( + { + "X-Company-Key": "sk-" + "1234", + "Authorization": "Bearer sk-" + "1234", + "x-tenant": "t", + }, + {"Authorization": "Bearer sk-" + "1234"}, + {"x-tenant": "t"}, + ), + ( + {"Authorization": "Bearer sk-" + "1234", "x-tenant": "t"}, + {"Authorization": "Bearer sk-" + "1234"}, + {"x-tenant": "t"}, + ), + ], +) +def test_custom_admission_header_controls_authorization_forwarding( + auth_type: MCPAuth, + raw_headers: dict[str, str], + oauth2_headers: dict[str, str], + expected: dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_forwarded_extra_headers + from litellm.proxy._experimental.mcp_server.operations import _prepare_mcp_server_headers + + monkeypatch.setattr(proxy_server, "master_key", "test-master-key") + monkeypatch.setattr(proxy_server, "general_settings", {"litellm_key_header_name": "x-company-key"}) + server: Final = MCPServer( + server_id="custom-admission-header", + name="custom-admission-header", + transport=MCPTransport.http, + auth_type=auth_type, + extra_headers=["Authorization", "x-tenant"], + ) + user: Final = UserAPIKeyAuth(api_key="hashed-key") + + upstream_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user, + )[1] + assert upstream_headers == expected + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected + + class TestMCPServerManager: """Test MCP Server Manager stdio functionality""" @@ -3559,14 +3757,14 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_call_regular_mcp_tool_passthrough_strips_authorization_when_admission_consumed_litellm_key( self, - ): + ) -> None: """OAuth pass-through must not forward the caller's Authorization to upstream when LiteLLM admission consumed the bearer as its API key — otherwise the LiteLLM key the caller used for admission would leak upstream.""" from litellm.proxy._types import UserAPIKeyAuth - manager = MCPServerManager() - server = MCPServer( + manager: Final = MCPServerManager() + server: Final = MCPServer( server_id="server-passthrough-call", name="passthrough-server", url="https://example.com", @@ -3576,20 +3774,19 @@ class TestMCPServerManager: oauth_passthrough=True, ) - mock_client = AsyncMock() + mock_client: Final = AsyncMock() mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) - captured_extra_headers = None + capture_extra_headers: Final = AsyncMock() async def capture_create_mcp_client( - server, - mcp_auth_header, - extra_headers, - stdio_env, - subject_token=None, - **kwargs, - ): # pragma: no cover - helper - nonlocal captured_extra_headers - captured_extra_headers = extra_headers + server: MCPServer, + mcp_auth_header: str | dict[str, str] | None, + extra_headers: dict[str, str] | None, + stdio_env: dict[str, str] | None, + subject_token: str | None = None, + **kwargs: object, + ) -> AsyncMock: + await capture_extra_headers(extra_headers) return mock_client manager.create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client) @@ -3610,8 +3807,64 @@ class TestMCPServerManager: user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), ) + assert capture_extra_headers.await_args is not None + captured_extra_headers: Final = capture_extra_headers.await_args.args[0] assert captured_extra_headers == {"x-request-id": "req-123"} + @pytest.mark.asyncio + async def test_call_regular_mcp_tool_never_forwards_admission_key_for_none_auth( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "master_key", "test-master-key") + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="server-none-admission-call", + name="none-auth-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization", "x-tenant"], + ) + + mock_client: Final = AsyncMock() + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) + capture_extra_headers: Final = AsyncMock() + + async def capture_create_mcp_client( + server: MCPServer, + mcp_auth_header: str | dict[str, str] | None, + extra_headers: dict[str, str] | None, + stdio_env: dict[str, str] | None, + subject_token: str | None = None, + **kwargs: object, + ) -> AsyncMock: + await capture_extra_headers(extra_headers) + return mock_client + + manager.create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client) + + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers={ + "Authorization": "Bearer sk-" + "1234", + "x-tenant": "tenant-control", + }, + proxy_logging_obj=None, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key"), + ) + + assert capture_extra_headers.await_args is not None + captured_extra_headers: Final = capture_extra_headers.await_args.args[0] + assert captured_extra_headers == {"x-tenant": "tenant-control"} + @pytest.mark.asyncio async def test_call_regular_mcp_tool_v2_authz_code_drops_caller_authorization( self, @@ -16193,6 +16446,7 @@ async def test_health_check_logs_upstream_exchange_only_for_http_failures( class _DiscoveryUpstream: def __init__(self) -> None: self.requests: tuple[tuple[str, str], ...] = () + self.tools: tuple[dict[str, object], ...] = () self.outcome = "supported" self.ttl_ms = 60000 self.entered = asyncio.Event() @@ -16254,7 +16508,7 @@ class _DiscoveryUpstream: "resources/templates/list": { "resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}] }, - "tools/list": {"tools": []}, + "tools/list": {"tools": list(self.tools)}, }[payload.method] continuation: Final = ( {"nextCursor": "last-page"} @@ -16291,6 +16545,64 @@ def _discovery_server() -> MCPServer: ) +@pytest.mark.asyncio +async def test_per_user_authorization_server_lists_without_static_credential() -> None: + upstream: Final = _DiscoveryUpstream() + upstream.tools = ({"name": "search", "inputSchema": {"type": "object"}},) + server: Final = MCPServer( + server_id="per-user-auth", + name="per-user-auth", + url="https://discovery.example/mcp", + transport=MCPTransport.http, + protocol_version="2026-07-28", + auth_type=MCPAuth.authorization, + ) + manager: Final = MCPServerManager() + manager.registry[server.server_id] = server + + with _mcp_upstream(upstream.respond): + tools: Final = await manager.get_tools_from_server(server, add_prefix=False) + + assert [tool.name for tool in tools] == ["search"] + assert ("tools/list", "") in upstream.requests + with pytest.raises(HTTPException, match="requires a usable upstream credential"): + await manager.create_mcp_client(server) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("auth_type", "credential", "upstream_token_header"), + ( + (MCPAuth.authorization, "Bearer", None), + (MCPAuth.authorization, "Bearer", "X-Token"), + (MCPAuth.authorization, " ", None), + (MCPAuth.authorization, " ", "X-Token"), + (MCPAuth.bearer_token, None, None), + ), +) +async def test_tool_listing_rejects_malformed_or_other_missing_static_credentials( + auth_type: MCPAuth, credential: str | None, upstream_token_header: str | None +) -> None: + upstream: Final = _DiscoveryUpstream() + server: Final = MCPServer( + server_id="invalid-static-auth", + name="invalid-static-auth", + url="https://discovery.example/mcp", + transport=MCPTransport.http, + protocol_version="2026-07-28", + auth_type=auth_type, + authentication_token=credential, + upstream_token_header=upstream_token_header, + ) + manager: Final = MCPServerManager() + manager.registry[server.server_id] = server + + with _mcp_upstream(upstream.respond), pytest.raises(MCPServerListError): + await manager.get_tools_from_server(server) + + assert upstream.requests == () + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", ("prompts", "resources", "templates")) async def test_discovery_cache_reuses_raw_results_and_expires(kind: str) -> None: diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 67ff7dd9344..ee1ae43fd3a 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3902,6 +3902,38 @@ class TestIsRequestBodySafeBlocksNVCFFunctionOverride: ) +def test_riva_offline_in_request_body_is_rejected() -> None: + with pytest.raises(ValueError, match="riva_offline"): + is_request_body_safe( + request_body={"model": "nvidia_riva/parakeet", "riva_offline": True}, + general_settings={}, + llm_router=None, + model="nvidia_riva/parakeet", + ) + + +def test_admin_opt_in_proxy_wide_allows_riva_offline() -> None: + assert ( + is_request_body_safe( + request_body={"model": "nvidia_riva/parakeet", "riva_offline": True}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="nvidia_riva/parakeet", + ) + is True + ) + + +def test_api_base_override_keeps_the_deployments_riva_offline_mode() -> None: + from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params + + out: Final = get_dynamic_litellm_params( + litellm_params={"model": "nvidia_riva/nvidia/parakeet-0.6b-tdt", "riva_offline": True}, + request_kwargs={"api_base": "caller-riva:50051"}, + ) + assert (out["api_base"], out["riva_offline"]) == ("caller-riva:50051", True) + + class TestIsRequestBodySafeBlocksRivaUseSsl: """``use_ssl`` is rejected as a request-body param unless the admin opted in proxy-wide or per-deployment.""" @@ -5088,7 +5120,7 @@ def test_get_model_from_request_vertex_ai_passthrough(request_data, route, expec @pytest.mark.parametrize( "field", - ["github_copilot_auth_type", "user_provider_credentials", "github_copilot_user_session"], + ["github_copilot_auth_type", "github_copilot_host", "user_provider_credentials", "github_copilot_user_session"], ) def test_per_user_credential_slots_in_request_body_are_rejected(field: str) -> None: """The per-user Copilot mode is decided by the stored credential's values and the caller's diff --git a/tests/unit/proxy/client/cli/test_pi.py b/tests/unit/proxy/client/cli/test_pi.py index ff4dc25cd0c..e6669564358 100644 --- a/tests/unit/proxy/client/cli/test_pi.py +++ b/tests/unit/proxy/client/cli/test_pi.py @@ -195,7 +195,12 @@ class TestProviderBlock: assert block == { "baseUrl": "http://localhost:4000/v1", "api": "openai-completions", - "compat": {"supportsStore": False, "supportsLongCacheRetention": False}, + "compat": { + "supportsStore": False, + "supportsLongCacheRetention": False, + "sendSessionAffinityHeaders": True, + "sessionAffinityFormat": "openrouter", + }, "apiKey": "$LITELLM_PROXY_API_KEY", "models": [{"id": "m-1"}, {"id": "m-2"}], } diff --git a/tests/unit/proxy/common_utils/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py index d78c3ef6185..b8d5ddef3d2 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -4,6 +4,8 @@ Unit tests for CheckResponsesCost class import asyncio import dataclasses +import hashlib +import json from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime @@ -127,10 +129,21 @@ class _InMemoryManagedObjectTable: return len(matched) +def _prisma_with_no_creator_rows() -> MagicMock: + prisma_client: Final = MagicMock() + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_deprecatedverificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_deletedteamtable.find_first = AsyncMock(return_value=None) + return prisma_client + + def _poller_over(table: _InMemoryManagedObjectTable, billed_logger: _RecordingBilledLogger): from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost - prisma_client: Final = MagicMock() + prisma_client: Final = _prisma_with_no_creator_rows() prisma_client.db.litellm_managedobjecttable = table poller: Final = CheckResponsesCost( proxy_logging_obj=MagicMock(), @@ -174,8 +187,7 @@ class TestCheckResponsesCost: @pytest.fixture def mock_prisma_client(self): """Create a mock Prisma client""" - client = MagicMock() - client.db = MagicMock() + client = _prisma_with_no_creator_rows() client.db.litellm_managedobjecttable = MagicMock() return client @@ -1702,26 +1714,60 @@ class _PricedBilling: table: _InMemoryManagedObjectTable response_id: str cost: float | None + logging_obj: object + + +@dataclass(frozen=True) +class _PricedPoller: + poller: object + prisma_client: MagicMock + table: _InMemoryManagedObjectTable + response_id: str + + async def bill(self, read: ResponsesAPIResponse) -> object: + from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + assert isinstance(self.poller, CheckResponsesCost) + charge: Final = await self.poller.bill_if_finished(response_id=self.response_id, response=read) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + return charge _PRICED_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType( {"model": "openai/gpt-5.6", "api_key": "sk-unused", "input_cost_per_token": 0.001, "output_cost_per_token": 0.002} ) +_CREATING_KEY_HASH: Final = hashlib.sha256(b"sk-alice").hexdigest() -async def _bill_through_a_priced_deployment( - read: ResponsesAPIResponse, +def _team_alpha_when_its_rows_exist( + team_row: object, deleted_team_row: object, default_team_settings: tuple[Mapping[str, object], ...] +) -> str | None: + return "team-alpha" if team_row is not None or deleted_team_row is not None or default_team_settings else None + + +def _priced_poller( litellm_params: Mapping[str, object] = _PRICED_LITELLM_PARAMS, model_info: Mapping[str, object] = MappingProxyType({}), + key_row: object = None, + team_row: object = None, + key_lookup_error: Exception | None = None, + default_team_settings: tuple[Mapping[str, object], ...] = (), + api_key: str = _CREATING_KEY_HASH, + rotated_to: Mapping[str, object] = MappingProxyType({}), + deleted_key_row: object = None, + deleted_team_row: object = None, model_name: str = "priced-gpt", requested_model: str = "priced-gpt", model_group_alias: Mapping[str, str] = MappingProxyType({}), team_id: str | None = None, -) -> _PricedBilling: +) -> _PricedPoller: + from types import SimpleNamespace + import litellm from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.responses.utils import ResponsesAPIRequestUtils response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( @@ -1733,17 +1779,31 @@ async def _bill_through_a_priced_deployment( id="job-priced", unified_object_id=response_id, status="queued", - api_key="hash-alice", + api_key=api_key, file_object={"model": requested_model}, - team_id=team_id, + team_id=team_id or _team_alpha_when_its_rows_exist(team_row, deleted_team_row, default_team_settings), ), ) ) - prisma_client: Final = MagicMock() + key_rows: Final = {api_key: key_row, **{token: row for token, row in rotated_to.items()}} + prisma_client: Final = _prisma_with_no_creator_rows() prisma_client.db.litellm_managedobjecttable = table prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) - prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) - prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=key_lookup_error or (lambda where, **_: key_rows.get(where["token"])) + ) + rotation_chain: Final = (api_key, *rotated_to) + rotated_from: Final = dict(zip(rotation_chain, rotation_chain[1:])) + prisma_client.db.litellm_deprecatedverificationtoken.find_unique = AsyncMock( + side_effect=lambda where: ( + SimpleNamespace(token=where["token"], active_token_id=rotated_from[where["token"]]) + if where["token"] in rotated_from + else None + ) + ) + prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=deleted_key_row) + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + prisma_client.db.litellm_deletedteamtable.find_first = AsyncMock(return_value=deleted_team_row) router: Final = litellm.Router( model_list=[ { @@ -1754,11 +1814,71 @@ async def _bill_through_a_priced_deployment( ], model_group_alias=dict(model_group_alias), ) - charge: Final = await CheckResponsesCost( - proxy_logging_obj=MagicMock(), prisma_client=prisma_client, llm_router=router - ).bill_if_finished(response_id=response_id, response=read) - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) - return _PricedBilling(table=table, response_id=response_id, cost=charge.cost if charge is not None else None) + return _PricedPoller( + poller=CheckResponsesCost( + proxy_logging_obj=MagicMock(), + prisma_client=prisma_client, + llm_router=router, + proxy_config=_proxy_config_with(default_team_settings), + ), + prisma_client=prisma_client, + table=table, + response_id=response_id, + ) + + +async def _bill_through_a_priced_deployment( + read: ResponsesAPIResponse, + litellm_params: Mapping[str, object] = _PRICED_LITELLM_PARAMS, + model_info: Mapping[str, object] = MappingProxyType({}), + key_row: object = None, + team_row: object = None, + key_lookup_error: Exception | None = None, + default_team_settings: tuple[Mapping[str, object], ...] = (), + api_key: str = _CREATING_KEY_HASH, + rotated_to: Mapping[str, object] = MappingProxyType({}), + deleted_key_row: object = None, + deleted_team_row: object = None, + model_name: str = "priced-gpt", + requested_model: str = "priced-gpt", + model_group_alias: Mapping[str, str] = MappingProxyType({}), + team_id: str | None = None, +) -> _PricedBilling: + priced: Final = _priced_poller( + litellm_params=litellm_params, + model_info=model_info, + key_row=key_row, + team_row=team_row, + key_lookup_error=key_lookup_error, + default_team_settings=default_team_settings, + api_key=api_key, + rotated_to=rotated_to, + deleted_key_row=deleted_key_row, + deleted_team_row=deleted_team_row, + model_name=model_name, + requested_model=requested_model, + model_group_alias=model_group_alias, + team_id=team_id, + ) + charge: Final = await priced.bill(read) + return _PricedBilling( + table=priced.table, + response_id=priced.response_id, + cost=getattr(charge, "cost", None), + logging_obj=getattr(charge, "logging_obj", None), + ) + + +def _proxy_config_with(default_team_settings: tuple[Mapping[str, object], ...]): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config: Final = ProxyConfig() + proxy_config.config = ( + {"litellm_settings": {"default_team_settings": [dict(settings) for settings in default_team_settings]}} + if default_team_settings + else {} + ) + return proxy_config @pytest.mark.asyncio @@ -1780,7 +1900,7 @@ async def test_the_billing_event_prices_the_fetched_usage_at_the_deployment_rate assert billing.cost == pytest.approx(_PRICED_COST) assert payload["model_id"] == "deployment-custom" assert payload["model_group"] == "priced-gpt" - assert payload["metadata"]["user_api_key_hash"] == "hash-alice" + assert payload["metadata"]["user_api_key_hash"] == _CREATING_KEY_HASH assert read.hidden_params == {"response_cost": 0.0} assert billing.table.statuses == {"job-priced": "completed"} @@ -1945,3 +2065,428 @@ async def test_the_billing_event_keys_its_spend_row_apart_from_the_read_it_bills assert billed_request_id != billing.response_id assert billed_request_id.startswith(billing.response_id) + +_CREATOR_LANGFUSE_VARS: Final = MappingProxyType( + { + "langfuse_public_key": "pk-lf-creator", + "langfuse_secret_key": "sk-lf-creator", + "langfuse_host": "http://127.0.0.1:9", + } +) + + +def _key_row(metadata: Mapping[str, object]) -> object: + from types import SimpleNamespace + + return SimpleNamespace(key_alias=None, organization_id=None, spend=0.0, max_budget=None, metadata=dict(metadata)) + + +def _team_row(metadata: Mapping[str, object]) -> object: + from types import SimpleNamespace + + return SimpleNamespace(team_alias=None, organization_id=None, spend=0.0, max_budget=None, metadata=dict(metadata)) + + +def _langfuse_callback_settings(callback_vars: Mapping[str, object]) -> Mapping[str, object]: + return {"callback_settings": {"success_callback": ["langfuse"], "callback_vars": dict(callback_vars)}} + + +@dataclass(frozen=True) +class _CreatorLogging: + key_row: object + team_row: object + default_team_settings: tuple[Mapping[str, object], ...] + + +def _creator_logging(logging_on: str) -> _CreatorLogging: + logging: Final = { + "logging": [ + {"callback_name": "langfuse", "callback_type": "success", "callback_vars": dict(_CREATOR_LANGFUSE_VARS)} + ] + } + team_metadata: Final = { + "team logging": logging, + "team callback_settings": _langfuse_callback_settings(_CREATOR_LANGFUSE_VARS), + } + return _CreatorLogging( + key_row=_key_row(logging if logging_on == "key logging" else {}), + team_row=_team_row(team_metadata.get(logging_on, {})), + default_team_settings=( + ({"team_id": "team-alpha", "success_callback": ["langfuse"], **_CREATOR_LANGFUSE_VARS},) + if logging_on == "config team settings" + else () + ), + ) + + +def _recording_langfuse(): + from litellm.integrations.custom_logger import CustomLogger + from litellm.integrations.langfuse.langfuse_prompt_management import LangfusePromptManagement + + class _RecordingLangfuse(LangfusePromptManagement): + def __init__(self) -> None: + CustomLogger.__init__(self) + self.events: tuple[tuple[object, object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + params: Final = kwargs["standard_callback_dynamic_params"] + assert isinstance(params, Mapping) + self.events = (*self.events, (kwargs.get("response_cost"), params.get("langfuse_public_key"))) + + return _RecordingLangfuse() + + +async def _log_a_request_from(creator: _CreatorLogging) -> None: + import litellm + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields + + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello", + **key_and_team_logging_fields( + user_api_key_dict=UserAPIKeyAuth( + metadata=getattr(creator.key_row, "metadata"), + team_id="team-alpha", + team_metadata=getattr(creator.team_row, "metadata"), + ), + proxy_config=_proxy_config_with(creator.default_team_settings), + ), + ) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_logs_to_langfuse", [False, True]) +@pytest.mark.parametrize( + "logging_on", ["key logging", "team logging", "team callback_settings", "config team settings"] +) +async def test_the_billing_event_reaches_the_langfuse_projects_the_creators_own_requests_log_to( + monkeypatch: pytest.MonkeyPatch, logging_on: str, proxy_logs_to_langfuse: bool +): + """A key or team that logs to its own Langfuse project must see the billed charge in the same projects + its own requests log to, wherever that logging is configured and whether or not the proxy logs to + Langfuse too, or the charge only shows in the proxy's spend logs.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + + langfuse: Final = _recording_langfuse() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [langfuse] if proxy_logs_to_langfuse else []) + creator: Final = _creator_logging(logging_on) + + await _log_a_request_from(creator) + request_projects: Final = tuple(public_key for _, public_key in langfuse.events) + await _bill_through_a_priced_deployment( + _terminal_response(), + key_row=creator.key_row, + team_row=creator.team_row, + default_team_settings=creator.default_team_settings, + ) + billed_events: Final = langfuse.events[len(request_projects) :] + + assert _CREATOR_LANGFUSE_VARS["langfuse_public_key"] in request_projects + assert tuple(public_key for _, public_key in billed_events) == request_projects + assert all(cost == pytest.approx(_PRICED_COST) for cost, _ in billed_events) + + +def _rotations_for(retired: str, key_logging: object) -> Mapping[str, object]: + rotated_once: Final = hashlib.sha256(b"sk-alice-rotated").hexdigest() + if retired == "rotated key": + return {rotated_once: key_logging} + if retired == "key rotated twice": + return {rotated_once: None, hashlib.sha256(b"sk-alice-rotated-again").hexdigest(): key_logging} + return {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retired", ["rotated key", "key rotated twice", "deleted key", "deleted team"]) +async def test_a_creator_rotated_or_deleted_before_the_bill_still_gets_the_charge_in_its_own_langfuse_project( + monkeypatch: pytest.MonkeyPatch, retired: str +): + """A key rotated or deleted, or a team deleted, while its background response ran logged that response's + create to its own project, so the charge follows the settings kept on the key it was rotated to or on + its deleted row instead of leaving that project for the proxy's callbacks alone.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + + langfuse: Final = _recording_langfuse() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + key_logging: Final = _creator_logging("key logging").key_row + + await _bill_through_a_priced_deployment( + _terminal_response(), + rotated_to=_rotations_for(retired, key_logging), + deleted_key_row=key_logging if retired == "deleted key" else SimpleNamespace(metadata={}), + deleted_team_row=_creator_logging("team callback_settings").team_row if retired == "deleted team" else None, + ) + + assert langfuse.events == ((pytest.approx(_PRICED_COST), _CREATOR_LANGFUSE_VARS["langfuse_public_key"]),) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failing_lookup, live_key_logging, team_logging", + [ + ("litellm_verificationtoken.find_unique", "key logging", "none"), + ("litellm_teamtable.find_unique", "none", "team callback_settings"), + ("litellm_deprecatedverificationtoken.find_unique", "none", "none"), + ("litellm_deletedverificationtoken.find_first", "none", "none"), + ], + ids=["live key", "live team", "rotated key", "deleted key"], +) +async def test_a_creator_lookup_that_keeps_failing_still_bills_the_response_once_without_the_creators_settings( + monkeypatch: pytest.MonkeyPatch, failing_lookup: str, live_key_logging: str, team_logging: str +): + """A creator lookup that fails on every read must not leave the response unbilled. It bills once under the + key and team the managed row stores, through the proxy's callbacks only, with the text left out since the + creator's redaction setting could not be read.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + + langfuse: Final = _recording_langfuse() + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + priced: Final = _priced_poller( + key_row=_creator_logging(live_key_logging).key_row if live_key_logging != "none" else None, + team_row=_creator_logging(team_logging).team_row if team_logging != "none" else None, + ) + table_name, method = failing_lookup.split(".") + setattr( + getattr(priced.prisma_client.db, table_name), + method, + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + + charge: Final = await priced.bill(_read_with_text("the stored answer")) + await priced.bill(_read_with_text("the stored answer")) + + assert getattr(charge, "cost", None) == pytest.approx(_PRICED_COST) + assert priced.table.statuses == {"job-priced": "completed"} + assert langfuse.events == () + (payload,) = capture.payloads + assert payload["response_cost"] == pytest.approx(_PRICED_COST) + assert "the stored answer" not in json.dumps(payload, default=str) + metadata: Final = payload["metadata"] + assert isinstance(metadata, Mapping) + assert metadata["user_api_key_hash"] == _CREATING_KEY_HASH + assert metadata["user_api_key_team_id"] == ("team-alpha" if team_logging != "none" else None) + + +_TEAM_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 2.0, "budget_duration": "7d"}} +_LOADED_KEY_ATTRIBUTION: Final = {"user_api_key_alias": "prod-key", "user_api_key_max_budget": 5.0} +_LOADED_TEAM_ATTRIBUTION: Final = { + "user_api_key_team_alias": "Team Alpha", + "user_api_key_team_max_budget": 50.0, + "user_api_key_team_model_max_budget": _TEAM_MODEL_BUDGET, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failing_lookup, key_is_live, loaded_attribution", + [ + ("litellm_verificationtoken.find_unique", True, _LOADED_TEAM_ATTRIBUTION), + ("litellm_teamtable.find_unique", True, _LOADED_KEY_ATTRIBUTION), + ("litellm_deprecatedverificationtoken.find_unique", False, _LOADED_TEAM_ATTRIBUTION), + ("litellm_deletedverificationtoken.find_first", False, _LOADED_TEAM_ATTRIBUTION), + ], + ids=["live key", "live team", "rotated key", "deleted key"], +) +async def test_a_creator_lookup_that_fails_keeps_the_aliases_and_budgets_of_the_rows_that_did_load( + failing_lookup: str, key_is_live: bool, loaded_attribution: Mapping[str, object] +): + """One failed creator lookup must not strip what the other rows give the charge: the team's per-model budget + only counts a billed charge whose event carries it, and the aliases name who spent it.""" + table: Final = _InMemoryManagedObjectTable( + rows=( + _StoredResponse( + id="job-attributed", + unified_object_id="resp_attributed", + status="queued", + api_key=_CREATING_KEY_HASH, + team_id="team-alpha", + ), + ) + ) + billed_logger: Final = _RecordingBilledLogger() + poller: Final = _poller_over(table, billed_logger) + poller.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=( + SimpleNamespace(key_alias="prod-key", organization_id=None, spend=4.0, max_budget=5.0, metadata={}) + if key_is_live + else None + ) + ) + poller.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=SimpleNamespace( + team_alias="Team Alpha", + organization_id=None, + spend=40.0, + max_budget=50.0, + model_max_budget=_TEAM_MODEL_BUDGET, + metadata={}, + ) + ) + table_name, method = failing_lookup.split(".") + setattr( + getattr(poller.prisma_client.db, table_name), + method, + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + + await poller.bill_if_finished(response_id="resp_attributed", response=_terminal_response()) + + assert table.statuses == {"job-attributed": "completed"} + (billed,) = billed_logger.billed + assert {field: billed.metadata.get(field) for field in loaded_attribution} == loaded_attribution + + +@pytest.mark.asyncio +async def test_a_team_logging_setting_the_logger_rejects_still_bills_the_response_once(monkeypatch: pytest.MonkeyPatch): + """A team whose stored callback settings the logger rejects, such as an os.environ reference it no + longer resolves, must not stop the bill: the charge reaches the proxy's callbacks and the response is + not put back for every later poll to fail on again.""" + import litellm + + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + + billing: Final = await _bill_through_a_priced_deployment( + _terminal_response(), + key_row=_key_row({}), + team_row=_team_row(_langfuse_callback_settings({"langfuse_public_key": "os.environ/QA_UNSET_LANGFUSE_KEY"})), + ) + + (payload,) = capture.payloads + assert payload["response_cost"] == pytest.approx(_PRICED_COST) + assert billing.cost == pytest.approx(_PRICED_COST) + assert billing.table.statuses == {"job-priced": "completed"} + + +def _read_with_text(text: str) -> ResponsesAPIResponse: + return _terminal_response().model_copy( + update={ + "output": [ + { + "type": "message", + "id": "msg_stored", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ] + } + ) + + +def _key_logging_to_langfuse(callback_vars: Mapping[str, object]) -> object: + return _key_row( + {"logging": [{"callback_name": "langfuse", "callback_type": "success", "callback_vars": dict(callback_vars)}]} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("key_row", "deleted_key_row", "text_logged"), + [ + (_key_logging_to_langfuse(_CREATOR_LANGFUSE_VARS), None, True), + (_key_logging_to_langfuse({**_CREATOR_LANGFUSE_VARS, "turn_off_message_logging": "true"}), None, False), + (None, _key_logging_to_langfuse({**_CREATOR_LANGFUSE_VARS, "turn_off_message_logging": "true"}), False), + (None, None, True), + ], + ids=[ + "creator keeps message logging", + "creator turns message logging off", + "deleted creator turned message logging off", + "creator has no key row anywhere", + ], +) +async def test_the_billing_event_logs_the_response_text_only_as_the_creators_logging_settings_allow( + monkeypatch: pytest.MonkeyPatch, key_row: object, deleted_key_row: object, text_logged: bool +): + """The billed event follows the creating key's or team's message logging setting, a deleted key's included, + which falls back to the proxy's the way the batch poll decides it, so a key with no stored settings + anywhere logs the text exactly as the proxy is set.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [_recording_langfuse()]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + + await _bill_through_a_priced_deployment( + _read_with_text("the stored answer"), key_row=key_row, deleted_key_row=deleted_key_row + ) + + (payload,) = capture.payloads + assert payload["response_cost"] == pytest.approx(_PRICED_COST) + assert ("the stored answer" in json.dumps(payload, default=str)) is text_logged + + +@pytest.mark.asyncio +@pytest.mark.parametrize("caller_token", ["oauth2-raw-access-token", "sk-custom-auth-1234"]) +async def test_a_response_created_without_a_virtual_key_bills_into_its_teams_langfuse_project( + monkeypatch: pytest.MonkeyPatch, caller_token: str +): + """Under OAuth2 the row keeps the caller's raw token, and under custom auth the hash of an sk- token that + has no key row, so the charge follows the team's logging settings like that caller's own requests.""" + from litellm.proxy._types import UserAPIKeyAuth + + import litellm + from litellm.litellm_core_utils import litellm_logging + + langfuse: Final = _recording_langfuse() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + + await _bill_through_a_priced_deployment( + _terminal_response(), + api_key=UserAPIKeyAuth(api_key=caller_token).api_key or caller_token, + team_row=_creator_logging("team callback_settings").team_row, + ) + + assert langfuse.events == ((pytest.approx(_PRICED_COST), _CREATOR_LANGFUSE_VARS["langfuse_public_key"]),) + + +@pytest.mark.asyncio +async def test_a_key_that_disabled_a_callback_has_its_charge_logged_wherever_its_own_requests_are( + monkeypatch: pytest.MonkeyPatch, +): + """A key's stored list of disabled callbacks must treat the billed charge like that key's own requests, or + a project that shows the key's requests at $0 never shows the charge they add up to.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", True) + langfuse: Final = _recording_langfuse() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [langfuse]) + creator: Final = _CreatorLogging( + key_row=_key_row({"litellm_disabled_callbacks": ["langfuse"]}), team_row=_team_row({}), default_team_settings=() + ) + + await _log_a_request_from(creator) + request_events: Final = langfuse.events + await _bill_through_a_priced_deployment(_terminal_response(), key_row=creator.key_row) + billed_events: Final = langfuse.events[len(request_events) :] + + assert len(billed_events) == len(request_events) + assert all(cost == pytest.approx(_PRICED_COST) for cost, _ in billed_events) diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 6aeceaf4356..c905d38dd41 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -4,7 +4,7 @@ import json from contextlib import contextmanager from types import SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import httpx import pytest @@ -14,11 +14,13 @@ from fastapi import HTTPException, Request from fastapi.testclient import TestClient import litellm -from litellm.models.credentials import CredentialSource +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE, _session_cache_key +from litellm.models.credentials import CredentialSource, UserConnectionFlowHandle from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper -from litellm.proxy.credential_endpoints.endpoints import get_llm_router +from litellm.proxy.credential_endpoints.endpoints import _decode_flow_handle, get_llm_router from litellm.proxy.proxy_server import app from litellm.types.utils import CredentialItem @@ -1428,6 +1430,298 @@ def test_create_credential_still_answers_200_for_a_name_that_is_free(credential_ find_by_name.assert_not_awaited(), "the unique index is the guard; create must not add a lookup" +@pytest.mark.parametrize( + "host", + ["https://x.ghe.com", "ghes.example.com", "x.ghe.com:443", "a.b.ghe.com"], +) +def test_create_credential_rejects_invalid_github_copilot_hosts(credential_store, host): + create: Final = AsyncMock() + credential_store(create=create) + + response: Final = _create_credential( + { + "credential_name": "copilot", + "credential_values": {"github_copilot_host": host}, + "credential_info": {"custom_llm_provider": "github_copilot"}, + } + ) + + assert response.status_code == 400 + create.assert_not_awaited() + + +def test_create_credential_normalizes_github_copilot_host(credential_store): + credential_store(create=AsyncMock()) + + response = _create_credential( + { + "credential_name": "copilot", + "credential_values": {"github_copilot_host": "OCTO.ghe.com"}, + "credential_info": {"custom_llm_provider": "github_copilot"}, + } + ) + + assert response.status_code == 200, response.text + assert CredentialAccessor.find_credential("copilot").credential_values["github_copilot_host"] == "octo.ghe.com" + + +@pytest.mark.parametrize( + "host", + ["https://x.ghe.com", "ghes.example.com", "x.ghe.com:443", "a.b.ghe.com"], +) +def test_update_credential_rejects_invalid_github_copilot_hosts(credential_store, host): + stored: Final = CredentialItem( + credential_name="copilot", + credential_values={"github_copilot_host": "octo.ghe.com"}, + credential_info={"custom_llm_provider": "github_copilot"}, + ) + update: Final = AsyncMock() + credential_store(in_memory=(stored,), find_by_name=AsyncMock(return_value=stored), update_by_name=update) + + response = _patch_credential( + "copilot", + { + "credential_name": "copilot", + "credential_values": {"github_copilot_host": host}, + "credential_info": {"custom_llm_provider": "github_copilot"}, + }, + ) + + assert response.status_code == 400 + update.assert_not_awaited() + + +def test_update_credential_normalizes_github_copilot_host(credential_store): + stored = CredentialItem( + credential_name="copilot", + credential_values={}, + credential_info={"custom_llm_provider": "github_copilot"}, + ) + credential_store( + in_memory=(stored,), + find_by_name=AsyncMock(return_value=stored), + update_by_name=AsyncMock(), + ) + + response = _patch_credential( + "copilot", + { + "credential_name": "copilot", + "credential_values": {"github_copilot_host": "OCTO.ghe.com"}, + "credential_info": {"custom_llm_provider": "github_copilot"}, + }, + ) + + assert response.status_code == 200, response.text + assert CredentialAccessor.find_credential("copilot").credential_values["github_copilot_host"] == "octo.ghe.com" + + +def _encrypted_per_user_copilot_credential(host: str) -> CredentialItem: + return CredentialItem( + credential_name="copilot", + credential_values={ + "github_copilot_auth_type": encrypt_value_helper("per_user_oauth"), + "github_copilot_host": encrypt_value_helper(host), + }, + credential_info={"custom_llm_provider": "github_copilot"}, + ) + + +def test_update_credential_purges_connections_when_github_host_changes(monkeypatch, restore_credential_list): + stored = _encrypted_per_user_copilot_credential("github.com") + with _repository_holding(stored): + from litellm.proxy import proxy_server + + table = MagicMock() + table.find_many = AsyncMock(return_value=[_connection_row(credential_name="copilot")]) + table.delete_many = AsyncMock(return_value=None) + proxy_server.prisma_client.db.litellm_userprovidercredentials = table + proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db + + cached_values: Final[dict[str, object]] = {} # mutable-ok: fake Redis stores invalidation writes + + async def set_cache(key: str, value: str, **kwargs: object) -> None: + cached_values[key] = value + + async def get_cache(key: str, **kwargs: object) -> object: + return cached_values.get(key) + + redis_cache: Final = MagicMock() + redis_cache.async_set_cache = AsyncMock(side_effect=set_cache) + redis_cache.async_get_cache = AsyncMock(side_effect=get_cache) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=redis_cache)) + + session_key: Final = _session_cache_key("user-a", "copilot", "github.com") + _SESSION_CACHE.set_cache( + session_key, + ("gho_secret", "copilot-token", "https://api.githubcopilot.com", 4102444800), + ttl=3600, + ) + response = _patch_credential( + "copilot", + { + "credential_name": "copilot", + "credential_values": {"github_copilot_host": "octo.ghe.com"}, + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + table.find_many.assert_has_awaits( + [ + call(where={"credential_name": "copilot"}), + call(where={"credential_name": "copilot"}), + ] + ) + table.delete_many.assert_awaited_once_with(where={"credential_name": "copilot"}) + assert cached_values[upc._cache_key("user-a", "copilot")] == upc._NOT_CONNECTED + assert _SESSION_CACHE.get_cache(session_key) is None + + +def test_update_credential_does_not_purge_connections_for_case_only_host_change(restore_credential_list): + stored = _encrypted_per_user_copilot_credential("octo.ghe.com") + with _repository_holding(stored): + from litellm.proxy import proxy_server + + table = MagicMock() + table.find_many = AsyncMock( + return_value=[_connection_row(credential_name="copilot", github_host="octo.ghe.com")] + ) + table.delete_many = AsyncMock(return_value=None) + proxy_server.prisma_client.db.litellm_userprovidercredentials = table + proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db + + session_key: Final = _session_cache_key("user-a", "copilot", "octo.ghe.com") + _SESSION_CACHE.set_cache( + session_key, + ("gho_secret", "copilot-token", "https://copilot-api.octo.ghe.com", 4102444800), + ttl=3600, + ) + response = _patch_credential( + "copilot", + { + "credential_name": "copilot", + "credential_values": {"github_copilot_host": "OCTO.ghe.com"}, + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + table.find_many.assert_not_awaited() + table.delete_many.assert_not_awaited() + assert _SESSION_CACHE.get_cache(session_key) is not None + _SESSION_CACHE.delete_cache(session_key) + + +def test_update_credential_does_not_purge_connections_for_legacy_host_case_difference(restore_credential_list): + stored = _encrypted_per_user_copilot_credential("OCTO.ghe.com") + with _repository_holding(stored): + from litellm.proxy import proxy_server + + table = MagicMock() + table.find_many = AsyncMock( + return_value=[_connection_row(credential_name="copilot", github_host="octo.ghe.com")] + ) + table.delete_many = AsyncMock(return_value=None) + proxy_server.prisma_client.db.litellm_userprovidercredentials = table + proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db + + session_key: Final = _session_cache_key("user-a", "copilot", "octo.ghe.com") + _SESSION_CACHE.set_cache( + session_key, + ("gho_secret", "copilot-token", "https://copilot-api.octo.ghe.com", 4102444800), + ttl=3600, + ) + response = _patch_credential( + "copilot", + { + "credential_name": "copilot", + "credential_values": {"github_copilot_host": "octo.ghe.com"}, + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + table.find_many.assert_not_awaited() + table.delete_many.assert_not_awaited() + assert _SESSION_CACHE.get_cache(session_key) is not None + _SESSION_CACHE.delete_cache(session_key) + + +def test_update_credential_purges_connections_when_github_host_is_deleted(restore_credential_list): + stored = _encrypted_per_user_copilot_credential("octo.ghe.com") + with _repository_holding(stored): + from litellm.proxy import proxy_server + + table = MagicMock() + table.find_many = AsyncMock(return_value=[_connection_row(credential_name="copilot")]) + table.delete_many = AsyncMock(return_value=None) + proxy_server.prisma_client.db.litellm_userprovidercredentials = table + proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db + + response = _patch_credential( + "copilot", + { + "credential_name": "copilot", + "credential_values": {}, + "credential_values_to_delete": ["github_copilot_host"], + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + table.delete_many.assert_awaited_once_with(where={"credential_name": "copilot"}) + + +def test_update_credential_purges_connections_for_invalid_legacy_github_host(restore_credential_list): + stored = _encrypted_per_user_copilot_credential("legacy.invalid") + with _repository_holding(stored): + from litellm.proxy import proxy_server + + table = MagicMock() + table.find_many = AsyncMock(return_value=[_connection_row(credential_name="copilot")]) + table.delete_many = AsyncMock(return_value=None) + proxy_server.prisma_client.db.litellm_userprovidercredentials = table + proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db + + response = _patch_credential( + "copilot", + { + "credential_name": "copilot", + "credential_values": {"api_key": "sk-updated"}, + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + table.delete_many.assert_awaited_once_with(where={"credential_name": "copilot"}) + + +def test_update_credential_does_not_purge_connections_for_unrelated_value_update(restore_credential_list): + stored = _encrypted_per_user_copilot_credential("octo.ghe.com") + with _repository_holding(stored): + from litellm.proxy import proxy_server + + table = MagicMock() + table.find_many = AsyncMock(return_value=[_connection_row(credential_name="copilot")]) + table.delete_many = AsyncMock(return_value=None) + proxy_server.prisma_client.db.litellm_userprovidercredentials = table + proxy_server.prisma_client.writer_db = proxy_server.prisma_client.db + + response = _patch_credential( + "copilot", + { + "credential_name": "copilot", + "credential_values": {"api_key": "sk-updated"}, + "credential_info": {}, + }, + ) + + assert response.status_code == 200, response.text + table.find_many.assert_not_awaited() + table.delete_many.assert_not_awaited() + + def test_update_credential_resolves_credential_values_from_model_id_like_create(credential_store): """Regression: PATCH dropped ``model_id`` from the body, so an update that named a deployment instead of raw values wrote whatever the caller sent, or nothing.""" @@ -1522,24 +1816,38 @@ def _as_user(user_id="user-a"): return lambda: UserAPIKeyAuth(api_key="test-key", user_role="internal_user", user_id=user_id) -def _per_user_credential(name="copilot-cred"): +def _per_user_credential( + name: str = "copilot-cred", + host: str | None = None, +) -> CredentialItem: return CredentialItem( credential_name=name, - credential_values={"github_copilot_auth_type": "per_user_oauth"}, + credential_values={ + "github_copilot_auth_type": "per_user_oauth", + **({"github_copilot_host": host} if host is not None else {}), + }, credential_info={}, ) -def _connection_row(user_id="user-a", credential_name="copilot-cred", login="octo"): +def _connection_row( + user_id: str = "user-a", + credential_name: str = "copilot-cred", + login: str = "octo", + github_host: str = "github.com", +) -> SimpleNamespace: from datetime import datetime, timezone - from types import SimpleNamespace return SimpleNamespace( user_id=user_id, credential_name=credential_name, provider="github_copilot", credential_b64=upc._encode( - upc.GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login=login) + upc.GithubCopilotUserConnectionPayload( + access_token="gho_secret", + github_login=login, + github_host=github_host, + ) ), created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), updated_at=datetime(2026, 1, 1, tzinfo=timezone.utc), @@ -1635,6 +1943,20 @@ def test_user_connection_start_returns_device_flow_fields(monkeypatch): } handle = response.json()["flow_handle"] assert isinstance(handle, str) and "device_code" not in handle and _DEVICE_FLOW_START["device_code"] not in handle + decoded_handle: Final = _decode_flow_handle(handle) + assert decoded_handle is not None + assert decoded_handle.github_host == "github.com" + + +def test_legacy_user_connection_flow_handle_defaults_to_github_com(): + handle: Final = UserConnectionFlowHandle( + user_id="user-a", + credential_name="copilot-cred", + device_code="DC-1", + interval=5, + expires_at=4102444800, + ) + assert handle.github_host == "github.com" def test_user_connection_start_404s_for_shared_credential(monkeypatch): @@ -1657,7 +1979,7 @@ def test_user_connection_poll_rejects_an_undecryptable_flow_handle(monkeypatch): def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch): - table = _patch_user_connection_env(monkeypatch, [_per_user_credential()]) + table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential()]) with _github_http( { _DEVICE_CODE_URL: _DEVICE_FLOW_START, @@ -1666,9 +1988,13 @@ def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch): _GITHUB_USER_URL: {"login": "octo"}, } ): - start = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user()) + start: Final = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user()) assert start.status_code == 200, start.text - poll = _call_as( + encrypted_handle: Final = start.json()["flow_handle"] + decoded_handle: Final = _decode_flow_handle(encrypted_handle) + assert decoded_handle is not None + assert decoded_handle.github_host == "github.com" + poll: Final = _call_as( "POST", "/credentials/copilot-cred/user_connection/poll", json_body={"flow_handle": start.json()["flow_handle"]}, @@ -1677,6 +2003,110 @@ def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch): assert poll.status_code == 200, poll.text assert poll.json() == {"status": "connected", "interval": None, "github_login": "octo"} table.upsert.assert_awaited_once() + stored_connection: Final = upc.decode_user_provider_credential( + table.upsert.await_args.kwargs["data"]["create"]["credential_b64"] + ) + assert stored_connection is not None + assert stored_connection.github_host == "github.com" + + +def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatch): + from litellm.llms.github_copilot.per_user_auth import attach_github_copilot_user_session + + table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")]) + + def change_host_after_token_exchange(request: httpx.Request) -> dict[str, str]: + monkeypatch.setattr(litellm, "credential_list", [_per_user_credential()]) + return {"access_token": "gho_1"} + + with _github_http( + { + "https://octo.ghe.com/login/device/code": _DEVICE_FLOW_START, + "https://octo.ghe.com/login/oauth/access_token": change_host_after_token_exchange, + "https://api.octo.ghe.com/copilot_internal/v2/token": { + "token": "copilot-tok", + "expires_at": 4102444800, + "endpoints": {"api": "https://copilot-api.octo.ghe.com"}, + }, + "https://api.octo.ghe.com/user": {"login": "octo"}, + } + ): + start: Final = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user()) + poll: Final = _call_as( + "POST", + "/credentials/copilot-cred/user_connection/poll", + json_body={"flow_handle": start.json()["flow_handle"]}, + auth=_as_user(), + ) + assert poll.status_code == 200, poll.text + stored_connection: Final = upc.decode_user_provider_credential( + table.upsert.await_args.kwargs["data"]["create"]["credential_b64"] + ) + assert stored_connection is not None + + kwargs: Final[dict[str, object]] = { + "litellm_credential_name": "copilot-cred", + "secret_fields": { + "user_provider_credentials_user_id": "user-a", + "user_provider_credentials": { + "copilot-cred": { + "access_token": stored_connection.access_token, + "github_host": stored_connection.github_host, + } + }, + }, + } + module_level_client: Final = MagicMock() + module_level_client.get.side_effect = AssertionError("unexpected token exchange") + with _github_http({}), patch.object(litellm, "module_level_client", module_level_client): + with pytest.raises(litellm.CallerCredentialAuthenticationError) as exc: + attach_github_copilot_user_session(kwargs) + assert "Connect GitHub Copilot for credential 'copilot-cred'" in str(exc.value) + assert stored_connection.github_host == "octo.ghe.com" + + +def test_ghe_user_connection_uses_tenant_urls_for_device_flow_exchange_and_login(monkeypatch): + table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")]) + with _github_http( + { + "https://octo.ghe.com/login/device/code": _DEVICE_FLOW_START, + "https://octo.ghe.com/login/oauth/access_token": {"access_token": "gho_1"}, + "https://api.octo.ghe.com/copilot_internal/v2/token": { + "token": "copilot-tok", + "expires_at": 4102444800, + "endpoints": {"api": "https://copilot-api.octo.ghe.com"}, + }, + "https://api.octo.ghe.com/user": {"login": "octo"}, + } + ): + start: Final = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user()) + assert start.status_code == 200, start.text + poll: Final = _call_as( + "POST", + "/credentials/copilot-cred/user_connection/poll", + json_body={"flow_handle": start.json()["flow_handle"]}, + auth=_as_user(), + ) + assert poll.status_code == 200, poll.text + assert poll.json() == {"status": "connected", "interval": None, "github_login": "octo"} + table.upsert.assert_awaited_once() + stored_connection: Final = upc.decode_user_provider_credential( + table.upsert.await_args.kwargs["data"]["create"]["credential_b64"] + ) + assert stored_connection is not None + assert stored_connection.github_host == "octo.ghe.com" + + +def test_user_connection_poll_rejects_handle_from_different_host(monkeypatch): + _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")]) + response = _call_as( + "POST", + "/credentials/copilot-cred/user_connection/poll", + json_body={"flow_handle": _flow_handle_for()}, + auth=_as_user(), + ) + assert response.status_code == 400 + assert response.json()["error"]["message"] == "invalid or expired flow_handle" def test_user_connection_poll_pending(monkeypatch): @@ -2038,6 +2468,23 @@ def test_pending_device_flow_survives_disconnect_via_stateless_handle(monkeypatc assert delete.json() == {"status": "disconnected"} +def test_delete_user_connection_evicts_copilot_session(monkeypatch): + table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential()], rows=[_connection_row()]) + session_key: Final = _session_cache_key("user-a", "copilot-cred", "github.com") + _SESSION_CACHE.set_cache( + session_key, + ("gho_secret", "copilot-token", "https://api.githubcopilot.com", 4102444800), + ttl=3600, + ) + + response: Final = _call_as("DELETE", "/credentials/copilot-cred/user_connection", auth=_as_user()) + + assert response.status_code == 200, response.text + assert response.json() == {"status": "disconnected"} + assert _SESSION_CACHE.get_cache(session_key) is None + table.delete_many.assert_awaited_once_with(where={"user_id": "user-a", "credential_name": "copilot-cred"}) + + def test_user_connection_cache_keys_cannot_collide_on_colons(): """user_id/credential_name are joined losslessly, so ('a:b','c') and ('a','b:c') can never share a cache or device-flow entry.""" diff --git a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py index b549ecc67a9..dae3a0459e2 100644 --- a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py +++ b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py @@ -1,5 +1,6 @@ """Tests for the per-user provider credentials DB/cache module.""" +import json from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -7,9 +8,11 @@ from unittest.mock import AsyncMock, MagicMock import pytest from litellm.caching.dual_cache import DualCache +from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.credential_endpoints.user_provider_credentials import ( GithubCopilotUserConnectionPayload, - aget_user_provider_tokens, + aget_user_provider_credentials, decode_user_provider_credential, delete_user_provider_credential, delete_user_provider_credentials_for_credential, @@ -110,7 +113,7 @@ async def test_delete_for_credential_returns_affected_user_ids(): @pytest.mark.asyncio -async def test_aget_tokens_returns_plaintext_and_caches_ciphertext(monkeypatch): +async def test_aget_credentials_returns_plaintext_and_caches_ciphertext(monkeypatch): payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") from litellm.proxy.credential_endpoints import user_provider_credentials as upc @@ -121,12 +124,12 @@ async def test_aget_tokens_returns_plaintext_and_caches_ciphertext(monkeypatch): redis = _fake_redis(monkeypatch) cache = DualCache(redis_cache=redis) - tokens = await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) - assert tokens == {"copilot-cred": "gho_secret"} + connections = await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) + assert connections == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}} # second call served from cache: no new DB read, and cache holds ciphertext not plaintext - tokens2 = await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) - assert tokens2 == {"copilot-cred": "gho_secret"} + connections2 = await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) + assert connections2 == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}} assert table.find_many.await_count == 1 cached = await redis.async_get_cache(upc._cache_key("user-a", "copilot-cred")) assert cached == ciphertext @@ -135,14 +138,59 @@ async def test_aget_tokens_returns_plaintext_and_caches_ciphertext(monkeypatch): @pytest.mark.asyncio -async def test_aget_tokens_negative_caches_missing_connection(monkeypatch): +async def test_aget_connections_preserves_github_host(monkeypatch): + payload: Final = GithubCopilotUserConnectionPayload( + access_token="gho_secret", + github_login="octo", + github_host="octo.ghe.com", + ) + from litellm.proxy.credential_endpoints import user_provider_credentials as upc + + table: Final = MagicMock() + table.find_many = AsyncMock(return_value=[_row(credential_b64=upc._encode(payload))]) + prisma_client: Final = _prisma(table) + cache: Final = DualCache(redis_cache=_fake_redis(monkeypatch)) + + connections: Final = await aget_user_provider_credentials( + prisma_client, + cache, + "user-a", + ["copilot-cred"], + ) + assert connections == {"copilot-cred": {"access_token": "gho_secret", "github_host": "octo.ghe.com"}} + + +@pytest.mark.asyncio +async def test_legacy_connection_without_github_host_defaults_to_github_com(monkeypatch): + ciphertext: Final = encrypt_value_helper(json.dumps({"access_token": "gho_secret", "github_login": "octo"})) + table: Final = MagicMock() + table.find_many = AsyncMock(return_value=[_row(credential_b64=ciphertext)]) + prisma_client: Final = _prisma(table) + cache: Final = DualCache(redis_cache=_fake_redis(monkeypatch)) + + connections: Final = await aget_user_provider_credentials( + prisma_client, + cache, + "user-a", + ["copilot-cred"], + ) + assert connections == { + "copilot-cred": { + "access_token": "gho_secret", + "github_host": GITHUB_COPILOT_DEFAULT_HOST, + } + } + + +@pytest.mark.asyncio +async def test_aget_credentials_negative_caches_missing_connection(monkeypatch): table = MagicMock() table.find_many = AsyncMock(return_value=[]) prisma_client = _prisma(table) cache = DualCache(redis_cache=_fake_redis(monkeypatch)) - assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {} - assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {} + assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {} + assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {} assert table.find_many.await_count == 1 @@ -156,10 +204,10 @@ async def test_invalidate_cache_forces_db_refetch(monkeypatch): prisma_client = _prisma(table) cache = DualCache(redis_cache=_fake_redis(monkeypatch)) - await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) + await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) await drop_user_provider_credential_cache(cache, "user-a", "copilot-cred") - tokens = await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) - assert tokens == {"copilot-cred": "gho_new"} + connections = await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) + assert connections == {"copilot-cred": {"access_token": "gho_new", "github_host": "github.com"}} assert table.find_many.await_count == 2 @@ -184,19 +232,19 @@ async def test_disconnect_on_one_worker_invalidates_other_workers_via_redis(monk prisma_client = _prisma(table) # worker A "connects" (DB read + cache write) - tokens_a = await aget_user_provider_tokens(prisma_client, cache_a, "user-a", ["copilot-cred"]) - assert tokens_a == {"copilot-cred": "gho_secret"} + connections_a = await aget_user_provider_credentials(prisma_client, cache_a, "user-a", ["copilot-cred"]) + assert connections_a == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}} # worker B reads the same connection straight from Redis (its in-memory layer is empty) - tokens_b = await aget_user_provider_tokens(prisma_client, cache_b, "user-a", ["copilot-cred"]) - assert tokens_b == {"copilot-cred": "gho_secret"} + connections_b = await aget_user_provider_credentials(prisma_client, cache_b, "user-a", ["copilot-cred"]) + assert connections_b == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}} # worker A disconnects: the key is overwritten with the not-connected tombstone await invalidate_user_provider_credential_cache(cache_a, "user-a", "copilot-cred") # worker B sees the tombstone from Redis, no worker-local staleness, no DB read needed - tokens_b_after = await aget_user_provider_tokens(prisma_client, cache_b, "user-a", ["copilot-cred"]) - assert tokens_b_after == {} + connections_b_after = await aget_user_provider_credentials(prisma_client, cache_b, "user-a", ["copilot-cred"]) + assert connections_b_after == {} assert table.find_many.await_count == 1 @@ -222,8 +270,8 @@ async def test_stale_fill_cannot_overwrite_a_disconnect_tombstone(monkeypatch): table.find_many = AsyncMock(side_effect=_read_while_disconnecting) prisma_client = _prisma(table) - await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) - assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == {} + await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) + assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {} @pytest.mark.asyncio @@ -240,8 +288,8 @@ async def test_redis_read_failure_falls_back_to_the_database(monkeypatch): table.find_many = AsyncMock(return_value=[_row(credential_b64=upc._encode(payload))]) prisma_client = _prisma(table) - assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == { - "copilot-cred": "gho_secret" + assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == { + "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} } @@ -259,8 +307,8 @@ async def test_redis_write_and_delete_failures_do_not_fail_the_request(monkeypat table.find_many = AsyncMock(return_value=[_row(credential_b64=upc._encode(payload))]) prisma_client = _prisma(table) - assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == { - "copilot-cred": "gho_secret" + assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == { + "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} } await invalidate_user_provider_credential_cache(cache, "user-a", "copilot-cred") @@ -279,14 +327,14 @@ async def test_without_redis_no_worker_local_entries_are_kept(): cache_a = DualCache() cache_b = DualCache() - assert await aget_user_provider_tokens(prisma_client, cache_a, "user-a", ["copilot-cred"]) == { - "copilot-cred": "gho_secret" + assert await aget_user_provider_credentials(prisma_client, cache_a, "user-a", ["copilot-cred"]) == { + "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} } - assert await aget_user_provider_tokens(prisma_client, cache_b, "user-a", ["copilot-cred"]) == { - "copilot-cred": "gho_secret" + assert await aget_user_provider_credentials(prisma_client, cache_b, "user-a", ["copilot-cred"]) == { + "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} } await invalidate_user_provider_credential_cache(cache_b, "user-a", "copilot-cred") - assert await aget_user_provider_tokens(prisma_client, cache_a, "user-a", ["copilot-cred"]) == {} + assert await aget_user_provider_credentials(prisma_client, cache_a, "user-a", ["copilot-cred"]) == {} assert table.find_many.await_count == 3 @@ -306,8 +354,8 @@ async def test_connect_overwrites_a_not_connected_tombstone(monkeypatch): table = MagicMock() table.find_many = AsyncMock(return_value=[]) prisma_client = _prisma(table) - assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == { - "copilot-cred": "gho_secret" + assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == { + "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} } table.find_many.assert_not_awaited() @@ -332,9 +380,9 @@ async def test_stale_not_connected_fill_cannot_overwrite_a_connect_write(monkeyp table.find_many = AsyncMock(side_effect=_read_while_connecting) prisma_client = _prisma(table) - await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) - assert await aget_user_provider_tokens(prisma_client, cache, "user-a", ["copilot-cred"]) == { - "copilot-cred": "gho_secret" + await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) + assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == { + "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} } @@ -391,7 +439,7 @@ async def test_reads_hit_the_writer_engine_not_a_stale_replica(): prisma_client.db.litellm_userprovidercredentials = reader_table prisma_client.writer_db.litellm_userprovidercredentials = writer_table - await aget_user_provider_tokens(prisma_client, DualCache(), "user-a", ["copilot-cred"]) + await aget_user_provider_credentials(prisma_client, DualCache(), "user-a", ["copilot-cred"]) writer_table.find_many.assert_awaited_once() reader_table.find_many.assert_not_awaited() diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 5ecc8784384..f30606ac2f1 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -6,10 +6,78 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy._types import DailyTagSpendTransaction +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.proxy_server import ProxyStartupEvent +def _daily_tag_spend_update_transactions() -> dict[str, DailyTagSpendTransaction]: + return { + "research_2025-01-15_api-key-123_gpt-4o_openai_/chat/completions": DailyTagSpendTransaction( + tag="research", + request_id="request-123", + date="2025-01-15", + api_key="api-key-123", + model="gpt-4o", + model_group=None, + mcp_namespaced_tool_name=None, + custom_llm_provider="openai", + endpoint="/chat/completions", + prompt_tokens=12, + completion_tokens=7, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + compression_saved_tokens=0, + compression_savings_spend=0.0, + prompt_caching_savings_spend=0.0, + gateway_injected_caching_savings_spend=0.0, + spend=0.0042, + api_requests=1, + successful_requests=1, + failed_requests=0, + ) + } + + +@pytest.mark.asyncio +async def test_daily_tag_spend_restored_to_in_memory_queue_on_rpush_failure( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock +): + mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away")) + queue: Final = DailySpendUpdateQueue() + original_transactions: Final = _daily_tag_spend_update_transactions() + await queue.add_update(original_transactions) + + await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) + + assert await queue.flush_and_get_aggregated_daily_spend_update_transactions() == original_transactions + + +@pytest.mark.asyncio +async def test_daily_tag_spend_restored_after_rpush_failure_is_pushed_once_on_next_tick( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock +): + mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away")) + queue: Final = DailySpendUpdateQueue() + original_transactions: Final = _daily_tag_spend_update_transactions() + await queue.add_update(original_transactions) + + await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) + mock_redis_cache.async_rpush = AsyncMock(return_value=1) + + await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) + + mock_redis_cache.async_rpush.assert_called_once() + redis_key: Final = mock_redis_cache.async_rpush.call_args.kwargs["key"] + values: Final = mock_redis_cache.async_rpush.call_args.kwargs["values"] + assert redis_key == REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY + assert json.loads(values[0]) == json.loads(safe_dumps(original_transactions)) + assert await queue.flush_and_get_aggregated_daily_spend_update_transactions() == {} + + @pytest.fixture def mock_redis_cache(): """Create a mock RedisCache instance""" diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index d440b99a632..7437042e4bd 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -19,7 +19,7 @@ from redis.exceptions import DataError import litellm from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes -from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS +from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_UPDATE_BUFFER_KEY, SPEND_ROLLUP_LOCK_TIMEOUT_MS from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import ( _TEAM_ADVISORY_LOCK_SQL, @@ -29,12 +29,14 @@ from litellm.proxy.db.db_spend_update_writer import ( _SpendTableName, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, ) from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -2365,6 +2367,58 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() +class _InMemoryRedis: + def __init__(self) -> None: + self.store: dict[str, str] = {} + self.lists: dict[str, list[object]] = {} + + async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None: + if nx and key in self.store: + return None + self.store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs: object) -> str | None: + raw = self.store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.store.pop(key, None) is not None else 0 + + async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]: + for op in rpush_list: + self.lists.setdefault(op["key"], []).extend(op["values"]) + return [len(self.lists[op["key"]]) for op in rpush_list] + + async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]: + return [self.lists.pop(op["key"], None) for op in lpop_list] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lock_held_by_another_pod", [False, True]) +async def test_commit_with_redis_writes_discovered_tools_on_leader_and_non_leader(lock_held_by_another_pod: bool): + redis = _InMemoryRedis() + if lock_held_by_another_pod: + assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=DB_SPEND_UPDATE_JOB_NAME) + db_writer = DBSpendUpdateWriter(redis_cache=redis) + db_writer.pod_lock_manager.redis_cache = redis + db_writer.tool_discovery_queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.litellm_tooltable.upsert = AsyncMock() + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + upserted = [c.kwargs["where"] for c in prisma_client.db.litellm_tooltable.upsert.await_args_list] + assert upserted == [{"tool_name": "get_weather"}] + assert (REDIS_UPDATE_BUFFER_KEY in redis.lists) is lock_held_by_another_pod, ( + f"spend buffer left in redis={list(redis.lists)}, only the lock holder may drain it" + ) + + @pytest.mark.asyncio async def test_commit_with_redis_requeues_all_on_db_failure(): """ diff --git a/tests/unit/proxy/db/test_health_check_latest.py b/tests/unit/proxy/db/test_health_check_latest.py index 29063d7b060..2f27be4fe61 100644 --- a/tests/unit/proxy/db/test_health_check_latest.py +++ b/tests/unit/proxy/db/test_health_check_latest.py @@ -7,8 +7,10 @@ import pytest from litellm.proxy.db.health_check_latest import ( LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, LATEST_HEALTH_CHECKS_SQL, + LatestHealthCheckRow, fetch_latest_health_checks, fetch_latest_health_checks_for_models, + latest_by_model_name, query_latest_health_checks, ) from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -129,3 +131,36 @@ async def test_the_latest_health_check_read_renders_a_postgres_select_span( ) -> None: assert await fetch_latest_health_checks(_prisma([])) == () assert await postgres_span_names() == ("postgres.select LiteLLM_HealthCheckTable",) + + +def _latest_row(**overrides) -> LatestHealthCheckRow: + return LatestHealthCheckRow.model_validate(_raw_row(**overrides)) + + +@pytest.mark.parametrize("reverse", [False, True], ids=["sql-order", "reversed"]) +def test_latest_by_model_name_keeps_each_groups_newest_row_whatever_its_model_id(reverse): + """model_id ASC puts a name probe's NULL row, or a deleted deployment's larger id, after the fresh rows.""" + fresh = _latest_row(health_check_id="hc-a", model_id="deployment-a", checked_at="2026-10-16T11:10:00+00:00") + deleted = _latest_row( + health_check_id="hc-z", model_id="deployment-z", status="unhealthy", checked_at="2026-09-04T11:00:00+00:00" + ) + name_probe = _latest_row( + health_check_id="hc-null", model_id=None, status="unhealthy", checked_at="2026-09-04T11:20:04+00:00" + ) + other_group = _latest_row(health_check_id="hc-c", model_name="claude", model_id="deployment-c") + rows = (other_group, fresh, deleted, name_probe) + + latest = latest_by_model_name(rows[::-1] if reverse else rows) + + assert latest == {"gpt-4": fresh, "claude": other_group} + + +def test_latest_by_model_name_lets_a_newer_name_probe_speak_for_the_group(): + per_deployment = _latest_row( + health_check_id="hc-a", model_id="deployment-a", checked_at="2026-10-16T11:10:00+00:00" + ) + name_probe = _latest_row( + health_check_id="hc-null", model_id=None, status="unhealthy", checked_at="2026-10-16T12:00:00+00:00" + ) + + assert latest_by_model_name((per_deployment, name_probe))["gpt-4"] is name_probe diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index fcf56eaea9a..021eb289ad7 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -8,7 +8,7 @@ import sys from collections.abc import AsyncGenerator, Callable, Iterator, Mapping from contextlib import asynccontextmanager from pathlib import Path -from typing import Final +from typing import Final, Literal import pytest import respx @@ -19,6 +19,8 @@ from pydantic import TypeAdapter from starlette.routing import Match import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.decisions_pricing import apply_decisions_pricing from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeature, attach_lazy_features from litellm.proxy._types import UserAPIKeyAuth @@ -32,6 +34,9 @@ from litellm.proxy.proxy_server import ( initialize, ) from litellm.types.decisions import DecisionsRequestBody, OpenAIDecisionRequestBody +from litellm.types.extraction import ExtractionRequestBody +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs from tests._master_key import MASTER_KEY if sys.version_info >= (3, 14): @@ -1296,6 +1301,328 @@ def test_the_two_decisions_routes_document_different_request_formats(decisions_o assert "input" in properties("/v1/decisions") and "input" not in properties("/v1/systemone") +_CLAIRIC_UPSTREAM: Final = "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/@cf/cloudflare/clairic" +_CLAIRIC_INPUT_TOKENS: Final = 416 +_CLAIRIC_RESULT: Final[Mapping[str, object]] = { + "model": "clairic", + "answers": { + "vendor": {"type": "text", "p_null": 0.0177, "span": [19, 29], "confidence": 0.9998, "value": "Acme Tools"}, + "days_to_pay": {"type": "estimate", "value": 31, "mean": 30.76, "median": 30, "p10": 29, "p90": 34}, + }, + "truncated": False, + "usage": {"input_tokens": _CLAIRIC_INPUT_TOKENS, "output_tokens": 0, "schema_tokens": 312, "passes": 2}, +} +_EXTRACTION_REQUEST: Final[Mapping[str, object]] = { + "model": "invoice-reader", + "state": "Invoice #4471 from Acme Tools. Payment due in 30 days.", + "fields": { + "vendor": {"type": "text", "instructions": "Who sent the invoice?"}, + "days_to_pay": {"type": "estimate", "min": 0, "max": 90}, + }, + "images": ["data:image/png;base64,AA=="], +} + + +def _route_invoice_reader_to_clairic( + monkeypatch: pytest.MonkeyPatch, drop_params: bool = False, additional_drop_params: tuple[str, ...] = () +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "invoice-reader", + "litellm_params": { + "model": "cloudflare/clairic", + "api_key": "cloudflare-key", + "api_base": "https://api.cloudflare.com/client/v4/accounts/acct/ai/run", + "drop_params": drop_params, + **({"additional_drop_params": list(additional_drop_params)} if additional_drop_params else {}), + }, + } + ] + ), + ) + + +@pytest.mark.parametrize("endpoint", ("/v1/extraction", "/extraction")) +def test_proxy_extraction_forwards_the_body_to_clairic_and_returns_its_answers_and_cost( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + endpoint: str, +) -> None: + _route_invoice_reader_to_clairic(monkeypatch) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond( + json={"result": _CLAIRIC_RESULT, "success": True, "errors": [], "messages": []} + ) + + response: Final = client.post(endpoint, json=_EXTRACTION_REQUEST) + + assert response.status_code == 200, response.text + assert response.json() == {**_CLAIRIC_RESULT, "model": "invoice-reader"} + expected_cost: Final = _CLAIRIC_INPUT_TOKENS * float( + litellm.model_cost["cloudflare/clairic"]["input_cost_per_token"] + ) + assert expected_cost > 0 + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected_cost) + assert json.loads(upstream.calls[0].request.content) == {**_EXTRACTION_REQUEST, "model": "clairic"} + assert upstream.calls[0].request.headers["authorization"] == "Bearer cloudflare-key" + + +@pytest.mark.parametrize( + "request_body", + ( + {"model": "invoice-reader", "state": "Invoice #4471"}, + {"model": "invoice-reader", "state": "Invoice #4471", "fields": {"total": {"type": "money"}}}, + {**_EXTRACTION_REQUEST, "images": ["https://example.com/invoice.png"]}, + ), + ids=("missing_fields", "unknown_field_type", "remote_image"), +) +def test_proxy_extraction_refuses_a_malformed_body_before_calling_clairic( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + request_body: Mapping[str, object], +) -> None: + _route_invoice_reader_to_clairic(monkeypatch) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond(json=_CLAIRIC_RESULT) + + response: Final = client.post("/v1/extraction", json=request_body) + + assert response.status_code == 400, response.text + assert "Invalid Extraction request" in response.text + assert not upstream.called + + +@pytest.mark.parametrize( + ("key", "value"), + ( + ("language", "en"), + ("priority", 1), + ("extra_body", {"language": "en"}), + ("temperature", 0.5), + ("user", "end-user-1"), + ), +) +def test_proxy_extraction_refuses_a_key_clairic_does_not_take_before_calling_clairic( + client: TestClient, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, key: str, value: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + _route_invoice_reader_to_clairic(monkeypatch) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond(json={"result": _CLAIRIC_RESULT, "success": True}) + + response: Final = client.post("/v1/extraction", json={**_EXTRACTION_REQUEST, key: value}) + + assert response.status_code == 400, response.text + assert f"does not support parameters: ['{key}']" in response.text + assert not upstream.called + + +@pytest.mark.parametrize( + ("drop_params", "additional_drop_params"), + ((True, ()), (False, ("language", "temperature", "user"))), + ids=("drop_params", "additional_drop_params"), +) +def test_proxy_extraction_sends_only_the_clairic_request_when_the_deployment_drops_other_keys( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + drop_params: bool, + additional_drop_params: tuple[str, ...], +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + _route_invoice_reader_to_clairic(monkeypatch, drop_params, additional_drop_params) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond(json={"result": _CLAIRIC_RESULT, "success": True}) + + response: Final = client.post( + "/v1/extraction", json={**_EXTRACTION_REQUEST, "language": "en", "temperature": 0.2, "user": "end-user-1"} + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == {**_EXTRACTION_REQUEST, "model": "clairic"} + + +def test_proxy_extraction_leaves_out_the_cli_generation_defaults( + client: TestClient, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_max_tokens", 256) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_temperature", 0.2) + _route_invoice_reader_to_clairic(monkeypatch) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond(json={"result": _CLAIRIC_RESULT, "success": True}) + + response: Final = client.post("/v1/extraction", json=_EXTRACTION_REQUEST) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == {**_EXTRACTION_REQUEST, "model": "clairic"} + + +class _PreCallRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: object, data: dict, call_type: str + ) -> dict: + self.calls += 1 + return data + + +def test_proxy_extraction_refuses_a_remote_image_before_any_pre_call_hook_runs( + client: TestClient, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + recorder: Final = _PreCallRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + _route_invoice_reader_to_clairic(monkeypatch) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond(json={"result": _CLAIRIC_RESULT, "success": True}) + + response: Final = client.post( + "/v1/extraction", json={**_EXTRACTION_REQUEST, "images": ["https://example.com/invoice.png"]} + ) + + assert response.status_code == 400, response.text + assert recorder.calls == 0 + assert not upstream.called + + +def test_proxy_extraction_refuses_stream_before_calling_clairic_unless_the_deployment_drops_params( + client: TestClient, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + _route_invoice_reader_to_clairic(monkeypatch) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond(json={"result": _CLAIRIC_RESULT, "success": True}) + + response: Final = client.post("/v1/extraction", json={**_EXTRACTION_REQUEST, "stream": True}) + + assert response.status_code == 400, response.text + assert "does not support parameters: ['stream']" in response.text + assert not upstream.called + + +def test_proxy_extraction_answers_a_dropped_stream_with_one_json_body( + client: TestClient, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + _route_invoice_reader_to_clairic(monkeypatch, drop_params=True) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond(json={"result": _CLAIRIC_RESULT, "success": True}) + + response: Final = client.post("/v1/extraction", json={**_EXTRACTION_REQUEST, "stream": True}) + + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("application/json") + assert response.json() == {**_CLAIRIC_RESULT, "model": "invoice-reader"} + assert json.loads(upstream.calls[0].request.content) == {**_EXTRACTION_REQUEST, "model": "clairic"} + + +class _RequestEmailMasker(CustomGuardrail): + def __init__(self, email: str) -> None: + super().__init__(guardrail_name="mask-email", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.email = email + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs( + texts=[text.replace(self.email, "[EMAIL]") for text in inputs.get("texts", [])] + ) + + +_EMAIL_STATE: Final = "Questions go to billing@acme.example today. The vendor is Acme Tools." +_MASKED_EMAIL_STATE: Final = "Questions go to [EMAIL] today. The vendor is Acme Tools." +_VENDOR_IN_MASKED_STATE: Final = { + "type": "text", + "value": "Acme Tools", + "span": [_MASKED_EMAIL_STATE.index("Acme Tools"), _MASKED_EMAIL_STATE.index("Acme Tools") + len("Acme Tools")], +} + + +def _clairic_vendor_result(state_text: str | None = None) -> Mapping[str, object]: + return { + "model": "clairic", + "answers": {"vendor": _VENDOR_IN_MASKED_STATE}, + "truncated": False, + "usage": {"input_tokens": 52, "output_tokens": 0}, + **({"state_text": state_text} if state_text is not None else {}), + } + + +@pytest.mark.parametrize( + ("guardrail_email", "sent_state", "returned_vendor"), + ( + ("billing@acme.example", _MASKED_EMAIL_STATE, {"type": "text", "value": "Acme Tools"}), + ("nobody@example.com", _EMAIL_STATE, _VENDOR_IN_MASKED_STATE), + ), + ids=("state_masked", "state_untouched"), +) +def test_proxy_extraction_drops_spans_into_a_state_only_its_request_guardrail_saw( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + guardrail_email: str, + sent_state: str, + returned_vendor: Mapping[str, object], +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_RequestEmailMasker(guardrail_email)]) + _route_invoice_reader_to_clairic(monkeypatch) + upstream: Final = respx_mock.post(_CLAIRIC_UPSTREAM).respond( + json={"result": _clairic_vendor_result(), "success": True} + ) + + response: Final = client.post( + "/v1/extraction", + json={"model": "invoice-reader", "state": _EMAIL_STATE, "fields": {"vendor": {"type": "text"}}}, + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content)["state"] == sent_state + assert response.json()["answers"] == {"vendor": returned_vendor} + + +def test_proxy_extraction_keeps_spans_into_the_masked_state_it_returns( + client: TestClient, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_RequestEmailMasker("billing@acme.example")]) + _route_invoice_reader_to_clairic(monkeypatch) + respx_mock.post(_CLAIRIC_UPSTREAM).respond( + json={"result": _clairic_vendor_result(state_text=_MASKED_EMAIL_STATE), "success": True} + ) + + response: Final = client.post( + "/v1/extraction", + json={ + "model": "invoice-reader", + "state": _EMAIL_STATE, + "fields": {"vendor": {"type": "text"}}, + "options": {"return_state_text": True}, + }, + ) + + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["answers"] == {"vendor": _VENDOR_IN_MASKED_STATE} + assert body["state_text"][slice(*body["answers"]["vendor"]["span"])] == "Acme Tools" + + +@pytest.mark.parametrize("path", ("/v1/extraction", "/extraction")) +def test_swagger_documents_the_extraction_request_and_response(decisions_openapi: dict[str, object], path: str) -> None: + operation = decisions_openapi["paths"][path]["post"] + media = operation["requestBody"]["content"]["application/json"] + + assert {"model", "fields"} <= set(media["schema"]["required"]) + assert "$ref" not in json.dumps(media["schema"]) + example = {key: value for key, value in media["example"].items() if key != "model"} + assert len(TypeAdapter(ExtractionRequestBody).validate_python(example).fields) == 5 + response_ref = operation["responses"]["200"]["content"]["application/json"]["schema"]["$ref"] + assert response_ref == "#/components/schemas/ExtractionResponse" + assert {"answers", "usage"} <= set(decisions_openapi["components"]["schemas"]["ExtractionResponse"]["properties"]) + + def test_proxy_systemone_reads_a_zstd_compressed_request_body( client: TestClient, respx_mock: respx.MockRouter, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 63fcdfbb0c5..df2fc8e1da2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1,6 +1,7 @@ import asyncio import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock import pytest @@ -1975,6 +1976,20 @@ class TestProxyWiring: class TestStreamUsage: + @pytest.mark.asyncio + async def test_trailing_flush_survives_a_usage_chunk_with_no_choices(self): + guardrail, _ = _shielded({"[EMAIL_1]": "a@example.com"}) + held: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Mail [EMAI"))]) + usage_chunk: Final = ModelResponseStream( + choices=[], usage=Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12) + ) + + out: Final = await _restore_stream(guardrail, [held, usage_chunk]) + + assert out[1] is not None and out[1].choices == [] + assert [chunk.choices[0].delta.content for chunk in out[2:]] == ["[EMAI"] + assert all(getattr(chunk, "usage", None) is None for chunk in out[2:]) + @pytest.mark.asyncio async def test_trailing_flush_does_not_repeat_usage(self): """With n>=2 and include_usage, the last chunk carries usage; a flush copied from it must not. diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index d5288a7fc12..950d1742d49 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -5297,9 +5297,49 @@ def test_test_model_connection_accepts_image_edit_mode(monkeypatch): assert response.json()["status"] == "success" -@pytest.mark.parametrize("mode", ("decisions", "evaluation"), ids=("decisions", "legacy-evaluation")) -def test_test_model_connection_accepts_decision_modes( - monkeypatch, mode: Literal["evaluation", "decisions"] +_TYPESAFE_TEST_CONNECTION: Final = ( + {"model": "typesafe/jev-latest", "api_key": "sk-test"}, + "https://api.typesafe.ai/v1/systemone", + { + "model": "jev-latest", + "answers": {"reachable": {"type": "noul", "noul": 1.0}}, + "usage": {"input_tokens": 12, "output_tokens": 1}, + }, +) +_CLAIRIC_TEST_CONNECTION: Final = ( + { + "model": "cloudflare/clairic", + "api_key": "sk-test", + "api_base": "https://api.cloudflare.com/client/v4/accounts/acct/ai/run", + }, + "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/@cf/cloudflare/clairic", + { + "result": { + "model": "clairic", + "answers": {"status": {"type": "text", "value": "ok"}}, + "truncated": False, + "usage": {"input_tokens": 9, "output_tokens": 0}, + }, + "success": True, + }, +) + + +@pytest.mark.parametrize( + ("mode", "litellm_params", "upstream_url", "upstream_json"), + ( + ("decisions", *_TYPESAFE_TEST_CONNECTION), + ("evaluation", *_TYPESAFE_TEST_CONNECTION), + ("extraction", *_CLAIRIC_TEST_CONNECTION), + ), + ids=("decisions", "legacy-evaluation", "extraction"), +) +def test_test_model_connection_accepts_decision_and_extraction_modes( + monkeypatch, + mode: Literal["evaluation", "decisions", "extraction"], + litellm_params: Mapping[str, str], + upstream_url: str, + upstream_json: Mapping[str, object], ): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() @@ -5315,20 +5355,8 @@ def test_test_model_connection_accepts_decision_modes( ), respx.mock(assert_all_called=True) as respx_mock, ): - upstream = respx_mock.post("https://api.typesafe.ai/v1/systemone").respond( - json={ - "model": "jev-latest", - "answers": {"reachable": {"type": "noul", "noul": 1.0}}, - "usage": {"input_tokens": 12, "output_tokens": 1}, - } - ) - response = client.post( - "/health/test_connection", - json={ - "mode": mode, - "litellm_params": {"model": "typesafe/jev-latest", "api_key": "sk-test"}, - }, - ) + upstream = respx_mock.post(upstream_url).respond(json=upstream_json) + response = client.post("/health/test_connection", json={"mode": mode, "litellm_params": litellm_params}) assert response.status_code == 200, response.text assert response.json()["status"] == "success" diff --git a/tests/unit/proxy/hooks/test_key_management_event_hooks.py b/tests/unit/proxy/hooks/test_key_management_event_hooks.py index 7a9d155cd67..8952951351c 100644 --- a/tests/unit/proxy/hooks/test_key_management_event_hooks.py +++ b/tests/unit/proxy/hooks/test_key_management_event_hooks.py @@ -6,6 +6,7 @@ Validates that email and secret manager operations are independent and non-block import asyncio import json +from collections.abc import Sequence from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -715,14 +716,14 @@ async def test_key_deleted_hook_writes_audit_log_for_alias_deletion(): captured: Final[list[LiteLLM_AuditLogs]] = [] - async def capture_audit_log(request_data: LiteLLM_AuditLogs) -> None: - captured.append(request_data) + async def capture_audit_logs(rows: Sequence[LiteLLM_AuditLogs]) -> None: + captured.extend(rows) with ( patch("litellm.store_audit_logs", True), patch( - "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", - new=capture_audit_log, + "litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates", + new=capture_audit_logs, ), patch.object( KeyManagementEventHooks, @@ -747,3 +748,34 @@ async def test_key_deleted_hook_writes_audit_log_for_alias_deletion(): assert audit_row.object_id == "hashed" assert audit_row.table_name == LitellmTableNames.KEY_TABLE_NAME assert audit_row.changed_by == "admin" + + +@pytest.mark.asyncio +async def test_deleting_many_keys_writes_every_audit_row_in_one_insert(): + from litellm.proxy._types import LiteLLM_VerificationToken, LitellmTableNames, UserAPIKeyAuth + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks + + tokens: Final = tuple(f"hashed-{index}" for index in range(50)) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_auditlog.create = AsyncMock() + prisma_client.db.litellm_auditlog.create_many = AsyncMock(return_value=len(tokens)) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + ): + KeyManagementEventHooks.create_key_deleted_audit_logs( + keys_being_deleted=[LiteLLM_VerificationToken(token=token) for token in tokens], + user_api_key_dict=UserAPIKeyAuth(user_id="admin", token="callertok"), + ) + await drain_audit_tasks() + + prisma_client.db.litellm_auditlog.create.assert_not_called() + prisma_client.db.litellm_auditlog.create_many.assert_awaited_once() + rows: Final = prisma_client.db.litellm_auditlog.create_many.await_args.kwargs["data"] + assert [row["object_id"] for row in rows] == list(tokens) + assert [json.loads(row["before_value"])["token"] for row in rows] == list(tokens) + assert {(row["table_name"], row["action"], row["changed_by"]) for row in rows} == { + (LitellmTableNames.KEY_TABLE_NAME, "deleted", "admin") + } diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 7af747fc26d..1ec505362dd 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3,6 +3,7 @@ Unit Tests for the max parallel request limiter v3 for the proxy """ import asyncio +import contextvars import logging import sys import time @@ -48,11 +49,13 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_E from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.extraction import ExtractionResponse from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from litellm.types.utils import ( + CallTypes, EmbeddingResponse, ModelResponse, PromptTokensDetailsWrapper, @@ -4205,6 +4208,175 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: ) == 0 +@pytest.mark.asyncio +async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() -> None: + api_key: Final = hash_token("sk-mcp-auto-execute") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = "mcp-auto-execute-owner" + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + def mcp_tool_call() -> dict[str, object]: + return { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Tool: add"}], + "mcp_server_name": "math", + } + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + async def pre_call_from_another_request( + data: dict[str, object], call_type: str + ) -> None: + await contextvars.Context().run(asyncio.create_task, pre_call(data, call_type)) + + await pre_call( + {"model": "gpt-4o-mini", "input": "add 2 and 3", "litellm_call_id": owner_call_id}, + CallTypes.aresponses.value, + ) + await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value) + await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value) + assert await in_flight() == 1 + + with pytest.raises(HTTPException): + await pre_call_from_another_request( + mcp_tool_call(), CallTypes.call_mcp_tool.value + ) + + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id}, + response_obj=None, + start_time=None, + end_time=None, + ) + assert await in_flight() == 0 + await pre_call_from_another_request( + {"model": "gpt-4o-mini", "input": "next request"}, + CallTypes.aresponses.value, + ) + assert await in_flight() == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stream", "in_flight_after_first_round_logs"), + [(False, 1), (True, 0)], + ids=["non_stream_round_logged_late_keeps_slot", "stream_logged_at_its_end_releases_slot"], +) +async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request_end( + stream: bool, in_flight_after_first_round_logs: int +) -> None: + reused_at: Final = datetime(2026, 1, 5, 12, 0, 0) + api_key: Final = hash_token(f"sk-mcp-late-round-log-{stream}") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=lambda: reused_at, + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = f"mcp-late-round-log-{stream}" + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + async def round_logged(ended_at: datetime) -> None: + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id, "stream": stream}, + response_obj=None, + start_time=None, + end_time=ended_at, + ) + + await pre_call( + { + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "add 2 and 3"}], + "litellm_call_id": owner_call_id, + }, + CallTypes.anthropic_messages.value, + ) + await pre_call( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ) + + await round_logged(reused_at - timedelta(seconds=1)) + assert await in_flight() == in_flight_after_first_round_logs + + await round_logged(reused_at + timedelta(seconds=1)) + assert await in_flight() == 0 + + +@pytest.mark.asyncio +async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None: + api_key: Final = hash_token("sk-mcp-reuse-headers") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) + data: Final[dict[str, object]] = { + "model": "gpt-4o-mini", + "input": "add 2 and 3", + "litellm_call_id": "mcp-reuse-headers-owner", + } + for request_data, call_type in ( + (data, CallTypes.aresponses.value), + ( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ), + ): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=request_data, + call_type=call_type, + ) + response: Final = ModelResponse( + id="mcp-reuse-headers", + object="chat.completion", + created=0, + model="gpt-4o-mini", + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + choices=[], + ) + response._hidden_params = {} + + await handler.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + headers: Final = response._hidden_params["additional_headers"] + assert headers["x-ratelimit-api_key-limit-max_parallel_requests"] == 2 + assert headers["x-ratelimit-api_key-remaining-max_parallel_requests"] == 1 + + @pytest.mark.asyncio async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None: from unittest.mock import AsyncMock @@ -8686,3 +8858,78 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( await _rpm_request(limiter, cache, auth, "a2a/cheap") assert denied.value.status_code == 429 assert "expensive" in str(denied.value.detail) + + +def _extraction_request(state: str, call_id: str) -> dict[str, object]: + return { + "model": "invoice-reader", + "state": state, + "fields": {"vendor": {"type": "text", "instructions": "Who sent the invoice"}}, + "litellm_call_id": call_id, + } + + +@pytest.mark.asyncio +async def test_an_extraction_request_passes_a_tpm_limited_key_without_an_injected_output_cap(): + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + key: Final = UserAPIKeyAuth(api_key=hash_token("sk-extraction-cap"), tpm_limit=2000) + data: Final = _extraction_request("Invoice from Acme Tools", "extraction-cap") + sent: Final = dict(data) + + await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=data, call_type="aextraction") + + assert data == sent + + +@pytest.mark.asyncio +async def test_an_extraction_request_reserves_its_state_and_no_output_budget(monkeypatch): + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + key: Final = UserAPIKeyAuth(api_key=hash_token("sk-extraction-estimate"), tpm_limit=1_000_000) + state_tokens: Final = 500 + + reserved: Final = await _reserved_tokens_for( + handler, cache, key, _extraction_request("word" * state_tokens, "extraction-estimate"), call_type="aextraction" + ) + + assert state_tokens <= reserved < 2 * state_tokens + + +@pytest.mark.asyncio +async def test_a_successful_extraction_counts_its_input_tokens_toward_the_key_tpm_limit(monkeypatch): + monkeypatch.setattr(_PROXY_MaxParallelRequestsHandler, "get_rate_limit_type", lambda self: "total") + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + api_key: Final = hash_token("sk-extraction-usage") + key: Final = UserAPIKeyAuth(api_key=api_key, tpm_limit=500) + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=_extraction_request("Invoice", "extraction-1"), call_type="aextraction" + ) + + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": "extraction-1", + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + }, + response_obj=ExtractionResponse.model_validate( + { + "model": "invoice-reader", + "answers": {"vendor": {"type": "text", "value": "Acme Tools"}}, + "truncated": False, + "usage": {"input_tokens": 480, "output_tokens": 0}, + } + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + with pytest.raises(HTTPException) as refused: + await handler.async_pre_call_hook( + user_api_key_dict=key, + cache=cache, + data=_extraction_request("Invoice " * 40, "extraction-2"), + call_type="aextraction", + ) + assert refused.value.status_code == 429 diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 289f8cd4a1a..bd997f10322 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -16,6 +16,7 @@ from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( _get_budget_reservation_from_metadata, @@ -29,6 +30,7 @@ from litellm.proxy.utils import ProxyUpdateSpend from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage @@ -154,6 +156,127 @@ async def test_each_earlier_billed_attempt_gets_its_own_spend_row(): assert row["team_id"] == "t1" +def _deployment(model_name: str, model_id: str, model: str, api_base: str | None) -> Deployment: + return Deployment( + model_name=model_name, + litellm_params=LiteLLM_Params(model=model, api_base=api_base), + model_info={"id": model_id}, + ) + + +_TIMED_OUT_AT: Final = datetime(2026, 10, 20, 21, 24, 48, tzinfo=timezone.utc) + + +async def _spend_log_payloads_of_a_retried_timeout_across_a_fallback() -> list[SpendLogsPayload]: + """The spend-log payloads a timed-out request writes when attempt 0 and 1 ran on the anthropic + deployment of group ``essay`` and the request then fell back to the openai deployment of group + ``essay-fallback``: the final row first, then one row per earlier billed attempt.""" + from litellm.proxy.utils import BilledFailedAttempt + + deployments = { + "dep-anthropic": _deployment("essay", "dep-anthropic", "anthropic/claude-sonnet-4-6", None), + "dep-openai": _deployment("essay-fallback", "dep-openai", "openai/gpt-4.1", "https://api.openai.com/v1"), + } + usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + request_data = { + "model": "anthropic/claude-sonnet-4-6", + "custom_llm_provider": "anthropic", + "litellm_call_id": "call-1", + "litellm_params": {"api_base": "https://api.anthropic.com/v1/messages"}, + "standard_logging_object": { + "id": "call-1", + "model": "anthropic/claude-sonnet-4-6", + "custom_llm_provider": "anthropic", + "api_base": "https://api.anthropic.com/v1/messages", + "model_id": "dep-anthropic", + "model_group": "essay", + "metadata": {}, + "hidden_params": {}, + "model_map_information": {}, + }, + "metadata": {"attempted_retries": 1, "model_group": "essay", "model_info": {"id": "dep-anthropic"}}, + "combined_usage_object": usage, + "response_cost": 7.2e-05, + "billed_failed_attempts": ( + BilledFailedAttempt(attempted_retries=1, deployment_id="dep-anthropic", usage=usage, response_cost=7.2e-05), + BilledFailedAttempt(attempted_retries=0, deployment_id="dep-openai", usage=usage, response_cost=4.8e-05), + ), + } + exception = litellm.Timeout("Request timed out.", model="gpt-4.1", llm_provider="openai") + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await ProxyDBLogger(deployment_lookup=deployments.get).async_post_call_failure_hook( + request_data=request_data, + original_exception=exception, + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key", user_id="u1"), + ) + return [ + get_logging_payload( + kwargs=call.kwargs["kwargs"], + response_obj=exception, + start_time=_TIMED_OUT_AT, + end_time=_TIMED_OUT_AT + timedelta(seconds=11), + ) + for call in mock_update_database.call_args_list + ] + + +@pytest.mark.asyncio +async def test_attempt_rows_are_attributed_to_the_deployment_that_billed_them(): + """Regression since #47416: an earlier attempt that ran on another deployment was written with the + final row's model, provider and api_base. Each attempt row names the deployment it ran on.""" + final_row, attempt_0, attempt_1 = await _spend_log_payloads_of_a_retried_timeout_across_a_fallback() + + assert (final_row["model"], final_row["custom_llm_provider"], final_row["model_id"]) == ( + "anthropic/claude-sonnet-4-6", + "anthropic", + "dep-anthropic", + ) + assert final_row["api_base"] == "https://api.anthropic.com/v1/messages" + assert (attempt_0["model"], attempt_0["custom_llm_provider"], attempt_0["model_id"], attempt_0["model_group"]) == ( + "anthropic/claude-sonnet-4-6", + "anthropic", + "dep-anthropic", + "essay", + ) + assert (attempt_1["model"], attempt_1["custom_llm_provider"], attempt_1["model_id"], attempt_1["model_group"]) == ( + "openai/gpt-4.1", + "openai", + "dep-openai", + "essay-fallback", + ) + assert attempt_1["api_base"] == "https://api.openai.com/v1" + assert attempt_1["request_id"] == "call-1-attempt-1" + assert attempt_1["spend"] == 4.8e-05 + + +@pytest.mark.asyncio +async def test_attempt_rows_add_spend_but_not_requests_to_daily_spend(): + """Regression since #47416: one caller request with N billed earlier attempts counted N+1 requests in + the daily spend tables. The attempt rows keep their spend, tokens and model usage and count zero requests.""" + payloads = await _spend_log_payloads_of_a_retried_timeout_across_a_fallback() + prisma = MagicMock() + prisma.get_request_status = MagicMock(return_value="failure") + + daily = [ + await DBSpendUpdateWriter()._common_add_spend_log_transaction_to_daily_transaction( + payload=payload, prisma_client=prisma, type="user" + ) + for payload in payloads + ] + + assert all(row is not None for row in daily) + assert [row["api_requests"] for row in daily if row is not None] == [1, 0, 0] + assert [row["failed_requests"] for row in daily if row is not None] == [1, 0, 0] + assert [row["spend"] for row in daily if row is not None] == [7.2e-05, 7.2e-05, 4.8e-05] + assert [row["prompt_tokens"] for row in daily if row is not None] == [24, 24, 24] + model_usage = [build_model_usage_transaction(payload) for payload in payloads] + assert all(usage is not None for usage in model_usage), "attempt spend must still reach daily model usage" + assert [usage.spend for usage in model_usage if usage is not None] == [7.2e-05, 7.2e-05, 4.8e-05] + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_metadata(): """ @@ -3202,6 +3325,7 @@ def _grounded_gemini_kwargs(response_cost: float) -> dict: "response_cost": response_cost, "request_tags": None, "cache_hit": None, + "custom_llm_provider": "vertex_ai", "cost_breakdown": { "input_cost": 0.0, "output_cost": 0.0, diff --git a/tests/unit/proxy/hooks/test_search_free_tier.py b/tests/unit/proxy/hooks/test_search_free_tier.py index 3dc89c0dc72..b57908378e5 100644 --- a/tests/unit/proxy/hooks/test_search_free_tier.py +++ b/tests/unit/proxy/hooks/test_search_free_tier.py @@ -31,9 +31,9 @@ class _BrokenCounter: raise ConnectionError("redis is down") -def _model_info(billing_unit: str, count: int, period: str = "month") -> dict: +def _model_info(billing_unit: str, count: int, period: str = "month", provider: str = "vertex_ai") -> dict: return { - "litellm_provider": "vertex_ai", + "litellm_provider": provider, "mode": "chat", "web_search_billing_unit": billing_unit, "search_context_cost_per_query": {"search_context_size_medium": RATE}, @@ -52,6 +52,7 @@ def _call(model_info: dict, response_cost: float, cache_hit: bool | None = None) sl_object = { "response_cost": response_cost, "cache_hit": cache_hit, + "custom_llm_provider": model_info["litellm_provider"], "cost_breakdown": {"input_cost": 0.1, "output_cost": 0.2, "tool_usage_cost": 0.5, "total_cost": response_cost}, "model_map_information": {"model_map_key": "gemini-test", "model_map_value": model_info}, } @@ -73,12 +74,12 @@ def test_freed_units_covers_only_what_landed_under_the_allowance(allowance, unit assert freed_units(allowance, units, total_after) == expected -def test_counter_key_is_per_group_and_utc_period(): - assert search_free_tier_key({"count": 1, "period": "month", "shared_group": "g3"}, NOW) == ( - "litellm:search_free_tier:g3:2026-10" +def test_counter_key_is_per_provider_group_and_utc_period(): + assert search_free_tier_key("gemini", {"count": 1, "period": "month", "shared_group": "g3"}, NOW) == ( + "litellm:search_free_tier:gemini:g3:2026-10" ) - assert search_free_tier_key({"count": 1, "period": "day", "shared_group": "g2"}, NOW) == ( - "litellm:search_free_tier:g2:2026-10-20" + assert search_free_tier_key("vertex_ai", {"count": 1, "period": "day", "shared_group": "g2"}, NOW) == ( + "litellm:search_free_tier:vertex_ai:g2:2026-10-20" ) @@ -107,8 +108,23 @@ async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_cross assert c2 == k2 == s2 == pytest.approx(1.0 - 2 * RATE) assert b2["discount_amount"] == pytest.approx(2 * RATE) assert c3 == k3 == s3 == 1.0 and "discount_amount" not in b3 - assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 15} - assert counter.ttls["litellm:search_free_tier:test-group:2026-10"] > 0 + assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 15} + assert counter.ttls["litellm:search_free_tier:vertex_ai:test-group:2026-10"] > 0 + + +@pytest.mark.asyncio +async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances(): + counter = _Counter() + gemini_api = _model_info("per_query", count=3, provider="gemini") + vertex_ai = _model_info("per_query", count=3, provider="vertex_ai") + kwargs, sl_object = _call(gemini_api, response_cost=1.0) + await apply_search_free_tier(kwargs, sl_object, 1.0, _response(2), counter, NOW) + kwargs, sl_object = _call(vertex_ai, response_cost=1.0) + vertex_cost = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(4), counter, NOW) + kwargs, sl_object = _call(gemini_api, response_cost=1.0) + gemini_cost = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(1), counter, NOW) + assert vertex_cost == pytest.approx(1.0 - 3 * RATE), "vertex ai had its whole allowance left" + assert gemini_cost == pytest.approx(1.0 - RATE), "gemini api had one free query left" @pytest.mark.asyncio @@ -121,7 +137,7 @@ async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_q second = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) assert first == pytest.approx(1.0 - RATE) assert second == 1.0 - assert counter.totals == {"litellm:search_free_tier:test-group:2026-10-20": 2} + assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10-20": 2} @pytest.mark.asyncio @@ -172,7 +188,7 @@ async def test_a_second_callback_run_on_the_same_request_counts_it_once(): first = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW) second = await apply_search_free_tier(kwargs, sl_object, first, _response(5), counter, NOW) assert first == second == pytest.approx(1.0 - 5 * RATE) - assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 5} + assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5} @pytest.mark.asyncio @@ -191,4 +207,4 @@ async def test_freed_cost_never_exceeds_the_billed_cost(): kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=billed) assert await apply_search_free_tier(kwargs, sl_object, billed, _response(5), counter, NOW) == 0.0 assert sl_object["cost_breakdown"]["discount_amount"] == pytest.approx(billed) - assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 5} + assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5} diff --git a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index cf505e2f98e..6014c9229f6 100644 --- a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -3,7 +3,7 @@ import asyncio import json import logging import time -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from itertools import chain from types import MappingProxyType from typing import Final, assert_never @@ -58,6 +58,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( update_user, user_api_key_auth, ) +from litellm.proxy.management_endpoints.team_endpoints import DeletedTeams from litellm.proxy.management_helpers.bulk_user_creation import BULK_NEW_USER_CONCURRENCY from litellm.proxy.management_helpers.team_roster_sync import ( MembersMissing, @@ -2831,16 +2832,24 @@ async def test_patch_user_grants_admin_by_team_display_name(mocker, monkeypatch) assert call_args[1]["data"]["user_role"] == LitellmUserRoles.PROXY_ADMIN -def _scim_admin_prisma(mocker, *, user_teams): - """Prisma double and its one transaction, in which the user resolves to user_teams - and teams expose an alias equal to their id. Only the transaction answers reads, - so a role written outside it fails the role-recompute helper tests.""" - user = mocker.MagicMock() - user.user_id = "member-1" - user.teams = user_teams +def _member_row(mocker, user_id: str, teams: Sequence[str]) -> MagicMock: + row = mocker.MagicMock() + row.user_id = user_id + row.teams = list(teams) + return row + + +def _scim_admin_prisma(mocker, *, user_teams, teammates: Sequence[tuple[str, Sequence[str]]] = ()): + """Prisma double and its one transaction, in which the user resolves to user_teams, + ``teammates`` are further (user_id, teams) rows, and teams expose an alias equal to + their id. Only the transaction answers reads, so a role written outside it fails the + role-recompute helper tests.""" + user = _member_row(mocker, "member-1", user_teams) + members = (user, *(_member_row(mocker, user_id, teams) for user_id, teams in teammates)) def _team_find_unique(where): team = mocker.MagicMock() + team.team_id = where["team_id"] team.team_alias = where["team_id"] return team @@ -2848,7 +2857,7 @@ def _scim_admin_prisma(mocker, *, user_teams): return tuple(_team_find_unique({"team_id": team_id}) for team_id in where["team_id"]["in"]) def _users_find_many(where): - return (user,) if user.user_id in where["user_id"]["in"] else () + return tuple(member for member in members if member.user_id in where["user_id"]["in"]) tx = mocker.MagicMock() tx.query_raw = AsyncMock(return_value=[]) @@ -2904,6 +2913,29 @@ async def test_recompute_scim_member_roles_grants_when_in_admin_group(mocker, mo tx.execute_raw.assert_awaited_once_with(ANY, LitellmUserRoles.PROXY_ADMIN, ["member-1"]) +@pytest.mark.asyncio +async def test_recompute_scim_member_roles_stops_serving_a_demoted_members_cached_admin_role(mocker, monkeypatch): + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + cache = UserApiKeyCache() + await cache.async_set_cache( + key="member-1", value={"user_id": "member-1", "user_role": LitellmUserRoles.PROXY_ADMIN.value} + ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + prisma, tx = _scim_admin_prisma(mocker, user_teams=["engineering"]) + + await _recompute_scim_member_roles(prisma, ["member-1"]) + + tx.execute_raw.assert_awaited_once_with(ANY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ["member-1"]) + assert await cache.async_get_cache(key="member-1") is None + + @pytest.mark.asyncio async def test_recompute_scim_member_roles_noop_when_admin_group_unset(mocker, monkeypatch): """With scim_admin_group unset the recompute helper must not touch any role, @@ -3043,44 +3075,128 @@ async def test_patch_group_recomputes_roles_for_changed_members(mocker): assert set(recompute_mock.call_args[0][1]) == {"user1"} -@pytest.mark.asyncio -async def test_delete_group_recomputes_roles_for_members(mocker): - """DELETE /Groups must recompute the global role for the team's members, so - deleting the admin group demotes everyone who was only admin through it.""" - from litellm.proxy._types import Member +_SCIM_CALLER: Final = UserAPIKeyAuth(api_key="hashed-scim-token", user_id="idp-service-user") +_ADMIN_GROUP: Final = "litellm-admins" +_DEMOTED: Final = call(ANY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ["member-1"]) - existing_team = mocker.MagicMock() - existing_team.members_with_roles = [ - Member(user_id="user1", role="user"), - Member(user_id="user2", role="user"), - ] - member = mocker.MagicMock() - member.teams = ["test-team-123"] +def _group_delete_mocks( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + team: LiteLLM_TeamTable | None, + *, + teammates: Sequence[tuple[str, Sequence[str]]] = (), +) -> tuple[MagicMock, MagicMock, AsyncMock]: + """Prisma and its transaction, whose users are left in no team once the delete detached + them and whose group lookup answers ``team``, plus the patched team delete, which hands + its transaction and the roster it read under the lock, member-1 and every teammate, to + the caller's write.""" + from litellm.proxy.proxy_server import proxy_config - mock_prisma_client = mocker.MagicMock() - mock_prisma_client.db = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) - mock_prisma_client.db.litellm_teamtable.delete = AsyncMock() - mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=member) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) - mock_prisma_client.db.litellm_usertable.update = AsyncMock() + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": _ADMIN_GROUP}} + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + prisma, tx = _scim_admin_prisma(mocker, user_teams=[], teammates=teammates) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", - AsyncMock(return_value=mock_prisma_client), + AsyncMock(return_value=prisma), ) - recompute_mock = mocker.patch( - "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", - AsyncMock(), + roster: Final = ("member-1", *(user_id for user_id, _ in teammates)) + + async def delete_rows( + *, before_commit: Callable[[object, tuple[str, ...]], Awaitable[None]] | None = None, **_: object + ) -> DeletedTeams: + if before_commit is not None: + await before_commit(tx, roster) + return DeletedTeams(team_ids=(_ADMIN_GROUP,)) + + delete_team_mock: Final = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.delete_validated_team_rows", + AsyncMock(side_effect=delete_rows), ) + return prisma, tx, delete_team_mock - await delete_group(group_id="test-team-123") - recompute_mock.assert_awaited_once() - assert list(recompute_mock.call_args[0][1]) == ["user1", "user2"] +def _admin_group_team() -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=_ADMIN_GROUP, members_with_roles=[Member(user_id="member-1", role="user")]) + + +@pytest.mark.asyncio +async def test_delete_group_demotes_a_member_who_was_admin_only_through_the_deleted_group( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch +): + """DELETE /Groups answers 204 and writes the non-admin role for a member who was + admin only through the deleted group, once, inside the team delete.""" + _, tx, _ = _group_delete_mocks(mocker, monkeypatch, _admin_group_team()) + + response: Final = await delete_group(group_id=_ADMIN_GROUP, user_api_key_dict=_SCIM_CALLER) + + assert response.status_code == 204 + assert tx.execute_raw.await_args_list == [_DEMOTED] + + +@pytest.mark.asyncio +async def test_delete_group_demotes_every_member_on_the_roster_the_team_delete_read_under_the_lock( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch +): + """A member a concurrent PUT added to the admin group after this delete read the + group is on the roster the team delete read under the lock, so the one role write + covers that late member too.""" + _, tx, _ = _group_delete_mocks(mocker, monkeypatch, _admin_group_team(), teammates=(("member-2", ()),)) + + response: Final = await delete_group(group_id=_ADMIN_GROUP, user_api_key_dict=_SCIM_CALLER) + + assert response.status_code == 204 + assert tx.execute_raw.await_args_list == [ + call(ANY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ["member-1", "member-2"]), + ] + + +@pytest.mark.asyncio +async def test_delete_group_fails_when_the_role_writes_fail(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """The role writes run on the team delete's transaction, so when they fail the + request fails with them and the identity provider's retry finds the group.""" + _, tx, _ = _group_delete_mocks(mocker, monkeypatch, _admin_group_team()) + tx.execute_raw.side_effect = RuntimeError("database unavailable") + + with pytest.raises(ProxyException) as raised: + await delete_group(group_id=_ADMIN_GROUP, user_api_key_dict=_SCIM_CALLER) + + assert raised.value.code == "500" + + +@pytest.mark.asyncio +async def test_delete_group_answers_404_when_a_concurrent_delete_removed_the_team_first( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch +): + """A concurrent delete that removed the group while this one waited on the team lock + leaves the team delete nothing to delete or report, so this request answers 404, the + same as a group that never existed.""" + _, _, delete_team_mock = _group_delete_mocks(mocker, monkeypatch, _admin_group_team()) + delete_team_mock.side_effect = None + delete_team_mock.return_value = DeletedTeams(team_ids=()) + + with pytest.raises(ProxyException) as raised: + await delete_group(group_id=_ADMIN_GROUP, user_api_key_dict=_SCIM_CALLER) + + assert raised.value.code == "404" + + +@pytest.mark.asyncio +async def test_delete_group_of_an_unknown_group_answers_404_without_deleting_or_touching_roles( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch +): + _, tx, delete_team_mock = _group_delete_mocks(mocker, monkeypatch, None) + + with pytest.raises(ProxyException) as raised: + await delete_group(group_id="missing-team", user_api_key_dict=_SCIM_CALLER) + + assert raised.value.code == "404" + delete_team_mock.assert_not_awaited() + tx.execute_raw.assert_not_called() @pytest.mark.asyncio @@ -3121,12 +3237,10 @@ async def test_delete_group_is_refused_while_the_team_holds_a_ptu_share(mocker, monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", shared_router) with pytest.raises(ProxyException) as refused: - await delete_group(group_id="test-team-123") + await delete_group(group_id="test-team-123", user_api_key_dict=_SCIM_CALLER) assert refused.value.code == "400" assert "test-team-123" in str(refused.value.message) and "shared-ptu" in str(refused.value.message) - mock_prisma_client.db.litellm_teamtable.delete.assert_not_called() - mock_prisma_client.db.litellm_usertable.update.assert_not_called() @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 037ccd26236..cb8a73d581b 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -5123,11 +5123,11 @@ async def test_delete_user_writes_deleted_audit_log_for_user_keys(mocker): mocker.patch("litellm.store_audit_logs", True) captured: Final[list] = [] - async def _capture(request_data): - captured.append(request_data) + async def _capture(rows): + captured.extend(rows) mocker.patch( - "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + "litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates", new=_capture, ) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 5063ae0b357..37028558af3 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, + LiteLLM_VerificationToken, LitellmTableNames, LitellmUserRoles, Member, @@ -50,7 +51,9 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( SUPPORTED_TEAM_ADMIN_PERMISSIONS, ) from litellm.proxy.management_endpoints.team_endpoints import ( + _LOCK_DELETED_TEAM_MEMBERS_SQL, _STRIP_DELETED_TEAMS_FROM_USERS_SQL, + DeletedTeams, GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, _build_team_list_where_conditions, @@ -63,6 +66,8 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _validate_team_member_reset_spend_value, aggregated_date_range_error, delete_team, + delete_validated_team_rows, + delete_validated_teams, list_available_teams, new_team, reset_team_member_budget_fn, @@ -212,21 +217,50 @@ def _wire_member_delete_tx(prisma_client): return tx -def _wire_team_delete_tx(prisma_client): - """`/team/delete` deletes the team rows and runs its post-delete reference sweep under - every team's advisory lock in one transaction, so a mocked client has to hand its own - table mocks (and db-level execute_raw) back out of `tx()` for existing per-table - assertions on `prisma_client.db.*` to keep seeing those calls.""" +def _wire_team_delete_tx( + prisma_client, + *, + present_team_ids=None, + locked_teams: tuple[LiteLLM_TeamTable, ...] = (), + team_keys: tuple[LiteLLM_VerificationToken, ...] = (), +): + """`/team/delete` records, deletes and sweeps the team rows under every team's advisory + lock in one transaction, so a mocked client has to hand its own table mocks back out of + `tx()` for existing per-table assertions on `prisma_client.db.*` to keep seeing those + calls. Under the lock every requested team still exists unless `present_team_ids` names + the ones left, and reads as the matching `locked_teams` entry or a bare row. The keys the + lock finds on those teams are `team_keys`. The locked sweep reads the users it detaches + through a `query_raw` whose rows name them and their teams, so the default answers none.""" + locked_by_id = {team.team_id: team for team in locked_teams} + + async def teams_still_there(where: dict[str, dict[str, list[str]]]) -> list[LiteLLM_TeamTable]: + return [ + locked_by_id.get(team_id, LiteLLM_TeamTable(team_id=team_id)) + for team_id in where["team_id"]["in"] + if present_team_ids is None or team_id in present_team_ids + ] + + async def keys_on_teams(where: dict[str, dict[str, list[str]]]) -> list[LiteLLM_VerificationToken]: + return [key for key in team_keys if key.team_id in where["team_id"]["in"]] + tx = SimpleNamespace( - litellm_teamtable=prisma_client.db.litellm_teamtable, + litellm_teamtable=SimpleNamespace( + find_many=teams_still_there, + delete_many=prisma_client.db.litellm_teamtable.delete_many, + ), + litellm_verificationtoken=SimpleNamespace(find_many=keys_on_teams, delete_many=AsyncMock(return_value=0)), + litellm_deletedverificationtoken=SimpleNamespace(create_many=AsyncMock()), + litellm_jwtkeymapping=SimpleNamespace(find_many=AsyncMock(return_value=[])), litellm_teammembership=prisma_client.db.litellm_teammembership, + litellm_deletedteamtable=prisma_client.db.litellm_deletedteamtable, query_raw=AsyncMock(return_value=[]), - execute_raw=prisma_client.db.execute_raw, + execute_raw=AsyncMock(return_value=0), ) tx_cm = MagicMock() tx_cm.__aenter__ = AsyncMock(return_value=tx) tx_cm.__aexit__ = AsyncMock(return_value=None) prisma_client.tx = MagicMock(return_value=tx_cm) + return tx # Mock prisma_client @@ -5489,13 +5523,13 @@ async def test_team_member_delete_writes_deleted_audit_log_for_member_keys( captured: Final[list] = [] - async def _capture(request_data): - captured.append(request_data) + async def _capture(rows): + captured.extend(rows) with ( patch("litellm.store_audit_logs", True), patch( - "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + "litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates", new=_capture, ), ): @@ -5564,12 +5598,12 @@ async def test_delete_team_writes_deleted_audit_log_for_team_keys( ) captured: Final[list] = [] - async def _capture(request_data): - captured.append(request_data) + async def _capture(rows): + captured.extend(rows) monkeypatch.setattr("litellm.store_audit_logs", True) monkeypatch.setattr( - "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + "litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates", _capture, ) await delete_team( @@ -9050,7 +9084,7 @@ async def test_save_deleted_team_records(): }, ] - await _save_deleted_team_records(records=records, prisma_client=mock_prisma_client) + await _save_deleted_team_records(records=records, deleted_teams=mock_prisma_client.db.litellm_deletedteamtable) mock_create_many.assert_called_once_with(data=records) @@ -9061,7 +9095,7 @@ async def test_save_deleted_team_records_empty_list(): mock_create_many = AsyncMock() mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many - await _save_deleted_team_records(records=[], prisma_client=mock_prisma_client) + await _save_deleted_team_records(records=[], deleted_teams=mock_prisma_client.db.litellm_deletedteamtable) mock_create_many.assert_not_called() @@ -9091,7 +9125,7 @@ async def test_persist_deleted_team_records(): await _persist_deleted_team_records( teams=[team], - prisma_client=mock_prisma_client, + deleted_teams=mock_prisma_client.db.litellm_deletedteamtable, user_api_key_dict=user_api_key_dict, litellm_changed_by="admin-user", ) @@ -9295,7 +9329,7 @@ async def test_delete_team_sweeps_references_outside_members_with_roles( # The locked delete-and-sweep transaction /team/member_add serializes against, kept # separate from mock_tx above (the BYOK-model-cleanup transaction, unrelated to this lock). - _wire_team_delete_tx(mock_prisma_client) + _wire_team_delete_tx(mock_prisma_client, locked_teams=(doomed_team,)) mock_lock_tx = mock_prisma_client.tx.return_value.__aenter__.return_value fresh_cache = UserApiKeyCache() @@ -9325,24 +9359,23 @@ async def test_delete_team_sweeps_references_outside_members_with_roles( litellm_changed_by="admin-user", ) - assert mock_execute_raw.await_args_list == [ - call(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, ["team-doomed"]), - call(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, ["team-doomed"]), + assert mock_execute_raw.await_args_list == [], ( + "no sweep runs outside the lock: one there strips a member a concurrent roster write added " + "after the caller read the team, before the locked sweep can report them" + ) + + # one pass for the membership rows, under the lock alongside the delete + assert mock_membership_delete_many.await_args_list == [call(where={"team_id": {"in": ("team-doomed",)}})] + + assert mock_lock_tx.query_raw.await_args_list == [ + call(TEAM_ADVISORY_LOCK_SQL, "team-doomed"), + call(_LOCK_DELETED_TEAM_MEMBERS_SQL, ["team-doomed"]), ], ( - "the unlocked sweep must run once to catch pre-existing drift, and the locked sweep " - "(alongside the delete, under the same advisory lock member_add takes) must run again " - "so a member_add that wrote its reference just before losing the lock is still reaped" - ) - - # same two passes for the membership rows, the second under the lock alongside the delete - assert mock_membership_delete_many.await_args_list == [ - call(where={"team_id": {"in": ("team-doomed",)}}), - call(where={"team_id": {"in": ("team-doomed",)}}), - ] - - assert mock_lock_tx.query_raw.await_args_list == [call(TEAM_ADVISORY_LOCK_SQL, "team-doomed")], ( - "the advisory lock must be acquired before the team row is deleted" + "the advisory lock must be acquired before the team row is deleted, and the user sweep " + "runs under it (the same lock member_add takes) so a member_add that wrote its " + "reference just before losing the lock is still reaped" ) + assert mock_lock_tx.execute_raw.await_args_list == [call(_STRIP_DELETED_TEAMS_FROM_USERS_SQL, ["team-doomed"])] assert fresh_cache.get_cache(key="team_id:team-doomed") is None assert fresh_cache.get_cache(key="team_alias:doomed-team") is None @@ -9417,7 +9450,7 @@ async def test_delete_team_evicts_member_caches_with_one_transaction( mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) - _wire_team_delete_tx(mock_prisma_client) + _wire_team_delete_tx(mock_prisma_client, locked_teams=(team,)) fresh_cache = UserApiKeyCache() for user_id in member_user_ids: @@ -9461,6 +9494,234 @@ async def test_delete_team_evicts_member_caches_with_one_transaction( ) +@pytest.mark.asyncio +async def test_delete_validated_team_rows_records_the_roster_the_lock_read_and_writes_its_users_before_commit( + monkeypatch, + disable_audit_logging_for_mocked_team, +): + """ + A roster write that lands between the caller's read of the team and the locked delete adds + a member the caller never saw, and a user row can reference the team without a roster entry. + The deleted-team record and the caller's write, which runs on the delete's transaction before + it commits, both take the roster read under the lock, so they name the late member and not the + stray user, while both lose their cached user object. + """ + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + read_by_caller = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[Member(user_id="member-1", role="user")], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + read_under_lock = read_by_caller.model_copy( + update={ + "members_with_roles": [ + Member(user_id="member-1", role="user"), + Member(user_id="late-member", role="user"), + ] + } + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 0}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + tx = _wire_team_delete_tx(mock_prisma_client, locked_teams=(read_under_lock,)) + tx.query_raw = AsyncMock( + side_effect=[ + [], + [ + {"user_id": "member-1", "teams": ["team-kept", "team-doomed"]}, + {"user_id": "late-member", "teams": ["team-doomed"]}, + {"user_id": "stray-user", "teams": ["team-doomed"]}, + ], + ] + ) + + fresh_cache = UserApiKeyCache() + for user_id in ("member-1", "late-member", "stray-user", "bystander-user"): + fresh_cache.set_cache(key=user_id, value=UserAPIKeyAuth(user_id=user_id)) + + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + async def write_before_commit(_tx: object, _user_ids: tuple[str, ...]) -> None: + mock_prisma_client.tx.return_value.__aexit__.assert_not_awaited() + + before_commit = AsyncMock(side_effect=write_before_commit) + + deleted = await delete_validated_team_rows( + teams=(read_by_caller,), + prisma_client=mock_prisma_client, + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + before_commit=before_commit, + ) + + assert deleted == DeletedTeams(team_ids=("team-doomed",)) + before_commit.assert_awaited_once_with(tx, ("member-1", "late-member")) + (written,) = mock_prisma_client.db.litellm_deletedteamtable.create_many.await_args_list + (record,) = written.kwargs["data"] + assert [member["user_id"] for member in json.loads(record["members_with_roles"])] == ["member-1", "late-member"] + for user_id in ("member-1", "late-member", "stray-user"): + assert fresh_cache.get_cache(key=user_id) is None, user_id + assert fresh_cache.get_cache(key="bystander-user") is not None + + +@pytest.mark.asyncio +async def test_delete_validated_team_rows_deletes_only_the_keys_its_rostered_members_hold_on_the_team( + monkeypatch, + disable_audit_logging_for_mocked_team, +): + """ + The SCIM group delete removes the team the way a SCIM member removal removes each member: + the keys rostered members hold on the team are tombstoned and deleted, and a key with no + user or held by someone off the roster stays. + """ + team = LiteLLM_TeamTable( + team_id="team-doomed", + members_with_roles=[Member(user_id="member-1", role="user")], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + team_keys = tuple( + LiteLLM_VerificationToken(token=token, team_id="team-doomed", user_id=user_id) + for token, user_id in (("hashed-member", "member-1"), ("hashed-service-account", None), ("hashed-outsider", "admin-made")) + ) + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + mock_access_group_tx_cm = MagicMock() + mock_access_group_tx_cm.__aenter__ = AsyncMock(return_value=AsyncMock()) + mock_access_group_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_access_group_tx_cm) + tx = _wire_team_delete_tx(mock_prisma_client, locked_teams=(team,), team_keys=team_keys) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_validated_team_rows( + teams=(team,), + prisma_client=mock_prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-scim-token"), + litellm_changed_by=None, + ) + + (tombstones,) = tx.litellm_deletedverificationtoken.create_many.await_args_list + assert [record["token"] for record in tombstones.kwargs["data"]] == ["hashed-member"] + (deleted_keys,) = tx.litellm_verificationtoken.delete_many.await_args_list + assert deleted_keys.kwargs["where"] == {"token": {"in": ["hashed-member"]}} + + +@pytest.mark.asyncio +async def test_delete_validated_teams_writes_the_deleted_audit_row_only_for_a_team_still_there_under_the_lock( + monkeypatch: pytest.MonkeyPatch, +): + """Two deletes racing on one team both read it before the lock, and only the one that still + finds it under the lock deletes it, so only that one writes its deleted audit row.""" + audit_logger = _wire_audit_log_callback(monkeypatch) + teams = tuple( + LiteLLM_TeamTable(team_id=team_id, metadata={}, model_max_budget={}, model_spend={}) + for team_id in ("team-raced", "team-kept-for-us") + ) + mock_prisma = AsyncMock() + mock_prisma.delete_data = AsyncMock(return_value={"deleted_keys": 0}) + mock_prisma.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_auditlog.create = AsyncMock() + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma.db.tx = MagicMock(return_value=mock_tx_cm) + _wire_team_delete_tx(mock_prisma, present_team_ids={"team-kept-for-us"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_validated_teams( + teams=teams, + prisma_client=mock_prisma, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1", api_key="sk-a"), + litellm_changed_by=None, + ) + await _settle_audit_log_tasks() + + assert [ + (p["object_id"], p["action"]) + for p in audit_logger.payloads + if p["table_name"] == LitellmTableNames.TEAM_TABLE_NAME + ] == [("team-kept-for-us", "deleted")] + + +@pytest.mark.asyncio +async def test_delete_validated_teams_records_and_reports_only_the_teams_still_there_under_the_lock( + monkeypatch, + disable_audit_logging_for_mocked_team, +): + """ + A concurrent delete that removed one of the teams while this one waited on the lock + already wrote that team's deleted-team record, so this delete writes records only for + the teams it still finds under the lock and reports only those as deleted. + """ + teams = tuple( + LiteLLM_TeamTable( + team_id=team_id, + team_alias=team_id, + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + for team_id in ("team-raced", "team-kept-for-us") + ) + mock_prisma_client = AsyncMock() + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 0}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + _wire_team_delete_tx(mock_prisma_client, present_team_ids={"team-kept-for-us"}) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + deleted = await delete_validated_teams( + teams=teams, + prisma_client=mock_prisma_client, + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by=None, + ) + + assert deleted.team_ids == ("team-kept-for-us",) + (written,) = mock_prisma_client.db.litellm_deletedteamtable.create_many.await_args_list + assert [record["team_id"] for record in written.kwargs["data"]] == ["team-kept-for-us"] + + @pytest.mark.asyncio async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes( monkeypatch, @@ -9558,8 +9819,7 @@ async def test_delete_team_failing_locked_sweep_rolls_back_the_delete_and_leaves mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() - # the unlocked pre-delete sweep succeeds, the locked post-delete sweep blows up - mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[None, ConnectionError("db went away")]) + mock_prisma_client.db.execute_raw = AsyncMock() mock_tx = AsyncMock() mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) @@ -9568,6 +9828,10 @@ async def test_delete_team_failing_locked_sweep_rolls_back_the_delete_and_leaves mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) _wire_team_delete_tx(mock_prisma_client) + # the advisory lock is taken, then the locked post-delete user sweep blows up + mock_prisma_client.tx.return_value.__aenter__.return_value.query_raw = AsyncMock( + side_effect=[[], ConnectionError("db went away")] + ) fresh_cache = UserApiKeyCache() cached_obj = LiteLLM_TeamTableCachedObj(team_id="team-doomed", team_alias="doomed-team") @@ -9607,21 +9871,26 @@ class _BroadcastPipeline: return self async def execute(self) -> None: - for message in self._messages: - await self._redis.publish("", message) + await self._redis.send(tuple(self._messages)) class _BroadcastRecordingRedis: namespace = None - def __init__(self) -> None: + def __init__(self, round_trip_seconds: float = 0.0) -> None: self.broadcast_keys: list[str] = [] + self._round_trip_seconds = round_trip_seconds def init_pubsub_client(self) -> "_BroadcastRecordingRedis": return self + async def send(self, messages: Sequence[str]) -> None: + if self._round_trip_seconds: + await asyncio.sleep(self._round_trip_seconds) + self.broadcast_keys.extend(json.loads(message)["cache_key"] for message in messages) + async def publish(self, channel: str, message: str) -> int: - self.broadcast_keys.append(json.loads(message)["cache_key"]) + await self.send((message,)) return 1 def pipeline(self, transaction: bool = True) -> _BroadcastPipeline: @@ -9667,7 +9936,7 @@ async def test_delete_team_broadcasts_cache_invalidation_to_other_workers( mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) mock_tx_cm.__aexit__ = AsyncMock(return_value=False) mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) - _wire_team_delete_tx(mock_prisma_client) + _wire_team_delete_tx(mock_prisma_client, locked_teams=(team,)) import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module @@ -9697,6 +9966,44 @@ async def test_delete_team_broadcasts_cache_invalidation_to_other_workers( assert redis.broadcast_keys == ["hashed-doomed-key", "team_id:team-doomed", "team_alias:doomed-team"] +@pytest.mark.asyncio +async def test_deleting_a_500_member_team_reaches_every_member_cache_on_other_workers( + monkeypatch: pytest.MonkeyPatch, +): + from itertools import chain + + import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module + from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, + ) + from litellm.proxy.management_endpoints.team_endpoints import _evict_deleted_team_member_cache + + redis: Final = _BroadcastRecordingRedis(round_trip_seconds=0.001) + monkeypatch.setattr(pubsub_module, "coordination_redis_cache", lambda: redis) + monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) + monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16)) + members: Final = tuple(f"member-{index:03d}" for index in range(500)) + + await _evict_deleted_team_member_cache( + team_id="team-large", member_user_ids=members, user_api_key_cache=UserApiKeyCache() + ) + await asyncio.gather(*pubsub_module._pending_publishes) # pyright: ignore[reportPrivateUsage] # drain module-level tasks + + expected: Final = frozenset( + chain.from_iterable( + ( + member, + team_membership_auth_cache_key(team_id="team-large", user_id=member), + team_membership_reservation_cache_key(user_id=member, team_id="team-large"), + ) + for member in members + ) + ) + assert expected - frozenset(redis.broadcast_keys) == frozenset() + + @pytest.mark.asyncio async def test_delete_team_survives_a_failing_cache_backend( monkeypatch, @@ -15783,6 +16090,8 @@ async def test_new_team_and_delete_team_both_drive_the_mirror( prisma.delete_data = AsyncMock(return_value=[team_row]) prisma.db.execute_raw = AsyncMock(return_value=0) prisma.db.litellm_teammembership.delete_many = AsyncMock(return_value=0) + prisma.db.litellm_teamtable.delete_many = AsyncMock(return_value=1) + _wire_team_delete_tx(prisma, locked_teams=(team_row,)) await delete_team( data=DeleteTeamRequest(team_ids=["team-gone"]), diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 873182a5145..479391e54fe 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -33,6 +33,7 @@ from litellm.proxy.management_endpoints.ui_sso import ( insert_sso_user, normalize_email, process_sso_jwt_access_token, + read_return_to_cookie, sso_lookup_user_ids, ) from litellm.types.proxy.management_endpoints.ui_sso import ( @@ -10603,6 +10604,10 @@ class TestPersistReturnToCookieSharedHelper: def _cookie(resp) -> str: return resp.headers.get("set-cookie", "") + @pytest.fixture(autouse=True) + def _master_key(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-master") + def test_sets_cookie_for_same_origin_relative_path(self, monkeypatch): from fastapi import Response @@ -10648,6 +10653,62 @@ class TestPersistReturnToCookieSharedHelper: _persist_return_to_cookie(resp, "https://cp.example.com/ui?page=models", _make_https_request()) assert "litellm_cp_return_to=" in self._cookie(resp) + def test_cookie_is_opaque_and_round_trips_the_loopback_redirect_uri(self, monkeypatch: pytest.MonkeyPatch) -> None: + """The DCR authorize round-trip for a local coding agent stores a return_to carrying + ``redirect_uri=http://127.0.0.1:/...``. An edge firewall (AWS WAF + ``EC2MetaDataSSRF_COOKIE``) reads that loopback URL in a cookie as an SSRF probe and 403s the + callback, so the cookie must carry nothing a middlebox can pattern-match, while the proxy + itself still reads the exact return_to back.""" + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie, read_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + return_to: Final = "/mcp/authorize?client_id=llm_dcrc_abc&redirect_uri=http%3A%2F%2F127.0.0.1%3A1455%2Fauth%2Fcallback" + resp: Final = Response() + _persist_return_to_cookie(resp, return_to, _make_https_request()) + cookie: Final = self._cookie(resp) + value: Final = cookie.split("litellm_cp_return_to=", 1)[1].split(";", 1)[0] + assert "127.0.0.1" not in cookie and "authorize" not in cookie, cookie + request: Final = _make_https_request() + request.cookies = {"litellm_cp_return_to": value} + assert read_return_to_cookie(request) == return_to + + def test_a_cookie_this_proxy_did_not_seal_reads_as_nothing(self) -> None: + from litellm.proxy.management_endpoints.ui_sso import read_return_to_cookie + + request: Final = _make_https_request() + request.cookies = {"litellm_cp_return_to": "/mcp/authorize?client_id=llm_dcrc_abc"} + assert read_return_to_cookie(request) is None + + def test_no_master_key_means_no_cookie_and_no_error(self, monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + resp: Final = Response() + _persist_return_to_cookie(resp, "/mcp/authorize", _make_https_request()) + assert "litellm_cp_return_to=" not in self._cookie(resp) + + def test_sso_dashboard_fallback_clears_a_return_to_cookie_it_rejected(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A plaintext (pre-upgrade or client-planted) litellm_cp_return_to is not honored by /sso/callback, + and the dashboard landing must also expire it so the browser stops replaying the URL-bearing cookie + (and tripping the edge firewall) for the rest of its lifetime.""" + from litellm.proxy.management_endpoints.ui_sso import _dashboard_redirect_response + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + request: Final = _make_https_request() + request.cookies = {"litellm_cp_return_to": "/mcp/authorize?client_id=llm_dcrc_abc"} + assert read_return_to_cookie(request) is None + response: Final = _dashboard_redirect_response(request, "jwt", "/ui?login=success") + set_cookies: Final = response.headers.getlist("set-cookie") + assert response.status_code == 303 and response.headers["location"] == "/ui?login=success" + assert any(c.startswith("token=jwt") for c in set_cookies), set_cookies + assert any(c.startswith('litellm_cp_return_to=""') and "Max-Age=0" in c for c in set_cookies), set_cookies + def test_cookie_is_secure_and_httponly_over_https(self, monkeypatch): from fastapi import Response diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py b/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py index 74b4a3caf0a..e212a433393 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py @@ -1007,12 +1007,12 @@ async def test_bulk_delete_writes_deleted_audit_log_for_deleted_keys(mocker): captured: Final[list] = [] - async def _capture(request_data): - captured.append(request_data) + async def _capture(rows): + captured.extend(rows) mocker.patch("litellm.store_audit_logs", True) mocker.patch( - "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + "litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates", new=_capture, ) await _delete(prisma, ["u1"]) @@ -1041,12 +1041,12 @@ async def test_bulk_member_delete_writes_deleted_audit_log_for_removed_team_keys captured: Final[list] = [] - async def _capture(request_data): - captured.append(request_data) + async def _capture(rows): + captured.extend(rows) mocker.patch("litellm.store_audit_logs", True) mocker.patch( - "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + "litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates", new=_capture, ) await _remove(prisma, "t1", [{"user_id": "u1"}]) @@ -1076,7 +1076,7 @@ async def test_bulk_delete_skips_the_key_audit_log_when_the_tx_rolls_back(mocker mocker.patch("litellm.store_audit_logs", True) mock_audit_write = mocker.patch( - "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + "litellm.proxy.management_helpers.audit_logs.create_audit_logs_for_updates", new=mocker.AsyncMock(), ) results = await _delete(prisma, ["u1", "u2"], cache=cache) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f0baceecdaa..db0a74d360a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -8536,6 +8536,7 @@ PROXY_STAMPED_NAMES: Final = frozenset( "litellm_trusted_callback_vars", "_litellm_addressed_response_id", "_litellm_strip_stream_usage", + "_litellm_stream_usage_injected", "client_side_timeout", "model_file_id_mapping", ) diff --git a/tests/unit/proxy/proxy_server/test_routes_login_sso.py b/tests/unit/proxy/proxy_server/test_routes_login_sso.py index 75e14a89148..43629b926ad 100644 --- a/tests/unit/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/unit/proxy/proxy_server/test_routes_login_sso.py @@ -11,8 +11,12 @@ Routes covered: from __future__ import annotations from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock +import pytest +from fastapi.testclient import TestClient + from .conftest import normalize # --------------------------------------------------------------------------- @@ -398,6 +402,29 @@ def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatc } +def _sealed_return_to(return_to: str) -> str: + """The cookie value /sso/key/generate writes: return_to sealed under the master key, never plaintext.""" + from litellm.constants import CP_RETURN_TO_COOKIE_PREFIX + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token + + return encrypt_bearer_token(return_to, prefix=CP_RETURN_TO_COOKIE_PREFIX) + + +def test_login_form_ignores_a_plaintext_return_to_cookie(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None: + """A return_to the proxy did not seal (a client-planted or pre-upgrade plaintext cookie) is not + honored: /login lands on the dashboard and clears the cookie.""" + _install_login_mocks(monkeypatch) + response: Final = client.post( + "/login", + data={"username": "admin", "password": "password"}, + cookies={"litellm_cp_return_to": "/mcp/authorize?client_id=llm_dcrc_abc"}, + follow_redirects=False, + ) + assert response.status_code == 303 + assert "/ui" in response.headers.get("location", "") + assert 'litellm_cp_return_to=""' in response.headers.get("set-cookie", "") + + def test_login_form_honors_same_origin_return_to_cookie(client, monkeypatch): """The aggregate DCR connect flow preserves a same-origin return_to in the litellm_cp_return_to cookie; /login must RESUME there after password sign-in instead of dead-ending at the dashboard.""" @@ -406,7 +433,7 @@ def test_login_form_honors_same_origin_return_to_cookie(client, monkeypatch): response = client.post( "/login", data={"username": "admin", "password": "password"}, - cookies={"litellm_cp_return_to": return_to}, + cookies={"litellm_cp_return_to": _sealed_return_to(return_to)}, follow_redirects=False, ) assert response.status_code == 303 @@ -425,7 +452,7 @@ def test_login_form_honors_control_plane_return_to_cookie(client, monkeypatch): response = client.post( "/login", data={"username": "admin", "password": "password"}, - cookies={"litellm_cp_return_to": "https://cp.example.com/console"}, + cookies={"litellm_cp_return_to": _sealed_return_to("https://cp.example.com/console")}, follow_redirects=False, ) location = response.headers.get("location", "") @@ -447,7 +474,7 @@ def test_login_form_survives_stale_control_plane_return_to(client, monkeypatch): response = client.post( "/login", data={"username": "admin", "password": "password"}, - cookies={"litellm_cp_return_to": "https://old-cp.example.com/console"}, + cookies={"litellm_cp_return_to": _sealed_return_to("https://old-cp.example.com/console")}, follow_redirects=False, ) assert response.status_code == 303, "login must not break on a stale return_to cookie" @@ -463,7 +490,7 @@ def test_login_form_ignores_open_redirect_return_to(client, monkeypatch): response = client.post( "/login", data={"username": "admin", "password": "password"}, - cookies={"litellm_cp_return_to": "https://evil.example.com/steal"}, + cookies={"litellm_cp_return_to": _sealed_return_to("https://evil.example.com/steal")}, follow_redirects=False, ) assert response.status_code == 303 diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index 07a61c8587b..a86314e3063 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/tests/unit/proxy/proxy_server/test_routes_queue.py @@ -6,10 +6,13 @@ Pins (PR2): from __future__ import annotations +from collections.abc import Callable +from contextlib import AbstractContextManager from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi.testclient import TestClient from litellm import Router from litellm.proxy import proxy_server @@ -85,6 +88,25 @@ def test_queue_chat_completions_happy(client, auth_as, patched_queue): } +def test_queue_chat_completions_drops_client_sent_stream_usage_flags( + client: TestClient, auth_as: Callable[..., AbstractContextManager[None]], patched_queue: MagicMock +) -> None: + payload: Final = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "priority": 0, + "_litellm_strip_stream_usage": True, + "_litellm_stream_usage_injected": True, + } + with auth_as(): + response: Final = client.post("/queue/chat/completions", json=payload) + assert response.status_code == 200 + forwarded: Final = patched_queue.schedule_acompletion.call_args.kwargs + assert "_litellm_strip_stream_usage" not in forwarded + assert "_litellm_stream_usage_injected" not in forwarded + assert forwarded["messages"] == [{"role": "user", "content": "hi"}] + + def test_queue_chat_completions_no_router_error(client, auth_as, queue_no_router): """Pins ``POST /queue/chat/completions`` (error: no llm_router).""" payload = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} diff --git a/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py index 07886745222..ddb4b08ee9e 100644 --- a/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py +++ b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py @@ -10,7 +10,7 @@ from fastapi.testclient import TestClient import litellm from litellm.proxy._types import LiteLLMRoutes from litellm.proxy.proxy_server import app -from litellm.types.router import ModelGroupInfo +from litellm.types.router import DeploymentTypedDict, ModelGroupInfo from tests._master_key import MASTER_KEY client = TestClient(app) @@ -21,13 +21,17 @@ LEGACY_MODEL_HUB_PATH = "/public/model_hub" @dataclass(frozen=True, slots=True) class _FakeRouter: - """Stands in for the running Router: `_get_model_group_info` only ever asks it this.""" + """Stands in for the running Router: `_get_model_group_info` only ever asks it these.""" infos: Mapping[str, ModelGroupInfo] + model_group_alias: Mapping[str, str] = MappingProxyType({}) def get_model_group_info(self, model_group: str, team_id: str | None = None) -> ModelGroupInfo | None: return self.infos.get(model_group) + def get_model_list_of_routed_group(self, model_group: str) -> list[DeploymentTypedDict]: + return [] + def _info( name: str, @@ -61,13 +65,19 @@ def _groups(response) -> list[str]: return [row["model_group"] for row in response.json()["data"]] -def _health_check(model_name: str, status: str = "healthy"): +def _health_check( + model_name: str, + status: str = "healthy", + *, + model_id: str | None = None, + checked_at: datetime = datetime(2026, 8, 1, 9, 30, tzinfo=timezone.utc), +): check = MagicMock() check.model_name = model_name - check.model_id = None + check.model_id = model_id check.status = status check.response_time_ms = 12.5 - check.checked_at = datetime(2026, 8, 1, 9, 30, tzinfo=timezone.utc) + check.checked_at = checked_at return check @@ -154,6 +164,21 @@ def test_the_latest_health_check_lands_on_its_row(monkeypatch): assert rows["model-000"]["health_status"] is None +def test_a_stale_row_sorting_after_the_fresh_one_does_not_speak_for_the_group(monkeypatch): + """The read orders by model_id with NULL last; that order used to pick the group's health.""" + fresh = _health_check( + "model-001", model_id="deployment-a", checked_at=datetime(2026, 10, 16, 11, 10, tzinfo=timezone.utc) + ) + stale = _health_check("model-001", status="unhealthy", checked_at=datetime(2026, 9, 4, 11, 20, tzinfo=timezone.utc)) + prisma_client, _ = _recording_prisma([fresh, stale]) + _publish(monkeypatch, _named(3), prisma_client=prisma_client) + + rows = {row["model_group"]: row for row in _get().json()["data"]} + + assert rows["model-001"]["health_status"] == "healthy" + assert rows["model-001"]["health_checked_at"] == "2026-10-16T11:10:00+00:00" + + def test_a_health_read_that_returns_nothing_still_serves_the_page(monkeypatch): prisma_client, read = _recording_prisma() read.return_value = [] diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index d5bacd9fd3c..8f27365fcc1 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -15,6 +15,7 @@ from pydantic import TypeAdapter import litellm from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.health_check_latest import LatestHealthCheckRow from litellm.proxy.public_endpoints import router from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -852,6 +853,46 @@ def test_public_model_hub_mixed_health_statuses(): app.dependency_overrides.clear() +def _stored_health_check(model_id: str | None, status: str, checked_at: datetime) -> LatestHealthCheckRow: + return LatestHealthCheckRow( + health_check_id=f"hc-{model_id}", + model_name="gpt-4", + model_id=model_id, + status=status, + checked_at=checked_at, + created_at=checked_at, + updated_at=checked_at, + ) + + +def test_public_model_hub_reports_the_newest_check_of_a_group_not_the_last_row_read(): + """Rows come back in model_id order with NULL last, so a stale name-probe row used to overwrite the fresh one.""" + app = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: MagicMock() + client = TestClient(app) + fresh = _stored_health_check("deployment-a", "healthy", datetime(2026, 10, 16, 11, 10, tzinfo=timezone.utc)) + stale = _stored_health_check(None, "unhealthy", datetime(2026, 9, 4, 11, 20, tzinfo=timezone.utc)) + mock_prisma = MagicMock() + mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=(fresh, stale)) + + with ( + patch("litellm.public_model_groups", ["gpt-4"]), + patch("litellm.proxy.proxy_server.get_model_group_info") as mock_get_info, + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + mock_get_info.return_value = [ + ModelGroupInfoProxy(model_group="gpt-4", providers=["openai"], is_public_model_group=True) + ] + response = client.get("/public/model_hub", headers={"Authorization": "Bearer test-key"}) + + assert response.status_code == 200 + assert response.json()[0]["health_status"] == "healthy" + assert response.json()[0]["health_checked_at"] == "2026-10-16T11:10:00+00:00" + app.dependency_overrides.clear() + + # --------------------------------------------------------------------------- # /public/agent_hub # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 32b77c3b564..abc1fe88c7a 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1079,6 +1079,9 @@ def _managed_files_over(table: _UniqueManagedObjectTable) -> tuple[MagicMock, Ma prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_deprecatedverificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_deletedteamtable.find_first = AsyncMock(return_value=None) proxy_logging: Final = MagicMock() proxy_logging.get_proxy_hook = MagicMock( return_value=_PROXY_LiteLLMManagedFiles(internal_usage_cache=DualCache(), prisma_client=prisma_client) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index ed8ae2db8b1..23bbe521670 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3725,7 +3725,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"user_api_key_user_email": null, "actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "used_client_oauth_token": null, "litellm_roi_estimator": false, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"user_api_key_user_email": null, "actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "used_client_oauth_token": null, "litellm_roi_estimator": false, "billed_retry_attempt": false, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index ebef1fbbef7..facb4c71169 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -3551,6 +3551,30 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == "downstream-model" + def test_frozen_response_comes_back_as_a_copy_carrying_the_requested_model(self): + from litellm.types.extraction import ExtractionResponse + + response_obj = ExtractionResponse.model_validate( + { + "model": "cloudflare/@cf/cloudflare/clairic", + "answers": {"vendor": {"type": "text", "value": "Acme Tools"}}, + "truncated": False, + "usage": {"input_tokens": 52, "output_tokens": 0}, + } + ) + response_obj.set_hidden_params({"response_cost": 0.0000052}) + + stamped = _override_openai_response_model( + response_obj=response_obj, + requested_model="invoice-reader", + log_context="test_context", + ) + + assert isinstance(stamped, ExtractionResponse) + assert stamped.model == "invoice-reader" + assert stamped.answers == response_obj.answers + assert stamped.hidden_params["response_cost"] == 0.0000052 + class TestIsAzureModelRouterRequest: """Tests for _is_azure_model_router_request helper""" @@ -9157,6 +9181,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_usage": True} assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_default_preserves_other_client_stream_options_keys(self): data = {"stream": True, "stream_options": {"include_obfuscation": True}} @@ -9165,6 +9190,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_obfuscation": True, "include_usage": True} assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_client_requested_usage_is_left_untouched_and_not_stripped(self): data = {"stream": True, "stream_options": {"include_usage": True}} @@ -9181,6 +9207,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"]["include_usage"] is True assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_explicit_false_flag_disables_injection_entirely(self): data = {"stream": True} @@ -9206,6 +9233,36 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_usage": True} assert "_litellm_strip_stream_usage" not in data + @pytest.mark.parametrize("client_stream_options", [None, {"include_obfuscation": True}]) + def test_flag_true_marks_usage_injected_for_clients_that_did_not_ask( + self, client_stream_options: dict[str, bool] | None + ): + data: Final = { + "stream": True, + **({} if client_stream_options is None else {"stream_options": client_stream_options}), + } + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"]["include_usage"] is True + assert data["_litellm_stream_usage_injected"] is True + + def test_flag_true_leaves_usage_unmarked_when_client_asked_and_neutralizes_client_marker(self): + data: Final = {"stream": True, "stream_options": {"include_usage": True}, "_litellm_stream_usage_injected": True} + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"] == {"include_usage": True} + assert data["_litellm_stream_usage_injected"] is False + def test_flag_true_respects_client_explicit_include_usage_false(self): data = {"stream": True, "stream_options": {"include_usage": False}} diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 246496ec48f..3f8e3bc93d1 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -936,6 +936,26 @@ def test_body_snapshot_excludes_team_callback_credentials() -> None: }, proxy_request +def test_body_snapshot_keeps_the_stream_usage_injection_marker_out_of_the_stored_request() -> None: + from litellm.proxy.common_request_processing import stream_usage_tracking_updates + from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot + from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD + + request: Final = {"model": "gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True} + proxy_request: Final = {"body": dict(request)} + data: Final = { + **request, + **stream_usage_tracking_updates(request, {}, "acompletion", lambda: True), + "proxy_server_request": proxy_request, + } + + refresh_proxy_server_request_body_snapshot(data) + + assert data[STREAM_USAGE_INJECTED_FIELD] is True + assert proxy_request["body"]["stream_options"] == {"include_usage": True} + assert STREAM_USAGE_INJECTED_FIELD not in proxy_request["body"] + + @pytest.mark.asyncio @pytest.mark.parametrize("pre_call_ran", [False, True]) async def test_post_guardrail_snapshot_preserves_logging_only_masking_in_spend_logs( @@ -9347,7 +9367,9 @@ class TestResolveUserProviderCredentials: llm_router=self._router(), ) - assert data["secret_fields"]["user_provider_credentials"] == {"copilot-cred": "gho_secret"} + assert data["secret_fields"]["user_provider_credentials"] == { + "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} + } assert data["secret_fields"]["user_provider_credentials_user_id"] == "user-a" # the shared fallback credential is never looked up call_where = table.find_many.await_args.kwargs["where"] @@ -9804,7 +9826,9 @@ class TestResolveUserProviderCredentials: team_id=None, llm_router=router, ) - assert data["secret_fields"]["user_provider_credentials"] == {"copilot-cred": "gho_secret"} + assert data["secret_fields"]["user_provider_credentials"] == { + "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} + } call_where: Final = table.find_many.await_args.kwargs["where"] assert call_where == {"AND": ({"user_id": "user-a"}, {"credential_name": {"in": ["copilot-cred"]}})} @@ -10290,7 +10314,8 @@ async def test_secret_fields_user_provider_credentials_repr_hides_tokens(monkeyp assert isinstance(credentials, RedactedDict) assert "gho_secret" not in str(data) assert "gho_secret" not in repr(credentials) - assert credentials["copilot-cred"] == "gho_secret", "provider reads still get the token" + assert credentials["copilot-cred"]["access_token"] == "gho_secret", "provider reads still get the token" + assert credentials["copilot-cred"]["github_host"] == "github.com" @pytest.mark.asyncio @@ -10304,8 +10329,10 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request _per_user_env(monkeypatch) - load = AsyncMock(return_value={"copilot-cred": "gho_user_a_token"}) - monkeypatch.setattr(upc, "aget_user_provider_tokens", load) + load = AsyncMock( + return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}} + ) + monkeypatch.setattr(upc, "aget_user_provider_credentials", load) router = Router( model_list=[ { @@ -10344,7 +10371,8 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey called_user_ids = [call.kwargs["user_id"] for call in load.await_args_list] assert called_user_ids == ["user-a"], f"header-mapped user-b must never reach the token loader: {called_user_ids}" credentials = updated["secret_fields"]["user_provider_credentials"] - assert credentials["copilot-cred"] == "gho_user_a_token" + assert credentials["copilot-cred"]["access_token"] == "gho_user_a_token" + assert credentials["copilot-cred"]["github_host"] == "github.com" assert updated["secret_fields"]["user_provider_credentials_user_id"] == "user-a" @@ -10357,8 +10385,10 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request _per_user_env(monkeypatch) - load = AsyncMock(return_value={"copilot-cred": "gho_user_a_token"}) - monkeypatch.setattr(upc, "aget_user_provider_tokens", load) + load = AsyncMock( + return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}} + ) + monkeypatch.setattr(upc, "aget_user_provider_credentials", load) router = Router( model_list=[ { @@ -10394,7 +10424,8 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin called_user_ids = [call.kwargs["user_id"] for call in load.await_args_list] assert called_user_ids == ["user-a"] - assert updated["secret_fields"]["user_provider_credentials"]["copilot-cred"] == "gho_user_a_token" + assert updated["secret_fields"]["user_provider_credentials"]["copilot-cred"]["access_token"] == "gho_user_a_token" + assert updated["secret_fields"]["user_provider_credentials"]["copilot-cred"]["github_host"] == "github.com" @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_proxy_reject_logging.py b/tests/unit/proxy/test_proxy_reject_logging.py index 6e8545a58b3..b300d48808f 100644 --- a/tests/unit/proxy/test_proxy_reject_logging.py +++ b/tests/unit/proxy/test_proxy_reject_logging.py @@ -209,4 +209,4 @@ async def test_chat_completion_request_with_redaction(route, body, router, monke assert _test_logger.reaches_async_failure_event is True - assert _test_logger.reaches_sync_failure_event is True + assert _test_logger.reaches_sync_failure_event is False diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 2fc1b72c932..c77d7ce8512 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -1,7 +1,7 @@ import asyncio import json from datetime import datetime -from typing import Any, Dict, Final, List, Optional, Union +from typing import Any, Dict, Final, List, Optional, TypeAlias, Union from unittest.mock import Mock import pytest @@ -24,6 +24,7 @@ from litellm.proxy.litellm_pre_call_utils import ( get_dynamic_logging_metadata, add_litellm_data_to_request, ) +from litellm.router import Router from pydantic import ValidationError pytestmark = pytest.mark.xdist_group("proxy_heavy") @@ -1602,6 +1603,109 @@ def test_get_model_group_info(): assert len(model_list) == 1 +def _custom_metadata_router(*model_infos: dict[object, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": f"fake-{index}"}, + "model_info": model_info, + } + for index, model_info in enumerate(model_infos) + ] + ) + + +_CustomFieldsCase: TypeAlias = tuple[ + tuple[dict[object, object], ...], tuple[str, ...] | None, dict[str, str | int | float | bool] | None +] +_CUSTOM_FIELDS_CASES: Final[tuple[_CustomFieldsCase, ...]] = ( + ((), None, None), + (({},), None, None), + ( + ({"custom_tags": ["deprecated", "ptu"], "custom_metadata": {"retirement_date": "2027-03-01"}},), + ("deprecated", "ptu"), + {"retirement_date": "2027-03-01"}, + ), + ( + ({}, {"custom_tags": ["ptu"], "custom_metadata": {"provisioned_ptu": 100, "gsu": False}}), + ("ptu",), + {"provisioned_ptu": 100, "gsu": False}, + ), + ( + ( + {"custom_tags": ["ptu", "eu"], "custom_metadata": {"provisioned_ptu": 100, "region": "eu"}}, + {"custom_tags": ["eu", "deprecated"], "custom_metadata": {"provisioned_ptu": 50, "max_tpm": 1.5}}, + ), + ("ptu", "eu", "deprecated"), + {"provisioned_ptu": 100, "region": "eu", "max_tpm": 1.5}, + ), + ( + ({"custom_tags": ["ok", 7, None, {"x": 1}], "custom_metadata": {"owner": "ml", "nested": {"a": 1}}},), + ("ok",), + {"owner": "ml"}, + ), + (({"custom_tags": "ptu", "custom_metadata": ["not", "a", "map"]},), None, None), + ( + ( + { + "custom_metadata": { + 2027: "retirement", + "nan": float("nan"), + "inf": float("inf"), + "neg_inf": float("-inf"), + "max_tpm": 1.5, + } + }, + ), + None, + {"max_tpm": 1.5}, + ), +) + + +@pytest.mark.parametrize(("model_infos", "expected_tags", "expected_metadata"), _CUSTOM_FIELDS_CASES) +def test_get_model_group_info_merges_custom_tags_and_metadata_across_deployments( + model_infos: tuple[dict[object, object], ...], + expected_tags: tuple[str, ...] | None, + expected_metadata: dict[str, str | int | float | bool] | None, +) -> None: + from litellm.proxy.proxy_server import get_model_group_info + + router: Final = _custom_metadata_router({"description": "first"}, *model_infos) + + (group,) = get_model_group_info(llm_router=router, all_models_str=["gpt-4"], model_group="gpt-4") + + assert (group.custom_tags, group.custom_metadata) == (expected_tags, expected_metadata) + + +@pytest.mark.parametrize( + "alias_target", + [pytest.param("gpt-4", id="string"), pytest.param({"model": "gpt-4", "hidden": False}, id="item")], +) +def test_get_model_group_info_serves_the_target_groups_custom_fields_on_an_alias_row( + alias_target: str | dict[str, object], +) -> None: + from litellm.proxy.proxy_server import get_model_group_info + + router: Final = Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"custom_tags": ["deprecated"], "custom_metadata": {"retirement_date": "2027-03-01"}}, + } + ], + model_group_alias={"gpt-alias": alias_target}, + ) + + groups: Final = get_model_group_info(llm_router=router, all_models_str=["gpt-4", "gpt-alias"], model_group=None) + + assert {group.model_group: (group.custom_tags, group.custom_metadata) for group in groups} == { + model_group: (("deprecated",), {"retirement_date": "2027-03-01"}) for model_group in ("gpt-4", "gpt-alias") + } + + @pytest.fixture def mock_team_data(): return [ @@ -2652,6 +2756,58 @@ async def test_post_call_failure_hook_auth_error_llm_api_route(): mock_handle_logging.assert_called_once() +class _InlineThread: + def __init__(self, target, args=(), kwargs=None, daemon=None) -> None: + self._run = lambda: target(*args, **(kwargs or {})) + + def start(self) -> None: + self._run() + + +@pytest.mark.asyncio +async def test_proxy_only_failure_gives_a_logger_with_both_failure_hooks_only_the_async_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import types + + import litellm.proxy.utils as proxy_utils + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import Logging + + hooks: Final[list[str]] = [] + + class Recorder(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + hooks.append("sync_failure") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + hooks.append("async_failure") + + for name in ("failure_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(proxy_utils, "threading", types.SimpleNamespace(Thread=_InlineThread)) + recorder: Final = Recorder() + logging_obj: Final = Logging( + model="gpt-4o-mini-tts", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="aspeech", + start_time=datetime.now(), + litellm_call_id="proxy-only-failure", + function_id="proxy-only-failure", + dynamic_failure_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + await proxy_utils.ProxyLogging._dispatch_proxy_only_failure_handlers( + litellm_logging_obj=logging_obj, + original_exception=HTTPException(status_code=401, detail="Authentication Error: invalid user key"), + ) + + assert hooks == ["async_failure"] + + @pytest.mark.asyncio @pytest.mark.parametrize( "request_data, route, expected_call_type", diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 59e29c9b33c..7b18cb3350c 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -828,6 +828,125 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: assert expected_input_cost > 0 assert request_data["response_cost"] == pytest.approx(expected_input_cost) + @pytest.mark.asyncio + async def test_upstream_499_error_is_a_plain_failure_not_a_priced_disconnect(self): + """A provider answering HTTP 499 surfaces as litellm.APIError with status_code 499. Only the + proxy's own disconnect HTTPException means the caller hung up mid-generation, so the upstream + error keeps the estimated usage at zero cost like any other post-dispatch provider error.""" + from datetime import datetime + + from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED + from litellm.types.utils import Usage + + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 0), + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "count these input tokens please"}], + "call_type": "acompletion", + } + ), + "metadata": {}, + } + await self._run( + request_data, + original_exception=litellm.APIError( + status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, + message="client closed request", + llm_provider="openai", + model="gpt-3.5-turbo", + ), + ) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + assert estimated.prompt_tokens > 0 + assert request_data["response_cost"] == 0.0 + + @pytest.mark.asyncio + async def test_client_disconnect_after_the_call_already_completed_is_not_priced_again(self): + """The LLM task finished and claimed its success log (which carries the real usage and cost) + while a during_call guardrail task was still pending, then the caller hung up. The 499 raised + for that disconnect must not price the prompt a second time.""" + from datetime import datetime + + from fastapi import HTTPException + + from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "count these input tokens please"}], + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1, 12, 0, 0), + litellm_call_id="completed-before-disconnect", + function_id="fn", + ) + logging_obj.model_call_details.update( + { + "first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 1), + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "count these input tokens please"}], + "call_type": "acompletion", + } + ) + assert logging_obj.claim_async_success_log() + logging_obj.mark_async_success_log_enqueued() + request_data = {"litellm_logging_obj": logging_obj, "metadata": {}} + await self._run( + request_data, + original_exception=HTTPException( + status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request" + ), + ) + + assert request_data.get("response_cost", 0.0) == 0.0 + + @pytest.mark.asyncio + async def test_client_disconnect_with_a_claimed_but_deferred_success_log_is_still_priced(self): + """With a post_call guardrail registered the success log is claimed but parked in the deferred slot, + and a 499 raised before the flush means it never runs. Nothing else bills that request, so the + failure row keeps the estimated input cost.""" + from datetime import datetime + + from fastapi import HTTPException + + from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "count these input tokens please"}], + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1, 12, 0, 0), + litellm_call_id="deferred-success-never-flushed", + function_id="fn", + ) + logging_obj.model_call_details.update( + { + "first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 1), + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "count these input tokens please"}], + "call_type": "acompletion", + } + ) + logging_obj.defer_async_logging = True + assert logging_obj.claim_async_success_log() + logging_obj.enqueue_deferred_logging = lambda: None + request_data = {"litellm_logging_obj": logging_obj, "metadata": {}} + await self._run( + request_data, + original_exception=HTTPException( + status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request" + ), + ) + + assert request_data["response_cost"] > 0.0 + @pytest.mark.asyncio async def test_client_disconnect_before_provider_handoff_is_not_priced(self): """A caller that hung up while the request was still in pre-call hooks never reached the @@ -2160,6 +2279,35 @@ def test_a_timed_out_request_is_priced_at_the_rate_the_router_registered_for_the assert response_cost == pytest.approx(usage.prompt_tokens * 0.25) +@pytest.mark.parametrize(("prompt_tokens", "expected_rate"), [(1_000, 1e-7), (250_000, 5e-7)]) +def test_a_timed_out_request_is_priced_at_the_deployments_long_context_tier(monkeypatch, prompt_tokens, expected_rate): + """A deployment that declares ``input_cost_per_token_above_200k_tokens`` next to its standard rate + bills a 250k-token prompt at the tier on the success path; the failure row must not flatten it to + the standard rate.""" + from litellm.proxy.utils import _estimated_input_cost + + monkeypatch.setitem( + litellm.model_cost, + "dep-lit-f9-tiered", + { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-7, + "input_cost_per_token_above_200k_tokens": 5e-7, + "output_cost_per_token": 4e-7, + }, + ) + + response_cost = _estimated_input_cost( + model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"}, + prompt_tokens=prompt_tokens, + deployment_id="dep-lit-f9-tiered", + base_model=None, + ) + + assert response_cost == pytest.approx(prompt_tokens * expected_rate) + + def test_every_earlier_timed_out_attempt_is_billed_at_its_own_deployment_rate(monkeypatch): """The router leaves one RetryAttemptRecord per failed attempt in metadata.previous_models, the final attempt included; the earlier Timeout ones are billed (each at the deployment that served @@ -2199,6 +2347,74 @@ def test_every_earlier_timed_out_attempt_is_billed_at_its_own_deployment_rate(mo assert all(a.usage is usage for a in attempts) +def test_an_earlier_attempt_on_another_model_is_priced_at_that_models_rate_not_the_final_models(monkeypatch): + """After a fallback hop the final attempt runs on another model group, so the records of the group + that timed out first name their own deployment model (and its ``base_model`` when the deployment + name is not in the cost map). Each is billed at that model's input rate, not at the final model's.""" + from litellm.proxy.utils import _billed_earlier_attempts + + for pricing_model, provider, input_rate in ( + ("lit145-final-model", "openai", 0.5), + ("lit145-first-model", "openai", 0.02), + ("azure/lit145-first-model", "azure", 0.02), + ): + monkeypatch.setitem( + litellm.model_cost, + pricing_model, + {"litellm_provider": provider, "mode": "chat", "input_cost_per_token": input_rate, "output_cost_per_token": 1.0}, + ) + usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + + attempts = _billed_earlier_attempts( + model_call_details={"model": "openai/lit145-final-model", "custom_llm_provider": "openai"}, + litellm_params={ + "metadata": { + "attempted_retries": 1, + "model_info": {"id": "dep-lit145-final"}, + "previous_models": ( + { + "deployment_id": "dep-lit145-first", + "model": "openai/lit145-first-model", + "base_model": None, + "exception_type": "Timeout", + "attempted_retries": 0, + }, + { + "deployment_id": "dep-lit145-first-azure", + "model": "azure/lit145-first-deployment", + "base_model": "azure/lit145-first-model", + "exception_type": "Timeout", + "attempted_retries": 1, + }, + { + "deployment_id": "dep-lit145-final", + "model": "openai/lit145-final-model", + "base_model": None, + "exception_type": "Timeout", + "attempted_retries": 0, + }, + { + "deployment_id": "dep-lit145-final", + "model": "openai/lit145-final-model", + "base_model": None, + "exception_type": "Timeout", + "attempted_retries": 1, + }, + ), + } + }, + estimated_usage=usage, + base_model=None, + deployment_id="dep-lit145-final", + ) + + assert [(a.attempted_retries, a.deployment_id, a.response_cost) for a in attempts] == [ + (0, "dep-lit145-first", pytest.approx(24 * 0.02)), + (1, "dep-lit145-first-azure", pytest.approx(24 * 0.02)), + (0, "dep-lit145-final", pytest.approx(24 * 0.5)), + ] + + def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_streams_recovered_usage(): """A final attempt whose stream broke mid-flight lifts its recovered usage (prompt plus the chunks delivered); the earlier timed-out attempts never produced output, so they are priced from the diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 8928b79bcd4..6044475652c 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,19 +1,42 @@ +import asyncio import contextlib +import gc +import threading +import weakref import httpx import json import litellm import pytest +import redis import respx +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping +from datetime import datetime from types import SimpleNamespace +from uuid import uuid4 from typing import Final from unittest.mock import MagicMock +from fakeredis import TcpFakeServer from fastapi import HTTPException -from litellm.proxy.route_llm_request import JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, route_request -from litellm.types.router import RouterGeneralSettings +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.route_llm_request import ( + JSON_OBJECT_ADAPTER, + ProxyModelNotFoundError, + UserConfigRouters, + route_request, + user_config_routers, +) +from litellm.proxy.utils import ProxyLogging +from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ADAPTIVE_ROUTER_RESPONSE_HEADER, +) +from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError @pytest.mark.parametrize( @@ -2434,3 +2457,571 @@ async def test_team_key_search_of_registered_bedrock_store_reaches_bedrock_with_ posted_url ) assert "Credential=AKIAREGISTRYKEY00000/" in signed_auth, signed_auth + + +@pytest.fixture(autouse=True) +def discard_cached_user_config_routers() -> Iterator[None]: + yield + user_config_routers.clear() + + +@pytest.fixture +def fake_redis_port() -> Iterator[int]: + server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis") + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield server.server_address[1] + server.shutdown() + server.server_close() + + +async def _wait_out_a_rate_limit_minute_rollover() -> None: + seconds_left: Final = 60 - datetime.now().second + if seconds_left < 10: + await asyncio.sleep(seconds_left) + + +def _proxy_logging_obj(model: str = "rate-limited") -> LiteLLMLoggingObj: + logging_obj, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model=model, + messages=[{"role": "user", "content": "hi"}], + litellm_call_id=str(uuid4()), + ) + assert isinstance(logging_obj, LiteLLMLoggingObj) + return logging_obj + + +def _mock_deployment(model_name: str, **deployment_limits: int) -> dict[str, object]: + return { + "model_name": model_name, + **deployment_limits, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "mock_response": "ok"}, + } + + +async def _user_config_request_outcome( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + logging_obj: LiteLLMLoggingObj | None, +) -> str: + data: Final = { + "model": "rate-limited", + "messages": [{"role": "user", "content": "hi"}], + "num_retries": 0, + **({"litellm_logging_obj": logging_obj} if logging_obj is not None else {}), + "user_config": { + **routing_settings, + **({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}), + "model_list": [_mock_deployment("rate-limited", **deployment_limits)], + }, + } + try: + await (await route_request(data, None, None, "acompletion")) + except (litellm.RateLimitError, RouterRateLimitError): + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + +async def _tpm_for_three_requests() -> int: + messages: Final = [{"role": "user", "content": "hi"}] + probe: Final = await litellm.acompletion(model="openai/gpt-5.5", messages=messages, mock_response="ok") + assert isinstance(probe, litellm.ModelResponse) + probe_usage: Final = getattr(probe, "usage", None) + assert isinstance(probe_usage, litellm.Usage) + return 2 * probe_usage.total_tokens + litellm.token_counter(messages=messages) + + +async def _request_cost() -> float: + probe: Final = await litellm.acompletion( + model="openai/gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok" + ) + return litellm.completion_cost(completion_response=probe) + + +async def _proxy_router_outcomes( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + requests: int, +) -> list[str]: + router: Final = litellm.Router( + **routing_settings, + num_retries=0, + model_list=[_mock_deployment("proxy-rate-limited", **deployment_limits)], + **({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}), + ) + + async def outcome() -> str: + try: + await router.acompletion(model="proxy-rate-limited", messages=[{"role": "user", "content": "hi"}]) + except (litellm.RateLimitError, RouterRateLimitError): + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + try: + return [await outcome() for _ in range(requests)] + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() -> None: + data: Final = { + "model": "smart-cheap-router", + "messages": [{"role": "user", "content": "Write a one line python function that adds two numbers"}], + "metadata": {}, + "user_config": { + "model_list": [ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast", "smart"]}, + }, + }, + { + **_mock_deployment("fast"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 2, "strengths": []}}, + }, + { + **_mock_deployment("smart"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}}, + }, + ] + }, + } + + await (await route_request(data, None, None, "acompletion")) + headers: Final = await ProxyLogging(user_api_key_cache=DualCache()).post_call_response_headers_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), response=None + ) + + assert headers == {ADAPTIVE_ROUTER_RESPONSE_HEADER: data["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY]} + + +def _global_router_callbacks() -> tuple[list[object], list[object]]: + return list(litellm.callbacks), list(litellm.input_callback) + + +@pytest.mark.parametrize( + "routing_settings, deployment_limits", + [ + ({"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}), + ({"routing_strategy": "usage-based-routing-v2"}, {"tpm": 30}), + ({"routing_strategy": "usage-based-routing"}, {"rpm": 2}), + ({"optional_pre_call_checks": ["enforce_model_rate_limits"]}, {"rpm": 2}), + ], +) +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.parametrize("with_redis", [True, False]) +@pytest.mark.asyncio +async def test_user_config_router_enforces_its_limit_across_proxy_requests( + fake_redis_port: int, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + arrives_with_logging_obj: bool, + with_redis: bool, +) -> None: + await _wait_out_a_rate_limit_minute_rollover() + callbacks_before: Final = _global_router_callbacks() + outcomes: Final = [ + await _user_config_request_outcome( + fake_redis_port if with_redis else None, + routing_settings, + deployment_limits, + _proxy_logging_obj() if arrives_with_logging_obj else None, + ) + for _ in range(6) + ] + + assert "429" in outcomes, outcomes + assert _global_router_callbacks() == callbacks_before + assert outcomes == await _proxy_router_outcomes( + fake_redis_port if with_redis else None, routing_settings, deployment_limits, len(outcomes) + ) + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_limit( + fake_redis_port: int, arrives_with_logging_obj: bool +) -> None: + tpm_for_three_requests: Final = await _tpm_for_three_requests() + + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, + {"routing_strategy": "usage-based-routing-v2"}, + {"tpm": tpm_for_three_requests}, + _proxy_logging_obj() if arrives_with_logging_obj else None, + ) + + await _wait_out_a_rate_limit_minute_rollover() + concurrent_outcomes: Final = await asyncio.gather(request(), request()) + later_outcomes: Final = [await request() for _ in range(2)] + + assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"] + + +@pytest.mark.asyncio +async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None: + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}, _proxy_logging_obj() + ) + + await _wait_out_a_rate_limit_minute_rollover() + outcomes: Final = await asyncio.gather(*(request() for _ in range(4))) + + assert outcomes.count("200") <= 2, outcomes + assert "429" in outcomes, outcomes + + +@pytest.mark.parametrize( + "routing_strategy, limited_by", + [("usage-based-routing-v2", "rpm"), ("usage-based-routing-v2", "tpm"), ("usage-based-routing", "rpm")], +) +@pytest.mark.parametrize("proxy_router_shares_redis", [True, False]) +@pytest.mark.asyncio +async def test_a_proxy_router_on_the_same_strategy_leaves_the_user_config_limit_unchanged( + fake_redis_port: int, routing_strategy: str, limited_by: str, proxy_router_shares_redis: bool +) -> None: + deployment_limits: Final = {"rpm": 3} if limited_by == "rpm" else {"tpm": await _tpm_for_three_requests()} + + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, {"routing_strategy": routing_strategy}, deployment_limits, _proxy_logging_obj() + ) + + async def outcomes() -> list[str]: + first_instance: Final = [await request() for _ in range(2)] + user_config_routers.clear() + second_instance: Final = [await request() for _ in range(3)] + return [*first_instance, *second_instance] + + await _wait_out_a_rate_limit_minute_rollover() + alone: Final = await outcomes() + user_config_routers.clear() + await asyncio.sleep(0.3) + redis.Redis(host="127.0.0.1", port=fake_redis_port).flushall() + proxy_router: Final = litellm.Router( + routing_strategy=routing_strategy, + model_list=[_mock_deployment("proxy-model")], + redis_host="127.0.0.1" if proxy_router_shares_redis else None, + redis_port=fake_redis_port if proxy_router_shares_redis else None, + ) + try: + beside_proxy_router: Final = await outcomes() + finally: + proxy_router.discard() + + assert "429" in alone, alone + assert beside_proxy_router == alone + + +@pytest.mark.parametrize("user_config_settings", [{}, {"optional_pre_call_checks": ["router_budget_limiting"]}]) +@pytest.mark.asyncio +async def test_user_config_requests_stay_under_the_proxy_provider_budget( + user_config_settings: Mapping[str, object], +) -> None: + proxy_router: Final = litellm.Router( + model_list=[_mock_deployment("proxy-model")], + provider_budget_config={"openai": {"budget_limit": await _request_cost() / 2, "time_period": "1d"}}, + ) + try: + await proxy_router.acompletion(model="proxy-model", messages=[{"role": "user", "content": "hi"}]) + await asyncio.sleep(0.3) + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await _user_config_request_outcome(None, user_config_settings, {}, _proxy_logging_obj()) + finally: + proxy_router.discard() + + +@pytest.mark.asyncio +async def test_a_user_config_provider_budget_limits_its_own_requests() -> None: + budget_limit: Final = await _request_cost() / 2 + + def request() -> Coroutine[None, None, str]: + budget: Final = {"openai": {"budget_limit": budget_limit, "time_period": "1d"}} + return _user_config_request_outcome(None, {"provider_budget_config": budget}, {}, _proxy_logging_obj()) + + first: Final = await request() + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await request() + assert first == "200" + + +@pytest.mark.asyncio +async def test_a_request_in_flight_when_its_user_config_router_is_dropped_still_counts_against_its_budget( + fake_redis_port: int, +) -> None: + budget_limit: Final = 1.5 * await _request_cost() + + async def request() -> None: + await ( + await route_request( + { + "model": "budgeted", + "messages": [{"role": "user", "content": "hi"}], + "litellm_logging_obj": _proxy_logging_obj("budgeted"), + "user_config": { + "redis_host": "127.0.0.1", + "redis_port": fake_redis_port, + "provider_budget_config": {"openai": {"budget_limit": budget_limit, "time_period": "1d"}}, + "model_list": [ + { + "model_name": "budgeted", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "mock_response": "ok", + "mock_delay": 0.5, + }, + } + ], + }, + }, + None, + None, + "acompletion", + ) + ) + + await request() + in_flight: Final = asyncio.create_task(request()) + await asyncio.sleep(0.2) + user_config_routers.clear() + await in_flight + await asyncio.sleep(0.3) + + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await request() + + +async def _user_config_request(model: str, user_config: Mapping[str, object], arrives_with_logging_obj: bool) -> object: + return await ( + await route_request( + { + "model": model, + "messages": list(_HELLO_MESSAGES), + "num_retries": 0, + **({"litellm_logging_obj": _proxy_logging_obj(model)} if arrives_with_logging_obj else {}), + "user_config": user_config, + }, + None, + None, + "acompletion", + ) + ) + + +def _ordered_pair(model_name: str, first_base: str, second_base: str) -> list[dict[str, object]]: + return [ + { + "model_name": model_name, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "api_base": api_base, "order": order}, + } + for order, api_base in ((1, first_base), (2, second_base)) + ] + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_a_rate_limited_user_config_deployment_cools_down_for_the_next_requests( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + rate_limited_base: Final = "https://rate-limited.example/v1" + healthy_base: Final = "https://healthy.example/v1" + user_config: Final = {"model_list": _ordered_pair("flaky", rate_limited_base, healthy_base)} + + with respx.mock(assert_all_called=False) as respx_mock: + rate_limited: Final = respx_mock.post(f"{rate_limited_base}/chat/completions").mock( + return_value=httpx.Response(429, json={"error": {"message": "slow down"}}) + ) + respx_mock.post(f"{healthy_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5")) + for _ in range(3): + await _user_config_request("flaky", user_config, arrives_with_logging_obj) + await asyncio.sleep(0.3) + + assert rate_limited.call_count == 1 + + +async def _remaining_usage_headers_after_429_then_200s( + send: Callable[[], Awaitable[object]], api_base: str +) -> list[tuple[object, object] | str]: + async def remaining_usage() -> tuple[object, object] | str: + try: + response: Final = await send() + except litellm.RateLimitError: + return "429" + finally: + await asyncio.sleep(0.3) + assert isinstance(response, litellm.ModelResponse) + headers: Final = response._hidden_params["additional_headers"] + return headers.get("x-ratelimit-remaining-requests"), headers.get("x-ratelimit-remaining-tokens") + + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.post(f"{api_base}/chat/completions").mock( + side_effect=[ + httpx.Response(429, json={"error": {"message": "slow down"}}), + _chat_completion_response("gpt-5.5"), + _chat_completion_response("gpt-5.5"), + ] + ) + return [await remaining_usage() for _ in range(3)] + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_user_config_remaining_usage_headers_match_a_proxy_router_with_the_same_deployment( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await _wait_out_a_rate_limit_minute_rollover() + proxy_base: Final = "https://proxy-limited.example/v1" + user_config_base: Final = "https://user-config-limited.example/v1" + + def deployment(model_name: str, api_base: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "api_base": api_base, + "rpm": 10, + "tpm": 1000, + }, + } + + proxy_router: Final = litellm.Router(model_list=[deployment("proxy-limited", proxy_base)]) + try: + proxy_router_headers: Final = await _remaining_usage_headers_after_429_then_200s( + lambda: proxy_router.acompletion(model="proxy-limited", messages=list(_HELLO_MESSAGES), num_retries=0), + proxy_base, + ) + finally: + proxy_router.discard() + + user_config: Final = {"model_list": [deployment("user-config-limited", user_config_base)]} + user_config_headers: Final = await _remaining_usage_headers_after_429_then_200s( + lambda: _user_config_request("user-config-limited", user_config, arrives_with_logging_obj), user_config_base + ) + + assert user_config_headers == proxy_router_headers + + +async def _first_deployment_calls_over_successes_then_failures( + send: Callable[[], Awaitable[object]], first_base: str, second_base: str +) -> int: + successes: Final = [_chat_completion_response("gpt-5.5") for _ in range(6)] + failures: Final = [httpx.Response(500, json={"error": {"message": "boom"}}) for _ in range(6)] + with respx.mock(assert_all_called=False) as respx_mock: + first: Final = respx_mock.post(f"{first_base}/chat/completions").mock(side_effect=[*successes, *failures]) + respx_mock.post(f"{second_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5")) + for _ in range(len(successes) + len(failures)): + await send() + await asyncio.sleep(0.3) + return first.call_count + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_user_config_successes_keep_a_deployment_out_of_cooldown_like_a_proxy_router( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await _wait_out_a_rate_limit_minute_rollover() + proxy_router: Final = litellm.Router( + model_list=_ordered_pair("proxy-pair", "https://proxy-first.example/v1", "https://proxy-second.example/v1") + ) + try: + proxy_router_calls: Final = await _first_deployment_calls_over_successes_then_failures( + lambda: proxy_router.acompletion(model="proxy-pair", messages=list(_HELLO_MESSAGES), num_retries=0), + "https://proxy-first.example/v1", + "https://proxy-second.example/v1", + ) + finally: + proxy_router.discard() + + user_config: Final = { + "model_list": _ordered_pair("user-pair", "https://user-first.example/v1", "https://user-second.example/v1") + } + + user_config_calls: Final = await _first_deployment_calls_over_successes_then_failures( + lambda: _user_config_request("user-pair", user_config, arrives_with_logging_obj), + "https://user-first.example/v1", + "https://user-second.example/v1", + ) + + assert user_config_calls == proxy_router_calls + + +@pytest.mark.asyncio +async def test_user_config_adaptive_router_learns_from_proxy_requests() -> None: + user_config: Final = { + "model_list": [ + _mock_deployment("smart"), + _mock_deployment("fast"), + { + "model_name": "adaptive", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["smart", "fast"]}, + }, + }, + ] + } + await ( + await route_request( + { + "model": "adaptive", + "messages": [{"role": "user", "content": "write a sorting function"}], + "metadata": {"session_id": "session-1"}, + "litellm_logging_obj": _proxy_logging_obj("adaptive"), + "user_config": user_config, + }, + None, + None, + "acompletion", + ) + ) + await asyncio.sleep(0.3) + + router: Final = user_config_routers.get(user_config, lambda: pytest.fail("the user_config router was not reused")) + snapshot: Final = await router.adaptive_routers["adaptive"][0].strategy.get_state_snapshot() + assert snapshot["feedback_contexts_live"] == 1, snapshot + + +@pytest.mark.asyncio +async def test_user_config_routers_discard_the_least_recently_used_router_past_capacity() -> None: + routers: Final = UserConfigRouters(max_size=2) + + def config(model_name: str) -> dict[str, object]: + return {"routing_strategy": "usage-based-routing-v2", "model_list": [_mock_deployment(model_name)]} + + def build(model_name: str) -> litellm.Router: + return litellm.Router(**config(model_name)) + + def must_be_cached() -> litellm.Router: + pytest.fail("a cached router was rebuilt") + + first_handler: Final = weakref.ref(routers.get(config("first"), lambda: build("first")).lowesttpm_logger_v2) + second_handler: Final = weakref.ref(routers.get(config("second"), lambda: build("second")).lowesttpm_logger_v2) + routers.get(config("first"), must_be_cached) + routers.get(config("third"), lambda: build("third")) + await asyncio.sleep(0.2) + gc.collect() + + assert routers.get(config("first"), must_be_cached) is not None + assert first_handler() is not None + assert second_handler() is None + routers.clear() 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 b9cf201a47a..395a8984875 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1,7 +1,9 @@ import asyncio +import gc import json import logging import sys +import warnings from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager, nullcontext from datetime import datetime @@ -18,9 +20,11 @@ from mcp.types import Tool as MCPTool # Import required modules import litellm +from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.langfuse.langfuse import langfuse_trace_id, langfuse_trace_metadata from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.xai.common_utils import xai_reported_cost_in_usd from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk from litellm.proxy.utils import _failure_fields_to_lift @@ -31,7 +35,7 @@ from litellm.types.llms.openai import ( OpenAIMcpServerTool, ToolParam, ) -from litellm.types.utils import StandardLoggingPayload, Usage +from litellm.types.utils import ModelResponse, StandardLoggingPayload, Usage class MockUserAPIKeyAuth: @@ -1065,6 +1069,108 @@ async def test_mcp_tool_execution_events_creation(): print("MCP tool execution events creation test passed!") +@pytest.mark.asyncio +async def test_aresponses_mcp_call_does_not_forward_admission_authorization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import importlib + import sys + import types + from unittest.mock import AsyncMock + + from mcp.types import CallToolResult, TextContent, Tool as MCPTool + + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + responses_main: Final = importlib.import_module("litellm.responses.main") + manager: Final = mcp_server_manager.global_mcp_server_manager + server: Final = MCPServer( + server_id="responses-auth-strip", + name="echo_none", + server_name="echo_none", + url="http://mcp.example", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization", "x-tenant"], + allow_all_keys=True, + ) + monkeypatch.setattr(manager, "registry", {server.server_id: server}) + + mcp_client: Final = AsyncMock() + mcp_client.list_tools.return_value = [MCPTool(name="whoami", inputSchema={"type": "object", "properties": {}})] + mcp_client.call_tool.return_value = CallToolResult( + content=[TextContent(type="text", text="ok")], + isError=False, + ) + monkeypatch.setattr(manager, "create_mcp_client", AsyncMock(return_value=mcp_client)) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + types.SimpleNamespace( + general_settings={}, + master_key="test-master-key", + proxy_logging_obj=None, + prisma_client=None, + user_api_key_cache=None, + ), + ) + + raw_headers: Final = {"authorization": "Bearer sk-" + "1234", "x-tenant": "tenant-a"} + user_api_key_auth: Final = UserAPIKeyAuth(api_key="hashed-key", token="hashed-key") + tools: Final = [ + { + "type": "mcp", + "server_label": "echo_none", + "server_url": "litellm_proxy/mcp/echo_none", + "require_approval": "never", + } + ] + + monkeypatch.setattr( + responses_main, + "aresponses", + AsyncMock( + side_effect=[ + ResponsesAPIResponse( + id="resp_first", + created_at=0, + output=[ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call-1", + "name": "echo_none-whoami", + "arguments": "{}", + "status": "completed", + } + ], + ), + ResponsesAPIResponse(id="resp_final", created_at=0, output=[]), + ] + ), + ) + monkeypatch.setattr( + importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler"), + "aresponses", + AsyncMock(return_value=ResponsesAPIResponse(id="resp_final", created_at=0, output=[])), + ) + + await responses_main.aresponses_api_with_mcp( + input="Call whoami", + model="openai/gpt-5", + tools=tools, + litellm_metadata={"user_api_key_auth": user_api_key_auth}, + secret_fields={"raw_headers": raw_headers}, + ) + + assert mcp_client.call_tool.await_args is not None + assert manager.create_mcp_client.await_args is not None + assert manager.create_mcp_client.await_args.kwargs["extra_headers"] == {"x-tenant": "tenant-a"} + + @pytest.mark.asyncio async def test_no_duplicate_mcp_tools_in_streaming_e2e(): """ @@ -1416,25 +1522,39 @@ def _openai_answer_round() -> bytes: ) -def _anthropic_round(message_id: str, block: Mapping[str, object], delta: Mapping[str, object], stop: str) -> bytes: +def _anthropic_message_start( + message_id: str, model: str, input_tokens: int, cache_usage: Mapping[str, object] = MappingProxyType({}) +) -> dict[str, object]: + return { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": input_tokens, "output_tokens": 1, **cache_usage}, + }, + } + + +def _anthropic_round( + message_id: str, + block: Mapping[str, object], + delta: Mapping[str, object], + stop: str, + model: str = "claude-haiku-4-5", + usage: tuple[int, int] = (5, 5), + cache_usage: Mapping[str, object] = MappingProxyType({}), +) -> bytes: return _sse( ( - { - "type": "message_start", - "message": { - "id": message_id, - "type": "message", - "role": "assistant", - "model": "claude-haiku-4-5", - "content": [], - "stop_reason": None, - "usage": {"input_tokens": 5, "output_tokens": 1}, - }, - }, + _anthropic_message_start(message_id, model, usage[0], cache_usage), {"type": "content_block_start", "index": 0, "content_block": dict(block)}, {"type": "content_block_delta", "index": 0, "delta": dict(delta)}, {"type": "content_block_stop", "index": 0}, - {"type": "message_delta", "delta": {"stop_reason": stop}, "usage": {"output_tokens": 5}}, + {"type": "message_delta", "delta": {"stop_reason": stop}, "usage": {"output_tokens": usage[1]}}, {"type": "message_stop"}, ) ) @@ -1484,6 +1604,7 @@ class _RoundRecorder(CustomLogger): def __init__(self) -> None: super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () + self.results: tuple[object, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () self.hooked: tuple[object, ...] = () @@ -1491,6 +1612,7 @@ class _RoundRecorder(CustomLogger): self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + self.results = (*self.results, response_obj) async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime @@ -1522,12 +1644,15 @@ _ECHO_MCP_TOOL: Final = MappingProxyType( def _serve_responses( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, responses: Sequence[httpx.Response] + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + responses: Sequence[httpx.Response], + url_regex: str = r"https://mcp-rounds\.example\.com/.*responses", ) -> respx.Route: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() - return respx_mock.post(url__regex=r"https://mcp-rounds\.example\.com/.*responses").mock(side_effect=list(responses)) + return respx_mock.post(url__regex=url_regex).mock(side_effect=list(responses)) def _serve_rounds( @@ -2322,3 +2447,953 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( assert lifecycle_ids[0] == lifecycle_ids[1] assert lifecycle_ids[2] != lifecycle_ids[0] assert not [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] + + +_FIXED_MARGIN: Final = 0.001 + + +def _streamed(body: Mapping[str, object], **terminal: object) -> httpx.Response: + completed: Final = json.loads(json.dumps({**body, **terminal}, default=dict)) + opened: Final = {**completed, "status": "in_progress", "output": [], "usage": None} + items: Final = tuple(enumerate(completed["output"])) + events: Final = ( + {"type": "response.created", "response": opened}, + {"type": "response.in_progress", "response": opened}, + *({"type": "response.output_item.added", "output_index": index, "item": item} for index, item in items), + *({"type": "response.output_item.done", "output_index": index, "item": item} for index, item in items), + {"type": f"response.{completed['status']}", "response": completed}, + ) + return httpx.Response( + 200, + content=_sse(tuple({**event, "sequence_number": number} for number, event in enumerate(events))), + headers=SSE_HEADERS, + ) + + +class _CutAfterTheFirstEvent(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},)) + raise httpx.ReadError("connection reset mid-stream") + + +async def _open_stream_auto_execute(logging_obj: Logging) -> object: + return await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=True, + litellm_logging_obj=logging_obj, + ) + + +async def _drain(stream: object) -> tuple[ResponsesAPIStreamingResponse, ...]: + return tuple([event async for event in cast(Any, stream)]) + + +async def _stream_auto_execute(logging_obj: Logging) -> tuple[ResponsesAPIStreamingResponse, ...]: + with _echo_tool_served("PING"): + return await _drain(await _open_stream_auto_execute(logging_obj)) + + +def _final_response( + events: Sequence[ResponsesAPIStreamingResponse], terminal: str = "response.completed" +) -> ResponsesAPIResponse: + assert _event_type(events[-1]) == terminal + response: Final = getattr(events[-1], "response", None) + assert isinstance(response, ResponsesAPIResponse) + return response + + +def _assert_one_stream_row_bills_both_rounds_once( + logged: tuple[StandardLoggingPayload, ...], client: ResponsesAPIResponse +) -> None: + charged_once: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN + assert len(logged) == 1 + row: Final = logged[0] + assert _output_types(row["response"]["output"]) == _output_types(client.output) + assert "message" in _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["cache_hit"] is None + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["hidden_params"]["additional_headers"]["llm_provider-content-type"] == SSE_HEADERS["content-type"] + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (71 + 149, 19 + 5) + assert client.usage.cost == pytest.approx(charged_once) + + +def _assert_breakdown_agrees_with_the_billed_cost(row: StandardLoggingPayload) -> None: + breakdown: Final = row.get("cost_breakdown") + assert breakdown is None or breakdown["total_cost"] == pytest.approx(row["response_cost"]) + + +def _request_failure_row(logging_obj: Logging, success_row: StandardLoggingPayload) -> Mapping[str, object]: + failure_row: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) + failure_payload: Final = cast(StandardLoggingPayload, failure_row["standard_logging_object"]) + assert failure_payload["status"] == "failure" + assert failure_payload["id"] != success_row["id"] + return failure_row + + +def _assert_billed_tokens(row: Mapping[str, object], tokens: tuple[int, int]) -> None: + usage: Final = row["combined_usage_object"] + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens) == tokens + + +def _assert_one_row_bills(logged: tuple[StandardLoggingPayload, ...], body: Mapping[str, object]) -> None: + usage: Final = cast(Mapping[str, int], body["usage"]) + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == ( + usage["input_tokens"], + usage["output_tokens"], + ) + assert logged[0]["response_cost"] == pytest.approx(_round_cost(body)) + assert logged[0]["response_cost"] > 0 + + +def _assert_one_row_bills_round_one_twice_with_one_fixed_margin(logged: tuple[StandardLoggingPayload, ...]) -> None: + charged_once: Final = 2 * _round_cost(_ROUND_ONE_BODY) - _FIXED_MARGIN + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (2 * 71, 2 * 19) + assert logged[0]["response_cost"] == pytest.approx(charged_once) + assert logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_logs_the_final_response_billed_for_both_rounds_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-rounds")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_when_the_proxy_releases( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + 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), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-rounds-deferred") + + with _echo_tool_served("PING"): + stream: Final = await _open_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + await asyncio.sleep(0.2) + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + assert logged_before_release == () + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_paid_first_round_on_its_success_row_when_the_follow_up_stream_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), httpx.Response(200, stream=_CutAfterTheFirstEvent(), headers=SSE_HEADERS)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-breaks") + + with pytest.raises(litellm.ServiceUnavailableError): + 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 [row["response_cost"] for row in recorder.failed] == [0.0] + 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 +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + 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), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-rounds-blocked") + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served("PING"): + stream: Final = await _open_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + if not ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked): + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_paid_first_round_once_when_the_follow_up_is_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + context_window_error: Final = { + "error": { + "message": "Your input exceeds the context window of this model.", + "type": "invalid_request_error", + "param": "input", + "code": "context_length_exceeded", + } + } + _serve_responses( + respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), httpx.Response(400, json=context_window_error)) + ) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-follow-up-rejected")) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert _event_type(events[-1]) == "error" + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_every_paid_round_with_one_fixed_margin_when_the_third_round_is_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + ( + _streamed(_ROUND_ONE_BODY), + _streamed(_ROUND_ONE_BODY), + httpx.Response(400, json={"error": {"message": "Your input exceeds the context window of this model."}}), + ), + ) + recorder: Final = _RoundRecorder() + + await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-third-round-rejected")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one_twice_with_one_fixed_margin(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_every_finished_round_when_the_client_disconnects_during_a_later_tool_call( + 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), _streamed(_ROUND_ONE_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-disconnect-in-tool-call") + echo_then_disconnect: Final = AsyncMock( + side_effect=[[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}], asyncio.CancelledError()] + ) + + with ( + _echo_tool_served("PING"), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", echo_then_disconnect), + pytest.raises(asyncio.CancelledError), + ): + await _drain(await _open_stream_auto_execute(logging_obj)) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one_twice_with_one_fixed_margin(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_TWO_BODY),)) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-one-round")) + await _settle(lambda: len(recorder.logged) > 0) + + assert _output_types(_final_response(events).output) == ("message",) + _assert_one_row_bills(recorder.logged, _ROUND_TWO_BODY) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_both_rounds_with_one_fixed_margin_when_the_final_round_is_incomplete( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + truncated: Final = _streamed( + _ROUND_TWO_BODY, status="incomplete", incomplete_details={"reason": "max_output_tokens"} + ) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), truncated)) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-final-round-incomplete")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events, "response.incomplete")) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_a_failed_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + server_error: Final = {"code": "server_error", "message": "The server had an error processing your request."} + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY, status="failed", error=server_error)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-failed") + + with pytest.raises(litellm.ServiceUnavailableError): + await _stream_auto_execute(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 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (149, 5) + assert recorder.failed[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(recorder.failed[0]) + failure_row: Final = _request_failure_row(logging_obj, recorder.logged[0]) + _assert_billed_tokens(failure_row, (149, 5)) + assert failure_row["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN) + + +_XAI_MODEL: Final = "xai/grok-4.7" +_XAI_RESPONSES_URL: Final = r"https://api\.x\.ai/v1/responses" + + +def _xai_round(body: Mapping[str, object], cost_in_usd_ticks: int) -> Mapping[str, object]: + usage: Final = cast(Mapping[str, object], body["usage"]) + stated: Final = {**usage, "cost_in_usd_ticks": cost_in_usd_ticks, "context_details": {"cached_tokens": 3}} + return MappingProxyType({**body, "model": "grok-4.7", "usage": stated}) + + +def _xai_round_cost(body: Mapping[str, object]) -> float: + usage: Final = cast(Mapping[str, object], body["usage"]) + stated: Final = {**usage, "cost": xai_reported_cost_in_usd(usage["cost_in_usd_ticks"])} + return litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate({**body, "usage": stated}), + model=_XAI_MODEL, + custom_llm_provider="xai", + call_type="aresponses", + ) + + +@pytest.mark.asyncio +async def test_xai_stream_auto_execute_bills_each_round_at_its_own_price_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": {"fixed_amount": _FIXED_MARGIN}}) + rounds: Final = (_xai_round(_ROUND_ONE_BODY, 12_400_000), _xai_round(_ROUND_TWO_BODY, 8_800_000)) + _serve_responses(respx_mock, monkeypatch, tuple(_streamed(body) for body in rounds), _XAI_RESPONSES_URL) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await litellm.aresponses( + model=_XAI_MODEL, + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="xai-mcp-rounds", + stream=True, + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-xai-stream-rounds"), + ) + ) + await _settle(lambda: len(recorder.logged) > 0) + + charged_once: Final = _xai_round_cost(rounds[0]) + _xai_round_cost(rounds[1]) - _FIXED_MARGIN + client: Final = _final_response(events) + assert len(recorder.logged) == 1 + assert recorder.logged[0]["response_cost"] == pytest.approx(charged_once) + assert recorder.logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert client.usage.cost == pytest.approx(charged_once) + assert json.loads(events[-1].model_dump_json())["response"]["usage"]["context_details"] == {"cached_tokens": 6} + + +_UNPRICED_MODEL: Final = "litellm-unpriced-mcp-rounds" + + +@pytest.mark.asyncio +async def test_stream_auto_execute_on_an_unpriced_model_logs_both_rounds_as_a_failed_cost_lookup_not_a_free_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + rounds: Final = tuple( + MappingProxyType({**body, "model": _UNPRICED_MODEL}) for body in (_ROUND_ONE_BODY, _ROUND_TWO_BODY) + ) + _serve_rounds(respx_mock, monkeypatch, rounds) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await litellm.aresponses( + model=f"openai/{_UNPRICED_MODEL}", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=True, + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-stream-unpriced-rounds"), + ) + ) + await _settle(lambda: len(recorder.logged) > 0) + + client: Final = _final_response(events) + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["response_cost_failure_debug_info"] is not None + assert client.usage is not None + assert client.usage.cost is None + + +async def _stop_reading_once_the_follow_up_answers(logging_obj: Logging) -> None: + stream: Final = await _open_stream_auto_execute(logging_obj) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_item.added" and _output_types((event.item,)) == ("message",): + return + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_finished_first_round_when_the_caller_stops_reading_the_follow_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _stop_reading_once_the_follow_up_answers(_rounds_logging_obj(recorder, "mcp-stream-caller-stops")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + + +async def _stream_read_until_the_follow_up_answers(logging_obj: Logging) -> object: + stream: Final = await _open_stream_auto_execute(logging_obj) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_item.added" and _output_types((event.item,)) == ("message",): + return stream + pytest.fail("the follow-up round never answered") + + +_STREAM_BILLING_NAMES: Final = ("StreamedRoundBilling", "dispatch_success_handlers", "StreamingIterator") + + +def _names_stream_billing(text: str) -> bool: + return any(name in text for name in _STREAM_BILLING_NAMES) + + +class _UnraisableRecorder: + def __init__(self) -> None: + self.raised: tuple[str, ...] = () + + def __call__(self, unraisable: "sys.UnraisableHookArgs") -> None: + self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") + + +def test_stream_auto_execute_dropped_after_its_event_loop_closed_discards_its_parked_log_cleanly( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + unraisable: Final = _UnraisableRecorder() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + + with _echo_tool_served("PING"), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + asyncio.run(_stream_read_until_the_follow_up_answers(_rounds_logging_obj(recorder, "mcp-stream-loop-closed"))) + gc.collect() + + assert recorder.logged == () + assert [raised for raised in unraisable.raised if _names_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_stream_billing(str(warning.message))] == [] + + +_BRIDGED_MODEL: Final = "anthropic/claude-haiku-5-5" +_ANTHROPIC_MESSAGES_URL: Final = r"https://api\.anthropic\.com/v1/messages" + + +def _bridged_round(content: bytes) -> httpx.Response: + return httpx.Response(200, content=content, headers=SSE_HEADERS) + + +def _bridged_echo_call_round(cache_usage: Mapping[str, object] = MappingProxyType({})) -> httpx.Response: + return _bridged_round( + _anthropic_round( + "msg_echo_call", + {"type": "tool_use", "id": "call_echo", "name": "echo", "input": {}}, + {"type": "input_json_delta", "partial_json": json.dumps({"value": "PING"})}, + "tool_use", + model="claude-haiku-5-5", + usage=(533, 40), + cache_usage=cache_usage, + ) + ) + + +def _bridged_echo_answer_round( + stop: str = "end_turn", cache_usage: Mapping[str, object] = MappingProxyType({}) +) -> httpx.Response: + return _bridged_round( + _anthropic_round( + "msg_echo_answer", + {"type": "text", "text": ""}, + {"type": "text_delta", "text": "PING"}, + stop, + model="claude-haiku-5-5", + usage=(470, 4), + cache_usage=cache_usage, + ) + ) + + +def _anthropic_prompt_cache_usage(written: int, read: int) -> Mapping[str, object]: + return MappingProxyType( + { + "cache_creation_input_tokens": written, + "cache_read_input_tokens": read, + "cache_creation": {"ephemeral_5m_input_tokens": written, "ephemeral_1h_input_tokens": 0}, + "inference_geo": "us", + "speed": "standard", + } + ) + + +class _CutAfterTheMessageStart(httpx.AsyncByteStream): + async def __aiter__(self): + yield _sse((_anthropic_message_start("msg_echo_answer", "claude-haiku-5-5", 470),)) + raise httpx.ReadError("connection reset mid-stream") + + +def _bridged_round_cost(prompt_tokens: int, completion_tokens: int) -> float: + return litellm.completion_cost( + completion_response=ModelResponse( + model="claude-haiku-5-5", + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ), + model=_BRIDGED_MODEL, + ) + + +async def _open_bridged_stream_auto_execute(logging_obj: Logging) -> object: + return await litellm.aresponses( + model=_BRIDGED_MODEL, + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-ant-mcp-rounds", + stream=True, + litellm_logging_obj=logging_obj, + ) + + +def _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder: _RoundRecorder, answer: str) -> None: + assert len(recorder.results) == 1 + logged: Final = recorder.results[0] + assert isinstance(logged, ModelResponse) + assert logged.choices[0].message.content == answer + usage: Final[object] = getattr(logged, "usage", None) + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens) == (533 + 470, 40 + 4) + + +def _assert_one_bridged_row_bills_both_rounds_once(recorder: _RoundRecorder, client: ResponsesAPIResponse) -> None: + charged_once: Final = _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) - _FIXED_MARGIN + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder, "PING") + assert row["id"] == client.id + assert "message" in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (533 + 470, 40 + 4) + assert row["cache_hit"] is False + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (533 + 470, 40 + 4) + assert client.usage.cost == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_logs_the_final_response_billed_for_both_rounds_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-stream-rounds")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert _output_types(_final_response(events).output) == ("mcp_call", "message") + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_logs_the_provider_usage_of_both_rounds_summed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + ( + _bridged_echo_call_round(_anthropic_prompt_cache_usage(written=10, read=5)), + _bridged_echo_answer_round(cache_usage=_anthropic_prompt_cache_usage(written=3, read=7)), + ), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-provider-usage")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + usage: Final = row["metadata"]["usage_object"] + assert "choices" in row["response"] + assert (row["prompt_tokens"], row["completion_tokens"]) == (533 + 10 + 5 + 470 + 3 + 7, 40 + 4) + assert (usage["cache_creation_input_tokens"], usage["cache_read_input_tokens"]) == (10 + 3, 5 + 7) + assert usage["prompt_tokens_details"]["cache_creation_token_details"]["ephemeral_5m_input_tokens"] == 10 + 3 + assert (usage["inference_geo"], usage["speed"]) == ("us", "standard") + assert usage["cost"] == pytest.approx(row["response_cost"]) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_both_rounds_with_one_fixed_margin_when_the_final_round_is_incomplete( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), _bridged_echo_answer_round("max_tokens")), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-stream-incomplete")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events, "response.incomplete")) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_on_release( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-rounds-deferred") + + with _echo_tool_served("PING"): + stream: Final = await _open_bridged_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + await asyncio.sleep(0.2) + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + assert logged_before_release == () + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_rounds_on_the_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-rounds-blocked") + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served("PING"): + stream: Final = await _open_bridged_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _drain(stream) + ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await _settle(lambda: len(recorder.failed) > 0) + + charged_once: Final = _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) - _FIXED_MARGIN + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert recorder.failed[0]["response_cost"] == pytest.approx(charged_once) + assert recorder.failed[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stop", ("end_turn", "max_tokens"), ids=("completed", "incomplete")) +async def test_bridged_stream_auto_execute_bills_both_rounds_when_the_caller_stops_reading_after_the_answer_text( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stop: str +): + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round(stop)), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + + async def stop_reading_once_the_answer_text_is_done() -> None: + stream: Final = await _open_bridged_stream_auto_execute( + _rounds_logging_obj(recorder, "mcp-bridged-stream-caller-stops") + ) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_text.done": + return + + with _echo_tool_served("PING"): + await stop_reading_once_the_answer_text_is_done() + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder, "PING") + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert recorder.logged[0]["response_cost"] == pytest.approx( + _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) + ) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), httpx.Response(200, stream=_CutAfterTheMessageStart(), headers=SSE_HEADERS)), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-breaks") + + with _echo_tool_served("PING"), pytest.raises(MidStreamFallbackError): + await _drain(await _open_bridged_stream_auto_execute(logging_obj)) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert len(recorder.logged) == 1 + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (533, 40) + assert recorder.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(533, 40)) + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (470, 0) + assert recorder.failed[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 0) - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(recorder.failed[0]) + failure_row: Final = _request_failure_row(logging_obj, recorder.logged[0]) + _assert_billed_tokens(failure_row, (470, 0)) + assert failure_row["response_cost"] == pytest.approx(_bridged_round_cost(470, 0) - _FIXED_MARGIN) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_a_broken_follow_up_its_own_fixed_margin_after_a_cached_first_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + _serve_responses( + respx_mock, + monkeypatch, + ( + _bridged_echo_call_round(), + _bridged_echo_answer_round(), + httpx.Response(200, stream=_CutAfterTheMessageStart(), headers=SSE_HEADERS), + ), + _ANTHROPIC_MESSAGES_URL, + ) + warm: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(warm, "mcp-bridged-cache-then-break-warm")) + ) + await _settle(lambda: len(warm.logged) > 0) + + recorder: Final = _RoundRecorder() + with _echo_tool_served("PONG"), pytest.raises(MidStreamFallbackError): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-cache-then-break")) + ) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert [row["response_cost"] for row in recorder.logged] == [0.0] + assert len(recorder.failed) == 1 + assert recorder.failed[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 0)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_bridged_echo_answer_round(),), _ANTHROPIC_MESSAGES_URL) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _drain(await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-one-round"))) + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + assert "choices" in recorder.logged[0]["response"] + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (470, 4) + assert recorder.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 4)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_nothing_for_a_round_served_from_the_response_cache( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + provider: Final = _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), _bridged_echo_answer_round(), _bridged_echo_answer_round()), + _ANTHROPIC_MESSAGES_URL, + ) + warm: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain(await _open_bridged_stream_auto_execute(_rounds_logging_obj(warm, "mcp-bridged-cache-warm"))) + await _settle(lambda: len(warm.logged) > 0) + + first_cached: Final = _RoundRecorder() + with _echo_tool_served("PONG"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(first_cached, "mcp-bridged-cache-first-round")) + ) + await _settle(lambda: len(first_cached.logged) > 0) + + both_cached: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(both_cached, "mcp-bridged-cache-both-rounds")) + ) + await _settle(lambda: len(both_cached.logged) > 0) + + client: Final = _final_response(events) + assert provider.call_count == 3 + assert len(first_cached.logged) == 1 + assert (first_cached.logged[0]["prompt_tokens"], first_cached.logged[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert first_cached.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 4)) + assert first_cached.logged[0]["cache_hit"] is False + assert [(row["response_cost"], row["cache_hit"]) for row in both_cached.logged] == [(0.0, True)] + assert client.usage is not None + assert client.usage.cost == pytest.approx(_bridged_round_cost(470, 4)) diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 04d2b9ddf81..7445070bd3a 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -23,7 +23,8 @@ import asyncio import json from datetime import datetime from importlib import import_module -from typing import Final +from collections.abc import Mapping +from typing import Final, cast from unittest.mock import Mock, patch import httpx @@ -31,6 +32,7 @@ import pytest import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id @@ -44,6 +46,7 @@ from litellm.responses.streaming_iterator import ( _status_code_for_error_fields, ) from litellm.utils import ProviderConfigManager +from litellm.types.utils import StandardLoggingPayload from litellm.types.llms.openai import ( ErrorEvent, ErrorEventError, @@ -523,6 +526,82 @@ def test_handle_logging_failed_response_records_usage_and_cost(): iterator.logging_obj.response_cost_calculator.assert_called_once_with(result=chunk.response) +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failed: tuple[StandardLoggingPayload, ...] = () + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + + +async def _first_failure_row(recorder: _FailureRecorder) -> StandardLoggingPayload: + for _ in range(200): + if recorder.failed: + return recorder.failed[0] + await asyncio.sleep(0.01) + pytest.fail("the failed stream logged no failure row") + + +@pytest.mark.asyncio +async def test_failed_stream_bills_only_its_own_usage_over_an_earlier_attempts_stash(): + recorder: Final = _FailureRecorder() + logging_obj: Final = LiteLLMLoggingObj( + model="gpt-5.6", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime(2026, 1, 1), + litellm_call_id="responses-failed-over-stash", + function_id="responses-failed-over-stash", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.6", optional_params={}, litellm_params=get_litellm_params(custom_llm_provider="openai") + ) + logging_obj.model_call_details["combined_usage_object"] = litellm.Usage( + prompt_tokens=100, completion_tokens=50, total_tokens=150 + ) + logging_obj.model_call_details["response_cost"] = 0.5 + failed: Final = { + "type": "response.failed", + "sequence_number": 1, + "response": { + "id": "resp_failed_over_stash", + "object": "response", + "created_at": 1790000000, + "status": "failed", + "model": "gpt-5.6", + "output": [], + "error": {"code": "server_error", "message": "boom"}, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + } + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=f"data: {json.dumps(failed)}\n\n".encode() + ), + model="gpt-5.6", + responses_api_provider_config=ProviderConfigManager.get_provider_responses_api_config( + provider="openai", model="gpt-5.6" + ), + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + with pytest.raises(MidStreamFallbackError): + await anext(aiter(iterator)) + + row: Final = await _first_failure_row(recorder) + prompt_cost, completion_cost = litellm.cost_per_token( + model="gpt-5.6", prompt_tokens=10, completion_tokens=5, custom_llm_provider="openai" + ) + assert (row["prompt_tokens"], row["completion_tokens"]) == (10, 5) + assert row["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + def test_handle_logging_failed_response_without_usage_skips_recording(): iterator = _make_iterator() iterator.completed_response = _make_failed_chunk( diff --git a/tests/unit/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py index ad61f43c5a0..0ed22658144 100644 --- a/tests/unit/router_strategy/adaptive_router/test_hooks.py +++ b/tests/unit/router_strategy/adaptive_router/test_hooks.py @@ -10,6 +10,7 @@ from litellm.router_strategy.adaptive_router.config import ( ) from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, _recent_tool_results, _resolve_session_key, ) @@ -337,8 +338,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): """The header hook returns the `x-litellm-adaptive-router-model` header so proxy header construction picks it up (works for both streaming and non-streaming; `async_post_call_success_hook` is too late for streaming).""" - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"adaptive_router_chosen_model": "smart"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -348,8 +348,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"litellm_session_id": "sess-A"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -359,8 +358,7 @@ async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_no_metadata(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -370,8 +368,7 @@ async def test_post_call_response_headers_hook_noop_when_no_metadata(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_not_dict(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": "not-a-dict"}, user_api_key_dict=MagicMock(), response=MagicMock(), diff --git a/tests/unit/router_strategy/test_base_routing_strategy.py b/tests/unit/router_strategy/test_base_routing_strategy.py index dc75c3d2916..19ff11dd3f6 100644 --- a/tests/unit/router_strategy/test_base_routing_strategy.py +++ b/tests/unit/router_strategy/test_base_routing_strategy.py @@ -1,15 +1,20 @@ import json import logging -from typing import Any, Dict, List, Optional, Set, Union +from typing import Any, Dict, Final, List, Optional, Set, Union import pytest import asyncio +import threading from unittest.mock import MagicMock, patch +from fakeredis import TcpFakeServer + from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy @@ -162,3 +167,24 @@ async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_erro assert caplog.records == [] assert base_strategy.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_an_increment_counted_after_retire_still_reaches_redis(): + server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis") + threading.Thread(target=server.serve_forever, daemon=True).start() + redis_cache: Final = RedisCache(host="127.0.0.1", port=server.server_address[1]) + strategy: Final = BaseRoutingStrategy( + dual_cache=DualCache(redis_cache=redis_cache, in_memory_cache=InMemoryCache()), + should_batch_redis_writes=True, + default_sync_interval=60, + ) + try: + strategy.retire() + await strategy._increment_value_in_current_window("counted-after-retire", 3, 60) + await asyncio.sleep(0.2) + + assert float(await redis_cache.async_get_cache("counted-after-retire")) == 3 + finally: + server.shutdown() + server.server_close() diff --git a/tests/unit/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py index 4667ac67efa..10704a29347 100644 --- a/tests/unit/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/unit/router_strategy/test_budget_limiter_hotpath.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import gc import logging from types import SimpleNamespace @@ -492,6 +493,8 @@ def _new_router_budget_limiter( budget_limiter._redis_increment_queue_lock = queue_lock if queue_lock is not None else asyncio.Lock() budget_limiter._redis_increment_flush_lock = asyncio.Lock() budget_limiter._detached_increment_operations = None + budget_limiter._retired = False + budget_limiter._flush_after_retire = None return budget_limiter @@ -784,3 +787,57 @@ async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations: await limiter._push_in_memory_increments_to_redis() assert redis_cache.values[_SPEND_KEY] == 10.0 assert limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_queued_when_a_limiter_is_retired_reaches_redis_after_a_failed_write() -> None: + pipeline_started: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + pipeline_completed=pipeline_completed, + should_fail_pipeline=True, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(asyncio.sleep(3600)) + + budget_limiter.retire() + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert redis_cache.values == {} + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_in_flight_when_a_limiter_is_retired_reaches_redis_after_that_write_fails() -> None: + pipeline_started: Final = asyncio.Event() + allow_pipeline_to_complete: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + pipeline_completed=pipeline_completed, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(budget_limiter.periodic_sync_in_memory_spend_with_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + + budget_limiter.retire() + allow_pipeline_to_complete.set() + with contextlib.suppress(asyncio.CancelledError): + await budget_limiter._sync_task + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == [] diff --git a/tests/unit/router_strategy/test_lowest_cost.py b/tests/unit/router_strategy/test_lowest_cost.py index b53ec6e6ba2..e3fdcd69499 100644 --- a/tests/unit/router_strategy/test_lowest_cost.py +++ b/tests/unit/router_strategy/test_lowest_cost.py @@ -112,6 +112,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -164,6 +165,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -184,6 +186,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -206,6 +209,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_get_available_deployments(): @@ -235,6 +239,7 @@ async def test_get_available_deployments(): assert selected_model["model_info"]["id"] == "groq-llama" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_get_available_deployments_custom_price(): @@ -286,6 +291,7 @@ async def test_get_available_deployments_custom_price(): assert selected_model["model_info"]["id"] == "chatgpt-v-1" + async def _deploy(lowest_cost_logger, deployment_id, tokens_used, duration): kwargs = { "litellm_params": { @@ -307,6 +313,7 @@ async def _deploy(lowest_cost_logger, deployment_id, tokens_used, duration): end_time=end_time, ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize("ans_rpm", [1, 5]) # 1 should produce nothing, 10 should select first @pytest.mark.asyncio @@ -357,3 +364,42 @@ async def test_get_available_endpoints_tpm_rpm_check_async(ans_rpm): assert (d_ans and d_ans["model_info"]["id"]) == ans print("selected deployment:", d_ans) + + +_TWO_DOLLARS_PER_TOKEN = {"input_cost_per_token": 1.0, "output_cost_per_token": 1.0} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "azure_litellm_params, base_model", + [ + ({"model": "azure/my-gpt4-deploy"}, "azure/gpt-4.1"), + ({"model": "azure/gpt-4.1"}, "azure/not-a-real-model"), + ], + ids=["base_model_priced", "unmapped_base_model_falls_back_to_deployment_model"], +) +async def test_cost_based_routing_prices_an_azure_deployment_through_its_base_model( + azure_litellm_params: dict, base_model: str +): + handler = LowestCostLoggingHandler(router_cache=DualCache()) + healthy_deployments = [ + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1", **_TWO_DOLLARS_PER_TOKEN}, + "model_info": {"id": "openai-two-dollars"}, + }, + { + "model_name": "gpt-4.1", + "litellm_params": azure_litellm_params, + "model_info": {"id": "azure-cost-map-priced", "base_model": base_model}, + }, + ] + + picked = await handler.async_get_available_deployments( + model_group="gpt-4.1", healthy_deployments=healthy_deployments + ) + + assert picked is not None and picked["model_info"]["id"] == "azure-cost-map-priced", ( + "a cost-map priced Azure deployment must beat a $2 per token deployment, " + "it loses only when it is scored at the $10 unmapped default" + ) diff --git a/tests/unit/router_utils/test_deployment_pricing.py b/tests/unit/router_utils/test_deployment_pricing.py new file mode 100644 index 00000000000..45cc211a89c --- /dev/null +++ b/tests/unit/router_utils/test_deployment_pricing.py @@ -0,0 +1,53 @@ +import pytest + +from litellm.router_utils.deployment_pricing import deployment_pricing_models, pricing_models + + +@pytest.mark.parametrize( + "base_model, model, expected", + [ + ("azure/gpt-4.1", "azure/my-gpt4-deploy", ("azure/gpt-4.1", "azure/my-gpt4-deploy")), + (None, "azure/my-gpt4-deploy", ("azure/my-gpt4-deploy",)), + ("", "azure/my-gpt4-deploy", ("azure/my-gpt4-deploy",)), + ("azure/gpt-4.1", "azure/gpt-4.1", ("azure/gpt-4.1",)), + (None, None, ()), + ], + ids=["base_model_first", "no_base_model", "blank_base_model", "same_name_once", "nothing_named"], +) +def test_pricing_models_walks_base_model_then_deployment_model(base_model, model, expected): + assert pricing_models(base_model, model) == expected + + +@pytest.mark.parametrize( + "deployment, expected", + [ + ( + {"litellm_params": {"model": "azure/my-gpt4-deploy"}, "model_info": {"base_model": "azure/gpt-4.1"}}, + ("azure/gpt-4.1", "azure/my-gpt4-deploy"), + ), + ( + {"litellm_params": {"model": "azure/my-gpt4-deploy", "base_model": "azure/gpt-4.1"}}, + ("azure/gpt-4.1", "azure/my-gpt4-deploy"), + ), + ( + { + "litellm_params": {"model": "azure/my-gpt4-deploy", "base_model": "azure/gpt-4o"}, + "model_info": {"base_model": "azure/gpt-4.1"}, + }, + ("azure/gpt-4.1", "azure/my-gpt4-deploy"), + ), + ({"litellm_params": {"model": "openai/gpt-4.1"}, "model_info": {}}, ("openai/gpt-4.1",)), + ({"litellm_params": "not-a-dict", "model_info": None}, ()), + ("not-a-deployment", ()), + ], + ids=[ + "model_info_base_model", + "params_base_model", + "model_info_wins", + "no_base_model", + "malformed_blocks", + "not_a_dict", + ], +) +def test_deployment_pricing_models_reads_the_deployment_dict(deployment, expected): + assert deployment_pricing_models(deployment) == expected diff --git a/tests/unit/rust_bridge/AGENTS.md b/tests/unit/rust_bridge/AGENTS.md index 64c9fbbadd3..abf5ac6bed7 100644 --- a/tests/unit/rust_bridge/AGENTS.md +++ b/tests/unit/rust_bridge/AGENTS.md @@ -4,6 +4,6 @@ Test what each side of the bridge does, not the rollout policy that picks a side Call each path directly with an explicit decision instead. The Python path is the implementation the dispatcher falls back to, e.g. `litellm.responses.main.responses`. The Rust path is the native binding, e.g. `NATIVE_OCR.load()` from `litellm/rust_bridge/ocr/entrypoints.py`, called with the `NativeCall` envelope that dispatch would hand it for every public inference route. When the native side reads a policy-derived setting such as `settings.secret_manager().native`, pin that field in the test instead of deriving it from the catalog. `ocr/test_secrets.py` shows the pattern -Rollout policy itself, meaning which rule matches and what `LITELLM_RUST` changes, belongs in `test_catalog.py`, `test_configuration.py` and `test_dispatch.py`, tested against rules the test builds rather than the shipped `catalog.RULES` +Rollout policy itself, meaning which rule matches and what `LITELLM_RUST` changes, belongs in `test_catalog.py` and `test_configuration.py`, tested against rules the test builds rather than the shipped `catalog.RULES`. Which path a public call takes under those rules is a dispatch test and lives in `tests/test_litellm_rust/` for now Before adding a test here, ask whether it checks something Rust cannot. A `host/routes/.py` module is the Python half of a native route: it projects Python-only state (the cost map, `litellm.*` settings, request kwargs) into the plain values the Rust side consumes, and maps native failures back onto public exceptions. Those projections are what belongs here, because a wrong key or an ignored provider prefix ships the wrong value to Rust and no Rust test sees it. `host/routes/test_messages.py` shows the shape. Behavior that lives in Rust (a request transform given its inputs, header assembly, stream relay) is tested in the crate, and the route end to end is tested against a recording server in `tests/test_litellm_rust/`. A test that only re-checks a Python helper the route host happens to call is a duplicate of that helper's own test and should not be added diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index eb1bfd82715..a40f8a77105 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -9,6 +9,7 @@ import pytest from pydantic import TypeAdapter from litellm._internal_context import is_internal_call +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.rust_bridge.host import callbacks_legacy_python as legacy from litellm.rust_bridge.host.callbacks_legacy_python import failure_handler, setup @@ -167,6 +168,75 @@ def test_failure_handler_of_an_internal_call_leaves_the_outer_budget_reservation pending.close() +class _Hooks(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.hooks: list[str] = [] + + def log_success_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.hooks.append("sync_success") + + async def async_log_success_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.hooks.append("async_success") + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.hooks.append("sync_failure") + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.hooks.append("async_failure") + + +def _logger_with_hooks(hooks: _Hooks) -> Logging: + logger: Final = Logging( + model="mistral/mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="hooks", + function_id="hooks", + dynamic_success_callbacks=[hooks], + dynamic_async_success_callbacks=[hooks], + dynamic_failure_callbacks=[hooks], + dynamic_async_failure_callbacks=[hooks], + ) + logger.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + return logger + + +@pytest.mark.asyncio +async def test_failure_of_an_async_call_gives_a_logger_with_both_failure_hooks_only_the_async_one() -> None: + hooks: Final = _Hooks() + now: Final = datetime.datetime.now() + kwargs: Final = {**_OCR_KWARGS, "litellm_logging_obj": _logger_with_hooks(hooks)} + result: Final = setup("aocr", (), kwargs, now, asynchronous=True) + + assert failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=False) is None + pending: Final = failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=True) + assert pending is not None + await pending + + assert hooks.hooks == ["async_failure"] + + +@pytest.mark.parametrize("earlier_async_call", [False, True]) +def test_failure_of_a_sync_call_keeps_the_sync_failure_hook(earlier_async_call: bool) -> None: + hooks: Final = _Hooks() + now: Final = datetime.datetime.now() + kwargs: Final = {**_OCR_KWARGS, "litellm_logging_obj": _logger_with_hooks(hooks)} + if earlier_async_call: + setup("aocr", (), kwargs, now, asynchronous=True) + result: Final = setup("ocr", (), kwargs, now, asynchronous=False) + + assert failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=False) is None + + assert hooks.hooks == ["sync_failure"] + + CONTRACT_PATH: Final = Path(__file__).parents[4] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json" diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index b99342388d6..704a4f73fe8 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -7,10 +7,12 @@ import pytest from litellm.rust_bridge import catalog, configuration from litellm.rust_bridge.catalog import ( + Feature, Route, RouteContext, RouteRule, Rules, + Target, ) from litellm.rust_bridge.configuration import Decision, Rollout @@ -23,36 +25,6 @@ def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: configuration.reset_rust_configuration() -@pytest.mark.parametrize("route", tuple(Route)) -@pytest.mark.parametrize( - "provider", (None, "bedrock", "mistral", "anthropic", "deepseek", "vertex_ai", "openai", "azure_ai", "unknown") -) -@pytest.mark.parametrize("process", (None, False, True)) -@pytest.mark.parametrize("environment", (None, "0", "1")) -def test_shipped_decisions( - monkeypatch: pytest.MonkeyPatch, - route: Route, - provider: str | None, - process: bool | None, - environment: str | None, -) -> None: - configuration.rust(process) - if environment is not None: - monkeypatch.setenv("LITELLM_RUST", environment) - context: Final = RouteContext(route, provider=provider, model="test-model") - - if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): - assert catalog.rollout(context) is Rollout.RUST_REQUIRED - assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.MESSAGES and provider in ("anthropic", "deepseek", "vertex_ai"): - assert catalog.rollout(context) is Rollout.RUST_OPT_IN - opted_in: Final = environment == "1" or (environment is None and process is True) - assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON) - else: - assert catalog.rollout(context) is Rollout.PYTHON_ONLY - assert catalog.decision(context) is Decision.PYTHON - - @pytest.mark.parametrize("route", tuple(Route)) def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pytest.MonkeyPatch, route: Route) -> None: configuration.rust(True) @@ -72,6 +44,11 @@ def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pyt (RouteContext(Route.RESPONSES, provider="openai", model="other"), Decision.PYTHON), (RouteContext(Route.RESPONSES, provider="anthropic", model="m"), Decision.PYTHON), (RouteContext(Route.MESSAGES, provider="openai", model="m"), Decision.PYTHON), + (RouteContext(Route.MESSAGES, target=Target.MESSAGES), Decision.RUST_REQUIRED), + (RouteContext(Route.MESSAGES, target=Target.CHAT_COMPLETIONS), Decision.PYTHON), + (RouteContext(Route.MESSAGES), Decision.PYTHON), + (RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), Decision.PYTHON), + (RouteContext(Route.OCR), Decision.RUST_REQUIRED), ), ) def test_first_matching_rule_respects_every_constraint(context: RouteContext, expected: Decision) -> None: @@ -83,24 +60,14 @@ def test_first_matching_rule_respects_every_constraint(context: RouteContext, ex models=frozenset({"m"}), ), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, targets=frozenset({Target.MESSAGES})), + RouteRule(Route.OCR, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED), ) assert catalog.decision(context, rules) is expected -@pytest.mark.parametrize("process", (None, False, True)) -@pytest.mark.parametrize("environment", (None, "0", "1")) -def test_ocr_has_no_python_path_to_opt_out_to( - monkeypatch: pytest.MonkeyPatch, process: bool | None, environment: str | None -) -> None: - configuration.rust(process) - if environment is not None: - monkeypatch.setenv("LITELLM_RUST", environment) - - assert catalog.decision(RouteContext(Route.OCR, model="m")) is Decision.RUST_REQUIRED - assert catalog.decision(RouteContext(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED - - @pytest.mark.parametrize( ("context", "expected"), ( @@ -154,6 +121,10 @@ def test_empty_constraints_match_nothing() -> None: rules: Final[Rules] = ( RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset()), RouteRule(Route.OCR, Rollout.RUST_REQUIRED, models=frozenset()), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, features=frozenset()), ) assert catalog.decision(RouteContext(Route.OCR), rules) is Decision.PYTHON + assert catalog.decision( + RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), rules + ) is Decision.PYTHON diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py new file mode 100644 index 00000000000..726128386d5 --- /dev/null +++ b/tests/unit/rust_bridge/test_detectors.py @@ -0,0 +1,522 @@ +from collections.abc import Callable, Mapping +from copy import deepcopy +from datetime import datetime +from functools import reduce +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.rust_bridge import detectors +from litellm.rust_bridge.chat_completions import detectors as chat_detectors +from litellm.rust_bridge.messages import detectors as messages_detectors +from litellm.rust_bridge.public_call import NativeCall +from litellm.rust_bridge.responses import detectors as responses_detectors +from litellm.types.llms.anthropic import ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES + +GATEWAY_TOOL: Final = {"type": "mcp", "server_label": "gateway", "server_url": "litellm_proxy"} +REMOTE_TOOL: Final = {"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"} +ENCODED_ID: Final = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-1", "rs_1") # pyright: ignore[reportPrivateUsage] # the encoder litellm stamps on earlier turns +WRAPPED_CONTENT: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("opaque", "deployment-1") + + +MAPPED_FILE: Final = { + "model_file_id_mapping": {"litellm-file": {"deployment-1": "file-abc"}}, + "model_info": {"id": "deployment-1"}, +} +MAPPED_FILE_INPUT: Final = [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-file"}]}] +CACHED_TEXT: Final = {"type": "input_text", "text": "hi", "cache_control": {"type": "ephemeral"}} +THINKING_TURN: Final = {"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "s"}]} +REDACTED: Final = {"type": "redacted_thinking", "data": "x"} +TOOL_SEARCH_TYPE: Final = sorted(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES)[0] + + +def function_tool(parameters: object) -> Mapping[str, object]: + return {"type": "function", "name": "f", "parameters": parameters} + + +def call(**kwargs: object) -> NativeCall: + return NativeCall(args=(), kwargs=kwargs, base={"model": "gpt-5"}) + + +@pytest.mark.parametrize( + ("model", "provider", "tools", "expected"), + ( + ("deepseek/test-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], True), + ("anthropic/test-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], False), + ("test-model", "anthropic", [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], False), + ("unresolved-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], True), + ("deepseek/test-model", None, [{"type": "function", "name": "f"}], False), + ("deepseek/test-model", None, ["not-a-tool"], False), + ("deepseek/test-model", None, [], False), + ("deepseek/test-model", None, None, False), + ), +) +def test_messages_interception_matches_tools_and_resolved_provider( + model: str, provider: str | None, tools: object, expected: bool +) -> None: + assert messages_detectors.intercepted(call(model=model, custom_llm_provider=provider, tools=tools)) is expected + + +def _wrap_nested_level(value: object, depth: int) -> object: + return {"nested": value} if depth % 2 else [value] + + +@pytest.mark.parametrize( + ("detector", "kwargs", "expected"), + ( + (detectors.mock_response, {"mock_response": "hi"}, True), + (detectors.mock_response, {"mock_response": None}, False), + (detectors.mock_response, {"mock_tool_calls": [{"id": "call_1"}]}, False), + (chat_detectors.chat_mock_response, {"mock_response": "hi"}, True), + (chat_detectors.chat_mock_response, {"mock_tool_calls": [{"id": "call_1"}]}, True), + (chat_detectors.chat_mock_response, {"mock_timeout": True}, True), + (chat_detectors.chat_mock_response, {"mock_tool_calls": [], "mock_timeout": False}, False), + (responses_detectors.client_headers, {"headers": {"x-tenant": "a"}}, True), + (responses_detectors.client_headers, {"headers": {}}, False), + (responses_detectors.client_headers, {"headers": None}, False), + ( + chat_detectors.model_list, + {"model_list": [{"model_name": "gpt-5", "litellm_params": {"model": "gpt-5"}}]}, + True, + ), + (chat_detectors.model_list, {"model_list": None}, False), + (chat_detectors.chat_retries_fallbacks, {"context_window_fallback_dict": {"gpt-5": "gpt-5-mini"}}, True), + (chat_detectors.chat_retries_fallbacks, {"fallbacks": ["gpt-5-mini"]}, True), + (detectors.retries_fallbacks, {"context_window_fallback_dict": {"gpt-5": "gpt-5-mini"}}, False), + (detectors.prompt_management, {"model": "litellm_agent/gpt-5"}, True), + (detectors.prompt_management, {"model": "gpt-5"}, False), + (detectors.caller_client, {"client": object()}, True), + (detectors.caller_client, {"shared_session": object()}, True), + (detectors.caller_client, {"client": None, "shared_session": None}, False), + (detectors.caller_client, {}, False), + (detectors.extra_body, {"extra_body": {"stream": True}}, True), + (detectors.extra_body, {"extra_body": {"stream": False}}, True), + (detectors.extra_body, {"extra_body": {}}, False), + (detectors.extra_body, {"extra_body": None}, False), + (detectors.extra_body, {}, False), + ( + responses_detectors.provider_affinity, + {"provider_affinity_header": "X-Session-ID", "session_id": "session-1"}, + True, + ), + (responses_detectors.provider_affinity, {"provider_affinity_header": "X-Session-ID"}, True), + ( + responses_detectors.provider_affinity, + {"provider_affinity_header": "X-Session-ID", "extra_headers": {"x-session-id": "session-1"}}, + True, + ), + (responses_detectors.provider_affinity, {"provider_affinity_header": "Authorization"}, True), + (responses_detectors.provider_affinity, {"provider_affinity_header": "invalid header"}, True), + (responses_detectors.provider_affinity, {"provider_affinity_header": ""}, True), + (responses_detectors.provider_affinity, {"provider_affinity_header": None, "session_id": "session-1"}, False), + (responses_detectors.provider_affinity, {}, False), + (detectors.streaming, {"stream": True}, True), + (detectors.streaming, {"stream": False}, False), + (detectors.timeout_variants, {"request_timeout": 5}, True), + (detectors.timeout_variants, {"stream": True, "stream_timeout": 5}, True), + (detectors.timeout_variants, {"stream_timeout": 5}, False), + (detectors.timeout_variants, {"timeout": httpx.Timeout(60, connect=2)}, True), + (detectors.timeout_variants, {"timeout": 60}, False), + (detectors.proxy_mcp_tools, {"tools": [GATEWAY_TOOL]}, True), + (detectors.proxy_mcp_tools, {"tools": [REMOTE_TOOL]}, False), + (detectors.proxy_mcp_tools, {"tools": None}, False), + (detectors.prompt_management, {"prompt_id": "welcome"}, True), + (detectors.prompt_management, {"prompt_variables": {"name": "x"}}, False), + (responses_detectors.text_format, {"text_format": dict}, True), + (responses_detectors.text_format, {}, False), + (responses_detectors.chat_completions_bridge, {"use_chat_completions_api": True}, True), + (responses_detectors.chat_completions_bridge, {"use_chat_completions_api": False}, False), + (responses_detectors.chat_completions_bridge, {"model": "openai/chat_completions/gpt-5"}, True), + (responses_detectors.chat_completions_bridge, {"model": "openai/gpt-5"}, False), + (detectors.retries_fallbacks, {"num_retries": 0}, True), + (detectors.retries_fallbacks, {"max_retries": 2}, True), + (detectors.retries_fallbacks, {"fallbacks": ["gpt-5-mini"]}, True), + (detectors.retries_fallbacks, {"retry_policy": {"TimeoutErrorRetries": 2}}, True), + (detectors.retries_fallbacks, {}, False), + (detectors.prompt_management, {"knowledge_bases": ["kb-1"]}, True), + (detectors.prompt_management, {"vector_store_ids": []}, False), + (responses_detectors.managed_resource_ids, {"input": MAPPED_FILE_INPUT, **MAPPED_FILE}, True), + (responses_detectors.managed_resource_ids, {"input": MAPPED_FILE_INPUT}, False), + (responses_detectors.managed_resource_ids, {"input": "hi"}, False), + (responses_detectors.cache_control_marks, {"input": [{"role": "user", "content": [CACHED_TEXT]}]}, True), + ( + responses_detectors.cache_control_marks, + {"tools": [{"type": "function", "name": "f", "cache_control": {}}]}, + True, + ), + ( + responses_detectors.cache_control_marks, + {"input": [{"role": "user", "content": [{"type": "input_text"}]}]}, + False, + ), + ( + responses_detectors.anthropic_tool_search, + {"tools": [{"type": TOOL_SEARCH_TYPE, "name": "tool_search"}]}, + True, + ), + (responses_detectors.anthropic_tool_search, {"tools": [{"type": "function", "name": "f"}]}, False), + (responses_detectors.tool_schema_repair, {"tools": [function_tool('{"type": "object"}')]}, True), + ( + responses_detectors.tool_schema_repair, + {"tools": [function_tool({"type": "string", "pattern": r"\p{L}"})]}, + True, + ), + (responses_detectors.tool_schema_repair, {"tools": [function_tool({"type": "object"})]}, False), + (responses_detectors.tool_schema_repair, {}, False), + (responses_detectors.model_param_gating, {"max_output_tokens": 1}, True), + (responses_detectors.model_param_gating, {"max_output_tokens": 1024}, False), + (detectors.cache_control_injection, {"enable_prompt_caching": True}, True), + (detectors.cache_control_injection, {"cache_control_injection_points": [{"location": "message"}]}, True), + (detectors.cache_control_injection, {"cache_control_injection_points": []}, False), + (messages_detectors.thinking_signature_retry, {"messages": [THINKING_TURN]}, True), + ( + messages_detectors.thinking_signature_retry, + {"messages": [{"role": "assistant", "content": [REDACTED]}]}, + True, + ), + (messages_detectors.thinking_signature_retry, {"messages": [{"role": "assistant", "content": "done"}]}, False), + ( + messages_detectors.thinking_signature_retry, + {"messages": [{"role": "user", "content": [{"type": "text"}]}]}, + False, + ), + (detectors.drop_params, {"drop_params": True}, True), + (detectors.drop_params, {"drop_params": False}, False), + (detectors.drop_params, {"drop_params": " True "}, True), + (detectors.drop_params, {"drop_params": "false"}, False), + (detectors.drop_params, {"drop_params": "maybe"}, False), + (detectors.drop_params, {"additional_drop_params": ["metadata.user_id"]}, True), + (detectors.drop_params, {"additional_drop_params": []}, False), + ( + responses_detectors.deployment_response_id, + {"litellm_metadata": {"model_info": {"id": "deployment-1"}}}, + True, + ), + (responses_detectors.deployment_response_id, {"litellm_metadata": {"model_info": {}}}, False), + (responses_detectors.deployment_response_id, {"litellm_metadata": None}, False), + (responses_detectors.encoded_item_ids, {"input": [{"type": "reasoning", "id": ENCODED_ID}]}, True), + ( + responses_detectors.encoded_item_ids, + {"input": [{"type": "reasoning", "encrypted_content": WRAPPED_CONTENT}]}, + True, + ), + ( + responses_detectors.encoded_item_ids, + {"input": [{"type": "reasoning", "id": "rs_1", "encrypted_content": "x"}]}, + False, + ), + (responses_detectors.encoded_item_ids, {"input": "hi"}, False), + ), +) +def test_request_detectors( + detector: Callable[[NativeCall], bool], kwargs: Mapping[str, object], expected: bool +) -> None: + assert detector(call(**kwargs)) is expected + + +def test_cache_control_marks_handles_deeply_nested_input() -> None: + nested_input: Final = reduce(_wrap_nested_level, range(5_000), {"cache_control": {}}) + assert responses_detectors.cache_control_marks(call(input=nested_input)) + + +@pytest.mark.parametrize( + ("input_value", "expected"), + ( + pytest.param( + [{"type": "reasoning", "id": "rs_00000000-0000-4000-8000-000000000000", "summary": list[object]()}], + True, + id="bridge-minted-reasoning", + ), + pytest.param( + [{"type": "reasoning", "encrypted_content": '[{"type":"thinking","thinking":"x","signature":"s"}]'}], + True, + id="bridge-thinking-blocks", + ), + pytest.param( + [{"type": "function_call", "id": "foreign-id", "call_id": "call_1", "name": "f", "arguments": "{}"}], + True, + id="foreign-tool-call-id", + ), + pytest.param( + [{"type": "reasoning", "id": "rs_1", "summary": list[object](), "status": None}], + True, + id="null-reasoning-status", + ), + pytest.param( + [{"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "f", "arguments": "{}"}], + False, + id="native-tool-call-id", + ), + pytest.param( + [{"type": "custom_tool_call", "call_id": "call_1", "name": "f", "input": "x", "namespace": "ns"}], + True, + id="custom-tool-call-namespace", + ), + pytest.param( + [{"type": "custom_tool_call", "call_id": "call_1", "name": "f", "input": "x"}], + False, + id="custom-tool-call", + ), + pytest.param([{"type": "reasoning", "id": "rs_1", "summary": list[object]()}], False, id="native-reasoning"), + pytest.param([{"role": "user", "content": "hi"}], False, id="message"), + pytest.param("hi", False, id="string"), + pytest.param([], False, id="empty"), + pytest.param(None, False, id="absent"), + ), +) +def test_input_rewrite_preserves_caller_input(input_value: object, expected: bool) -> None: + original: Final = deepcopy(input_value) + assert responses_detectors.input_rewrite(call(input=input_value)) is expected + assert input_value == original + + +@pytest.mark.parametrize( + ("detector", "setting", "value", "kwargs"), + ( + (detectors.drop_params, "drop_params", True, {}), + (detectors.modify_params, "modify_params", True, {}), + (detectors.retries_fallbacks, "num_retries", 3, {}), + (detectors.cache_control_injection, "enable_anthropic_prompt_caching", True, {}), + (messages_detectors.stream_cost_injection, "include_cost_in_streaming_usage", True, {"stream": True}), + (messages_detectors.stream_cost_injection, "include_cost_in_streaming_usage", "true", {"stream": True}), + (chat_detectors.chat_retries_fallbacks, "model_fallbacks", [{"gpt-5": ["gpt-5-mini"]}], {}), + (chat_detectors.chat_retries_fallbacks, "num_retries", 3, {}), + (chat_detectors.model_alias_map, "model_alias_map", {"gpt-5": "openai/gpt-5"}, {}), + ), +) +def test_global_settings_are_detected( + monkeypatch: pytest.MonkeyPatch, + detector: Callable[[NativeCall], bool], + setting: str, + value: object, + kwargs: Mapping[str, object], +) -> None: + for name in ("drop_params", "modify_params", "enable_anthropic_prompt_caching", "include_cost_in_streaming_usage"): + monkeypatch.setattr(litellm, name, False) + for name in ("num_retries", "model_fallbacks"): + monkeypatch.setattr(litellm, name, None) + monkeypatch.setattr(litellm, "model_alias_map", {}) + assert not detector(call(**kwargs)) + monkeypatch.setattr(litellm, setting, value) + assert detector(call(**kwargs)) + + +def test_stream_cost_injection_needs_a_stream(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + assert not messages_detectors.stream_cost_injection(call()) + + +def test_model_alias_map_applies_only_to_the_requested_model(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"fast": "openai/gpt-5-mini"}) + assert chat_detectors.model_alias_map(call(model="fast")) + assert not chat_detectors.model_alias_map(call()) + + +@pytest.mark.parametrize( + ("kwargs", "expected"), + ( + pytest.param({"temperature": 0.2}, True, id="gpt-5-temperature"), + pytest.param({"max_tokens": 16}, False, id="gpt-5-max-tokens"), + pytest.param({"model": "gpt-4o", "temperature": 0.2}, False, id="gpt-4o-temperature"), + pytest.param({"model": "anthropic/claude-sonnet-4-5", "logit_bias": {"1": 1}}, True, id="unsupported-param"), + pytest.param({"model": "anthropic/claude-sonnet-4-5", "n": 1}, False, id="langchain-default-n"), + pytest.param({"model": "anthropic/claude-sonnet-4-5", "n": 2}, True, id="unsupported-n"), + pytest.param({"model": "anthropic/claude-sonnet-4-5", "max_tokens": 16, "stream": True}, False, id="supported"), + pytest.param({"model": "unknown-model", "temperature": 0.2}, False, id="unresolved-provider"), + pytest.param({}, False, id="no-params"), + ), +) +def test_chat_model_param_gating_follows_the_provider_param_rules( + monkeypatch: pytest.MonkeyPatch, kwargs: Mapping[str, object], expected: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + assert chat_detectors.chat_model_param_gating(call(**kwargs)) is expected + + +def test_openai_workload_identity_applies_only_without_a_static_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setattr(litellm, "api_base", None) + identity: Final = { + "openai_identity_provider_id": "idp-1", + "openai_service_account_id": "sa-1", + "openai_identity_token_file": "/var/run/token", + } + assert responses_detectors.openai_workload_identity(call(**identity)) + assert not responses_detectors.openai_workload_identity(call(**identity, api_key="sk-static")) + assert not responses_detectors.openai_workload_identity(call()) + + +def test_anthropic_workload_identity_applies_only_without_a_static_key(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_FEDERATION_RULE_ID", + "ANTHROPIC_ORGANIZATION_ID", + ): + monkeypatch.delenv(name, raising=False) + federation: Final = {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"} + assert messages_detectors.anthropic_workload_identity(call(**federation)) + assert not messages_detectors.anthropic_workload_identity(call(**federation, api_key="sk-ant-static")) + assert not messages_detectors.anthropic_workload_identity(call(anthropic_federation_rule_id="rule-1")) + + +class _Logger(CustomLogger): + pass + + +class _Rewrite(CustomLogger): + async def async_pre_request_hook(self, model: str, messages: list, kwargs: dict) -> dict | None: # pyright: ignore[reportMissingTypeArgument] # mirrors the CustomLogger signature + return kwargs + + +class _InheritedRewrite(_Rewrite): + pass + + +class _Agentic(CustomLogger): + async def async_should_run_agentic_loop(self, *args: object, **kwargs: object) -> tuple[bool, dict]: # pyright: ignore[reportIncompatibleMethodOverride, reportMissingTypeArgument] # test double only needs the override + return False, {} + + +class _ChatAgentic(CustomLogger): + async def async_should_run_chat_completion_agentic_loop( + self, *args: object, **kwargs: object + ) -> tuple[bool, dict[str, object]]: # pyright: ignore[reportIncompatibleMethodOverride] # test double only needs the override + return False, {} + + +class _Deployment(CustomLogger): + async def async_pre_call_deployment_hook(self, *args: object, **kwargs: object) -> dict[str, object] | None: # pyright: ignore[reportIncompatibleMethodOverride] # test double only needs the override + return None + + +@pytest.mark.parametrize( + ("detector", "callbacks", "expected"), + ( + (messages_detectors.pre_request_hook, [_Rewrite()], True), + (messages_detectors.pre_request_hook, [_InheritedRewrite()], True), + (messages_detectors.pre_request_hook, [_Logger(), _Agentic()], False), + (messages_detectors.pre_request_hook, ["not-a-logger"], False), + (detectors.agentic_loop_hook, [_Agentic()], True), + (detectors.agentic_loop_hook, [_Logger(), _Rewrite()], False), + (detectors.agentic_loop_hook, [_ChatAgentic()], False), + (chat_detectors.chat_agentic_loop_hook, [_ChatAgentic()], True), + (chat_detectors.chat_agentic_loop_hook, [_Agentic()], True), + (chat_detectors.chat_agentic_loop_hook, [_Logger(), _Deployment()], False), + (responses_detectors.deployment_hook, [_Deployment()], True), + (responses_detectors.deployment_hook, [_Logger(), _Agentic(), _Rewrite()], False), + ), +) +def test_callback_detectors_need_an_override( + monkeypatch: pytest.MonkeyPatch, detector: Callable[[NativeCall], bool], callbacks: list[object], expected: bool +) -> None: + monkeypatch.setattr(litellm, "callbacks", callbacks) + assert detector(call()) is expected + + +def _logging_with(callback: CustomLogger) -> Logging: + logging_obj: Final = Logging( + model="claude-sonnet-4-5", + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="function-1", + ) + logging_obj.add_dynamic_callback(callback) + return logging_obj + + +@pytest.mark.parametrize( + ("detector", "kwargs", "expected"), + ( + (detectors.agentic_loop_hook, {"success_callback": [_Agentic()]}, True), + (detectors.agentic_loop_hook, {"litellm_logging_obj": _logging_with(_Agentic())}, True), + (detectors.agentic_loop_hook, {"success_callback": [_Logger()]}, False), + (messages_detectors.pre_request_hook, {"success_callback": [_Rewrite()]}, False), + (chat_detectors.chat_agentic_loop_hook, {"litellm_logging_obj": _logging_with(_ChatAgentic())}, True), + (responses_detectors.deployment_hook, {"success_callback": [_Deployment()]}, True), + ), + ids=( + "success-callback", + "logging-object", + "plain-logger", + "pre-request-reads-globals-only", + "chat-agentic-logging-object", + "deployment-success-callback", + ), +) +def test_callback_detectors_read_the_request_callbacks_the_handler_reads( + monkeypatch: pytest.MonkeyPatch, + detector: Callable[[NativeCall], bool], + kwargs: Mapping[str, object], + expected: bool, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + assert detector(call(**kwargs)) is expected + + +def test_websearch_interception_needs_the_interception_logger(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger + + monkeypatch.setattr(litellm, "callbacks", [_Rewrite()]) + assert not messages_detectors.websearch_interception(call()) + monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger()]) + assert messages_detectors.websearch_interception(call()) + + +@pytest.mark.parametrize( + ("cache", "kwargs", "expected"), + ( + (None, {}, False), + (Cache(), {}, True), + (Cache(), {"cache": {"ttl": 60}}, True), + (Cache(), {"cache": {"no-cache": True}}, False), + ), +) +def test_response_cache_is_skipped_only_by_no_cache( + monkeypatch: pytest.MonkeyPatch, cache: Cache | None, kwargs: Mapping[str, object], expected: bool +) -> None: + monkeypatch.setattr(litellm, "cache", cache) + assert detectors.response_cache(call(**kwargs)) is expected + + +@pytest.mark.parametrize( + ("environment", "setting", "expected"), + ((None, None, False), ("DEFAULT@SECLEVEL=1", None, True), (None, "DEFAULT@SECLEVEL=1", True)), +) +def test_tls_policy_reads_the_security_level_environment_then_the_setting( + monkeypatch: pytest.MonkeyPatch, environment: str | None, setting: str | None, expected: bool +) -> None: + monkeypatch.delenv("SSL_SECURITY_LEVEL", raising=False) + monkeypatch.delenv("SSL_ECDH_CURVE", raising=False) + monkeypatch.setattr(litellm, "ssl_ecdh_curve", None) + if environment is not None: + monkeypatch.setenv("SSL_SECURITY_LEVEL", environment) + monkeypatch.setattr(litellm, "ssl_security_level", setting) + assert detectors.tls_policy(call()) is expected + + +@pytest.mark.parametrize( + ("environment", "setting", "expected"), + ((None, "X25519", False), (None, " P-256 ", False), ("secp521r1", None, True), (None, "secp521r1", True)), +) +def test_tls_policy_flags_only_curves_native_tls_lacks( + monkeypatch: pytest.MonkeyPatch, environment: str | None, setting: str | None, expected: bool +) -> None: + monkeypatch.delenv("SSL_SECURITY_LEVEL", raising=False) + monkeypatch.delenv("SSL_ECDH_CURVE", raising=False) + monkeypatch.setattr(litellm, "ssl_security_level", None) + if environment is not None: + monkeypatch.setenv("SSL_ECDH_CURVE", environment) + monkeypatch.setattr(litellm, "ssl_ecdh_curve", setting) + assert detectors.tls_policy(call()) is expected diff --git a/tests/unit/rust_bridge/test_public_call.py b/tests/unit/rust_bridge/test_public_call.py index 63029e288e2..5f50db6ea5a 100644 --- a/tests/unit/rust_bridge/test_public_call.py +++ b/tests/unit/rust_bridge/test_public_call.py @@ -3,7 +3,7 @@ from typing import Final import pytest -from litellm.rust_bridge.public_call import native_call, signature +from litellm.rust_bridge.public_call import AsyncFlag, native_call, signature def _messages( @@ -18,6 +18,7 @@ def _messages( _SIGNATURE: Final = signature(_messages) +_FLAG: Final = AsyncFlag("aflag", sync=_SIGNATURE) @pytest.mark.parametrize("supplied", ({}, {"api_key": None}, {"api_key": "explicit"})) @@ -62,3 +63,16 @@ def test_resolved_lays_the_keywords_over_the_base_and_equals_the_full_binding() } assert call.resolved["metadata"] is metadata assert supplied == {"model": "keyword-model", "metadata": metadata} + + +def test_flag_keywords_name_positionals_and_drop_only_the_flag() -> None: + messages: Final[Sequence[object]] = [{"role": "user", "content": "hello"}] + + keywords: Final = _FLAG.keywords((128, messages), {"model": "m", "stream": True, "aflag": True}) + + assert keywords == {"max_tokens": 128, "messages": messages, "model": "m", "stream": True} + + +def test_flag_keywords_reject_a_parameter_given_twice() -> None: + with pytest.raises(TypeError): + _FLAG.keywords((128,), {"max_tokens": 64}) diff --git a/tests/unit/test_circleci_path_filter.py b/tests/unit/test_circleci_path_filter.py index f94b2fa863a..c44560b96f1 100644 --- a/tests/unit/test_circleci_path_filter.py +++ b/tests/unit/test_circleci_path_filter.py @@ -96,7 +96,6 @@ CI = [".github/workflows/test-unit.yml"] ("ui", BACKEND + CLIENT, "run"), ("cost-map-only", ["model_prices_and_context_window.json"], "run"), ("cost-map-only", ["litellm/model_prices_and_context_window_backup.json"], "run"), - ("cost-map-only", ["model_prices_and_context_window.schema.json"], "run"), ( "cost-map-only", ["model_prices_and_context_window.json", "tests/unit/proxy/test_y.py"], diff --git a/tests/unit/test_constants.py b/tests/unit/test_constants.py index 0d0e77765b7..99b049f1803 100644 --- a/tests/unit/test_constants.py +++ b/tests/unit/test_constants.py @@ -1,6 +1,7 @@ import ast, asyncio, os import inspect import json +import sys from unittest import mock import httpx @@ -117,6 +118,33 @@ def test_cli_jwt_expiration_hours_from_environment( importlib.reload(litellm.constants) +@pytest.mark.parametrize( + ("raw", "expected"), + ( + (None, 64 * 1024 * 1024), + ("", 64 * 1024 * 1024), + ("invalid", 64 * 1024 * 1024), + ("0", 64 * 1024 * 1024), + ("-1", 64 * 1024 * 1024), + ("1", 1), + ("200000", 200000), + (str(sys.maxsize), sys.maxsize), + (str(sys.maxsize + 1), 64 * 1024 * 1024), + ), +) +def test_mcp_sse_event_limit_uses_positive_override_or_default(raw: str | None, expected: int) -> None: + try: + with pytest.MonkeyPatch.context() as environment: + if raw is None: + environment.delenv("LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", raising=False) + else: + environment.setenv("LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", raw) + importlib.reload(constants) + assert constants.MCP_MAX_SSE_EVENT_SIZE_BYTES == expected + finally: + importlib.reload(constants) + + def _is_blank(value: object) -> bool: return isinstance(value, str) and not value.strip() diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index a705597952b..578d5482822 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -186,13 +186,6 @@ def test_bot_may_only_touch_the_cost_map_files() -> None: ) -def test_bot_may_update_cost_map_and_legacy_schema() -> None: - head: Final = _snapshot({**BASE_MAP, "openrouter/c": _entry()}) - changed: Final = (*guard.GUARDED_PATHS, guard.LEGACY_SCHEMA_PATH) - - assert _failures(head, changed_files=changed) == () - - def test_bot_may_not_remove_models() -> None: head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}) assert _failures(head, bot=False) == () diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f5d13aecb6f..c3d3575b8b3 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -3868,6 +3868,9 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): model="gpt-3.5-turbo-instruct", choices=[{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}], ), + TextCompletionResponse( + id="cmpl-1", object="text_completion", created=1, model="gpt-3.5-turbo-instruct", choices=[] + ), ] response = stream_chunk_builder_text_completion(chunks=chunks, messages=[{"role": "user", "content": "say hello"}]) @@ -8853,7 +8856,8 @@ async def test_text_completion_stream_include_usage_returns_usage_on_final_chunk "stream": True, "stream_options": {"include_usage": True}, } - assert "".join(chunk.choices[0].text or "" for chunk in chunks) == "reply" + assert "".join(chunk.choices[0].text or "" for chunk in chunks if chunk.choices) == "reply" + assert chunks[-1].choices == [] final_usage: Final = chunks[-1].usage assert (final_usage.prompt_tokens, final_usage.completion_tokens, final_usage.total_tokens) == (3, 2, 5) diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 628870259a4..701b79d0ba9 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -272,6 +272,27 @@ def test_vertex_gemini_flash_rows_carry_the_documented_retirement_date(prices: d assert prices[f"{prefix}{model}"]["deprecation_date"] == VERTEX_GEMINI_FLASH_RETIREMENT_DATES[model] +def _sampling_flags_across_provider_twins(prices: Mapping[str, Mapping[str, object]], key: str) -> frozenset[object]: + bare: Final = key.split("/")[-1] + twins: Final = tuple(twin for twin in (bare, f"vertex_ai/{bare}", f"gemini/{bare}") if twin in prices) + return frozenset(prices[twin].get("supports_sampling_params") for twin in (key, *twins)) + + +def test_gemini_rows_agree_on_supports_sampling_params_across_their_provider_twins( + prices: Mapping[str, Mapping[str, object]], +) -> None: + """The bare, vertex_ai/ and gemini/ rows of one Gemini model gate sampling params identically, so a request + takes the same drop-or-400 path on the Vertex route and the AI Studio route""" + flagged: Final = tuple( + key + for key, row in prices.items() + if "supports_sampling_params" in row and key.split("/")[-1].startswith("gemini-") + ) + assert flagged + for key in flagged: + assert len(_sampling_flags_across_provider_twins(prices, key)) == 1, key + + OPENAI_REASONING_FAMILY_MARKERS = ("codex", "deep-research", "chat-latest") diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d3605fa22d2..0cb6f04b5b4 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -15321,7 +15321,7 @@ async def test_cyclic_fallback_graph_does_not_amplify_one_request(): _FLAT_ATTEMPT_RECORD_KEYS = frozenset( - {"model_group", "deployment_id", "exception_type", "exception_string", "attempted_retries"} + {"model_group", "deployment_id", "model", "base_model", "exception_type", "exception_string", "attempted_retries"} ) _BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip" @@ -29796,6 +29796,73 @@ def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: py assert [d["model_info"]["id"] for d in filtered] == ["only"] +@pytest.mark.parametrize( + ("deployment_stream_options", "default_stream_options", "request_stream_options", "usage_chunk_choice_count"), + [ + ({"include_usage": True}, None, None, 1), + (None, {"include_usage": True}, None, 1), + ({"include_usage": True}, None, {"include_usage": True}, 0), + (None, None, {"include_usage": True}, 0), + ], + ids=["deployment-config", "router-default", "request-over-deployment", "request-only"], +) +@pytest.mark.asyncio +async def test_router_configured_stream_usage_keeps_a_choice_on_the_usage_chunk( + deployment_stream_options: dict[str, bool] | None, + default_stream_options: dict[str, bool] | None, + request_stream_options: dict[str, bool] | None, + usage_chunk_choice_count: int, +): + deployment_params: Final = {"model": "openai/gpt-5-mini", "api_key": "fake", "mock_response": "hi there"} + router: Final = Router( + model_list=[ + { + "model_name": "m", + "litellm_params": ( + deployment_params + if deployment_stream_options is None + else {**deployment_params, "stream_options": deployment_stream_options} + ), + } + ], + default_litellm_params=None if default_stream_options is None else {"stream_options": default_stream_options}, + ) + request_kwargs: Final = {} if request_stream_options is None else {"stream_options": request_stream_options} + + response: Final = await router.acompletion( + model="m", messages=[{"role": "user", "content": "hi"}], stream=True, **request_kwargs + ) + chunks: Final = [chunk async for chunk in response] + + assert getattr(chunks[-1], "usage", None) is not None + assert len(chunks[-1].choices) == usage_chunk_choice_count + assert all(len(chunk.choices) == 1 for chunk in chunks[:-1]) + + +@pytest.mark.parametrize( + ("stream_options", "marked_injected"), + [({"include_usage": True}, True), (None, False)], + ids=["include-usage", "no-stream-options"], +) +def test_update_kwargs_with_deployment_reads_stream_options_from_pydantic_litellm_params( + stream_options: dict[str, bool] | None, marked_injected: bool +): + from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD + from litellm.types.router import LiteLLM_Params + + kwargs: Final[dict[str, object]] = {"stream": True, "metadata": {}} + Router()._update_kwargs_with_deployment( + deployment={ + "model_name": "gpt-5-mini", + "litellm_params": LiteLLM_Params(model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options), + "model_info": {"id": "pydantic-params"}, + }, + kwargs=kwargs, + ) + + assert (kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True) is marked_injected + + @pytest.mark.asyncio async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens(): """After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone.""" diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index e3061294bef..af8e7f53462 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1,8 +1,12 @@ +import asyncio import functools +import gc import json import os +import weakref from datetime import datetime, timezone -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -12,8 +16,12 @@ from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo +from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.types.litellm_params import RoutingStrategyName +from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( CallTypes, ModelResponse, @@ -837,7 +845,12 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> "model": "gpt-3.5-turbo", "api_key": "sk-must-not-be-recorded", "messages": [{"role": "user", "content": "hi"}], - metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"}, + metadata_key: { + "model_info": {"id": "deployment-1", "base_model": "azure/gpt-3.5-turbo"}, + "deployment": "azure/my-35-deployment", + "attempted_retries": 2, + "user_api_key": "sk-proxy", + }, }, e=rate_limit_error, ) @@ -845,6 +858,8 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> { "model_group": "gpt-3.5-turbo", "deployment_id": "deployment-1", + "model": "azure/my-35-deployment", + "base_model": "azure/gpt-3.5-turbo", "exception_type": "RateLimitError", "exception_string": "litellm.RateLimitError: slow down", "attempted_retries": 2, @@ -856,6 +871,38 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1 +def test_log_retry_reads_base_model_from_the_deployments_litellm_params() -> None: + """An Azure named deployment declares ``base_model`` under ``litellm_params``, not ``model_info``, and + that is the model the failure row is priced at, so the retry record has to carry it too""" + router = Router( + model_list=[ + { + "model_name": "azure-group", + "litellm_params": { + "model": "azure/my-named-deployment", + "api_key": "sk-fake", + "api_base": "https://example.openai.azure.com", + "base_model": "azure/gpt-4o-mini", + }, + "model_info": {"id": "dep-azure-named"}, + } + ] + ) + new_kwargs = router.log_retry( + kwargs={ + "model": "azure-group", + "metadata": { + "model_info": {"id": "dep-azure-named"}, + "deployment": "azure/my-named-deployment", + "attempted_retries": 0, + }, + }, + e=litellm.Timeout(message="slow", llm_provider="azure", model="azure/my-named-deployment"), + ) + record = new_kwargs["metadata"]["previous_models"][0] + assert (record["model"], record["base_model"]) == ("azure/my-named-deployment", "azure/gpt-4o-mini") + + @pytest.mark.usefixtures("router_minute_pinned") def test_update_usage(model_list): """Test if the '_update_usage' function is working correctly""" @@ -1813,6 +1860,315 @@ def test_discard(model_list): assert len(litellm.service_callback) == 0 +_CALLBACK_LIST_NAMES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +def _callback_ids(list_name: str) -> list[int]: + return [id(callback) for callback in getattr(litellm, list_name)] + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("least-busy", "leastbusy_logger"), + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_releases_strategy_selector_across_router_lifecycles( + routing_strategy: RoutingStrategyName, selector_attr: str +): + def retire() -> weakref.ReferenceType[object]: + router: Final = Router(model_list=[], routing_strategy=routing_strategy) + selector: Final = getattr(router, selector_attr) + assert any(callback is selector for callback in litellm.callbacks) + router.discard() + return weakref.ref(selector) + + selector_refs: Final = tuple(retire() for _ in range(3)) + await asyncio.sleep(0) + gc.collect() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + assert all(selector_ref() is None for selector_ref in selector_refs) + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("least-busy", "leastbusy_logger"), + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_purges_selector_from_event_lists_after_serving_a_request( + routing_strategy: RoutingStrategyName, selector_attr: str +): + router: Final = Router( + model_list=[{"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}], + routing_strategy=routing_strategy, + ) + selector: Final = getattr(router, selector_attr) + await router.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + assert any(callback is selector for callback in litellm.success_callback) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def test_discard_keeps_callbacks_owned_by_other_live_routers(): + survivor: Final = Router(model_list=[], routing_strategy="latency-based-routing") + unrelated: Final = CustomLogger() + litellm.logging_callback_manager.add_litellm_callback(unrelated) + before_ids: Final = {list_name: _callback_ids(list_name) for list_name in _CALLBACK_LIST_NAMES} + retired: Final = Router(model_list=[], routing_strategy="least-busy") + assert any(callback is retired.leastbusy_logger for callback in litellm.callbacks) + + retired.discard() + + assert {list_name: _callback_ids(list_name) for list_name in before_ids} == before_ids + assert any(callback is survivor.lowestlatency_logger for callback in litellm.callbacks) + survivor.discard() + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_leaves_same_strategy_survivor_receiving_events( + routing_strategy: RoutingStrategyName, selector_attr: str +): + model_list: Final[list[DeploymentTypedDict]] = [ + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} + ] + retired: Final = Router(model_list=model_list, routing_strategy=routing_strategy) + survivor: Final = Router(model_list=model_list, routing_strategy=routing_strategy) + + retired.discard() + await survivor.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + + assert any(callback is getattr(survivor, selector_attr) for callback in litellm.success_callback) + survivor.discard() + + +def test_discard_drops_routing_group_selectors(): + router: Final = Router( + model_list=[ + {"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} + ], + routing_groups=[ + {"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "latency-based-routing"} + ], + ) + _, group_selector = router._get_routing_context("grouped-model") + assert any(callback is group_selector for callback in litellm.callbacks) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +_GROUPED_MODEL_LIST: Final[list[DeploymentTypedDict]] = [ + {"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} +] + + +class _LabelledHook(CustomLogger): + def __init__(self, label: str) -> None: + super().__init__() + self.label = label + + +def _visit_hooks_reloading_routing_groups_mid_loop( + router: Router, hooks: tuple[_LabelledHook, ...] +) -> Iterator[CustomLogger]: + for callback in litellm.callbacks: + if callback is hooks[0]: + router.update_settings( + routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}] + ) + if callback in hooks: + yield callback + + +def test_routing_group_reload_keeps_every_hook_of_an_in_flight_callback_loop(): + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}], + ) + hooks: Final = tuple(_LabelledHook(label) for label in ("guardrail_a", "guardrail_b", "guardrail_c")) + for hook in hooks: + litellm.logging_callback_manager.add_litellm_callback(hook) + + visited: Final = tuple(_visit_hooks_reloading_routing_groups_mid_loop(router, hooks)) + + assert visited == hooks + router.discard() + + +def _routers_seen_in_success_loop_discarding_first(first: Router, second: Router) -> Iterator[Router]: + for callback in litellm.success_callback: + if not isinstance(callback, MethodType): + continue + if callback.__self__ is first: + first.discard() + yield first + if callback.__self__ is second: + yield second + + +def test_discard_mid_success_loop_keeps_the_next_router_callback(): + first: Final = Router(model_list=_GROUPED_MODEL_LIST) + second: Final = Router(model_list=_GROUPED_MODEL_LIST) + + seen: Final = tuple(_routers_seen_in_success_loop_discarding_first(first, second)) + + assert seen == (first, second) + second.discard() + + +@pytest.mark.asyncio +async def test_discard_purges_optional_pre_call_checks_from_event_lists_after_serving_a_request(): + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + optional_pre_call_checks=["enforce_model_rate_limits", "prompt_caching", "deployment_affinity"], + ) + await router.acompletion(model="grouped-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def test_discard_registers_same_config_survivor_optional_pre_call_checks(): + retired: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"]) + survivor: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"]) + + retired.discard() + + assert survivor.optional_callbacks + assert all(_is_registered(check) for check in survivor.optional_callbacks) + survivor.discard() + + +def _is_registered(callback: object) -> bool: + return any(registered is callback for registered in litellm.callbacks) + + +@pytest.mark.asyncio +async def test_discard_releases_per_request_override_selector(): + def build_override_then_discard() -> weakref.ReferenceType[object]: + router: Final = Router(model_list=_GROUPED_MODEL_LIST) + selector: Final = router._get_override_strategy_selector("usage-based-routing-v2") + assert selector is not None + router.discard() + return weakref.ref(selector) + + selector_ref: Final = build_override_then_discard() + await asyncio.sleep(0) + gc.collect() + + assert selector_ref() is None + + + +def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list(): + retired: Final = Router(model_list=[], routing_strategy="latency-based-routing") + survivor: Final = Router(routing_strategy="latency-based-routing") + assert not _is_registered(survivor.lowestlatency_logger) + + retired.discard() + + assert _is_registered(survivor.lowestlatency_logger) + survivor.discard() + + +def test_discard_drops_its_slack_alerting_callbacks(): + router: Final = Router(model_list=_GROUPED_MODEL_LIST, alerting_config=AlertingConfig(webhook_url="test")) + assert any(isinstance(callback, SlackAlerting) for callback in litellm.callbacks) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def _adaptive_router_model_list(name: str) -> list[DeploymentTypedDict]: + return [ + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}, + { + "model_name": name, + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["gpt-5.5"]}, + }, + }, + ] + + +def _registered_adaptive_router_strategies() -> tuple[object, ...]: + return tuple( + callback.adaptive_router for callback in litellm.callbacks if isinstance(callback, AdaptiveRouterPostCallHook) + ) + + +def test_discard_drops_its_adaptive_router_hook(): + router: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + assert _registered_adaptive_router_strategies() + + router.discard() + + assert _registered_adaptive_router_strategies() == () + + +def test_discard_registers_the_adaptive_router_hook_of_a_survivor(): + survivor: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + retired: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + + retired.discard() + + survivor_strategies: Final = tuple(tagged.strategy for tagged in survivor.adaptive_routers["adaptive"]) + assert _registered_adaptive_router_strategies() == survivor_strategies + survivor.discard() + + +@pytest.mark.asyncio +async def test_discard_releases_its_router_budget_limiter(): + def build_budget_limiter_then_discard() -> weakref.ReferenceType[object]: + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + budget_limiter: Final = router.get_router_deployment_budget_limiter() + assert budget_limiter is not None + router.discard() + return weakref.ref(budget_limiter) + + budget_limiter_ref: Final = build_budget_limiter_then_discard() + await asyncio.sleep(0) + gc.collect() + + assert budget_limiter_ref() is None + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list) diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 426481d2719..f001008ee2c 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -585,7 +585,7 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) -@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill"]) +@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params"]) def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: """A deployment's capability override must not land on the shared ``{provider}/{model}`` key, or every sibling deployment of that model would inherit it.""" diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 24d3ed678c9..cdff03af3b4 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1073,6 +1073,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "image_edit", "embedding", "evaluation", + "extraction", "guardrail", "image_generation", "video_generation", @@ -3890,6 +3891,114 @@ class TestCallbackAsyncSyncSeparation: assert my_sync_cb not in litellm._async_input_callback +class _DropsACallbackFromAnotherThreadWhenComparedToDynamodb: + def __init__(self, dropped: Callable[..., object]) -> None: + self.dropped = dropped + self.drop_thread: threading.Thread | None = None + + def __eq__(self, other: object) -> bool: + if isinstance(other, str) and other == "dynamodb" and self.drop_thread is None: + self.drop_thread = threading.Thread( + target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),) + ) + self.drop_thread.start() + self.drop_thread.join(timeout=0.2) + return self is other + + __hash__ = object.__hash__ + + +def test_function_setup_moves_async_callbacks_while_another_thread_drops_a_callback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def dropped_callback(*args: object, **kwargs: object) -> None: + return None + + async def async_callback(*args: object, **kwargs: object) -> None: + return None + + drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToDynamodb(dropped_callback) + monkeypatch.setattr(litellm.utils, "callback_list", ["initialized"]) + for list_name in ( + "callbacks", + "input_callback", + "failure_callback", + "service_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, []) + monkeypatch.setattr(litellm, "success_callback", [dropped_callback, drops_on_compare, async_callback]) + + litellm.utils.function_setup( + "completion", + litellm.utils.Rules(), + datetime.now(), + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="callback-drop-race", + ) + assert drops_on_compare.drop_thread is not None + drops_on_compare.drop_thread.join() + + assert litellm.success_callback == [drops_on_compare] + assert litellm._async_success_callback == [async_callback] + + +class _DropsACallbackFromAnotherThreadWhenComparedToSentry: + def __init__(self, dropped: Callable[..., object]) -> None: + self.dropped = dropped + self.drop_thread: threading.Thread | None = None + self.drop_finished_during_logger_setup = False + + def __eq__(self, other: object) -> bool: + if isinstance(other, str) and other == "sentry" and self.drop_thread is None: + self.drop_thread = threading.Thread( + target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),) + ) + self.drop_thread.start() + self.drop_thread.join(timeout=5) + self.drop_finished_during_logger_setup = not self.drop_thread.is_alive() + return self is other + + __hash__ = object.__hash__ + + +def test_function_setup_lets_other_threads_drop_callbacks_while_it_sets_up_loggers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def dropped_callback(*args: object, **kwargs: object) -> None: + return None + + drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToSentry(dropped_callback) + monkeypatch.setattr(litellm.utils, "callback_list", []) + for list_name in ( + "callbacks", + "input_callback", + "failure_callback", + "service_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, []) + monkeypatch.setattr(litellm, "success_callback", [drops_on_compare]) + + litellm.utils.function_setup( + "completion", + litellm.utils.Rules(), + datetime.now(), + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="logger-setup-outside-lock", + ) + assert drops_on_compare.drop_thread is not None + drops_on_compare.drop_thread.join() + + assert drops_on_compare.drop_finished_during_logger_setup + + class TestMetadataNoneHandling: """ Test that metadata=None in kwargs doesn't cause TypeError. @@ -9378,7 +9487,7 @@ def test_models_by_provider(): continue elif v["litellm_provider"] == "sagemaker" or v["litellm_provider"] == "bedrock_converse": continue - elif v.get("mode") in ("search", "decisions", "evaluation"): + elif v.get("mode") in ("search", "decisions", "evaluation", "extraction"): continue else: providers.add(v["litellm_provider"]) diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 84ba85f4a68..5a2cefd08b1 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -91,6 +91,7 @@ CONNECTION_NAMES: Final = ( "azure_ad_token_provider", "litellm_credential_name", "github_copilot_auth_type", + "github_copilot_host", "configurable_clientside_auth_params", "use_xai_oauth", "fireworks_forward_user_id", @@ -262,6 +263,7 @@ INTERNAL_STATE_NAMES: Final = ( "litellm_trusted_callback_vars", "_litellm_addressed_response_id", "_litellm_strip_stream_usage", + "_litellm_stream_usage_injected", "client_side_timeout", "model_file_id_mapping", "acompletion", diff --git a/tests/unit/types/test_router.py b/tests/unit/types/test_router.py index e15e95f9d55..2892a829eab 100644 --- a/tests/unit/types/test_router.py +++ b/tests/unit/types/test_router.py @@ -361,7 +361,7 @@ def test_server_owned_registry_includes_anthropic_openai_and_oauth_token_exchang "token_exchange_profile", "token_exchange_scope", ) - assert github_copilot_oauth_litellm_params == ("github_copilot_auth_type",) + assert github_copilot_oauth_litellm_params == ("github_copilot_auth_type", "github_copilot_host") assert server_owned_wif_litellm_params == ( anthropic_wif_litellm_params + openai_wif_litellm_params diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 423287e7c0f..4802aa057c6 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -73,21 +73,6 @@ "count": 1 } }, - "src/app/(dashboard)/budgets/_components/budget_modal.tsx": { - "local/filename-pascal-case": { - "count": 1 - } - }, - "src/app/(dashboard)/budgets/_components/budget_panel.tsx": { - "local/filename-pascal-case": { - "count": 1 - } - }, - "src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx": { - "local/filename-pascal-case": { - "count": 1 - } - }, "src/app/(dashboard)/caching/_components/cache_dashboard.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 7cc7415d52d..944455721cd 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -86,7 +86,7 @@ "tw-animate-css": "1.4.0", "typescript": "5.9.3", "typescript-eslint": "8.60.1", - "vite": "7.3.5", + "vite": "7.3.7", "vitest": "4.1.11" }, "engines": { @@ -12422,13 +12422,13 @@ } }, "node_modules/vite": { - "version": "7.3.5", - "resolved": "https://registry.npmjs.org/vite/-/vite-7.3.5.tgz", - "integrity": "sha512-KuOaNhcnGFN2zIPGA7wRmzF+lJA1sea7rHq17aiJ++9lzY1WWG6Jpwqwe1KNbRVPIqHmr8GLYx7jbrQcN/7/ww==", + "version": "7.3.7", + "resolved": "https://registry.npmjs.org/vite/-/vite-7.3.7.tgz", + "integrity": "sha512-AHEzTORMbi7/7aXp7w7KZ5+DCUdSeLsJz0grY8xHpc9DHdYb0vQu6DqJfm2tOw3gDExNgPgnJTXxM6syYFHaCA==", "dev": true, "license": "MIT", "dependencies": { - "esbuild": "^0.27.0", + "esbuild": "^0.27.0 || ^0.28.0", "fdir": "^6.5.0", "picomatch": "^4.0.3", "postcss": "^8.5.6", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index cd08995471f..f2a1a5e4e7f 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -103,7 +103,7 @@ "tw-animate-css": "1.4.0", "typescript": "5.9.3", "typescript-eslint": "8.60.1", - "vite": "7.3.5", + "vite": "7.3.7", "vitest": "4.1.11" }, "overrides": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx new file mode 100644 index 00000000000..9ffd23884b4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetFormFields.tsx @@ -0,0 +1,154 @@ +import React, { useId } from "react"; +import { Controller } from "react-hook-form"; +import { z } from "zod"; +import { FormField } from "@/components/shared/form/FormField"; +import { Badge } from "@/components/ui/badge"; +import { ButtonGroup } from "@/components/ui/button-group"; +import { + Field, + FieldDescription, + FieldGroup, + FieldLabel, + FieldLegend, + FieldSeparator, + FieldSet, +} from "@/components/ui/field"; +import { Input } from "@/components/ui/input"; +import { InputGroup, InputGroupAddon, InputGroupInput, InputGroupText } from "@/components/ui/input-group"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { useZodForm } from "@/lib/forms/useZodForm"; + +const budgetShape = { + budget_id: z.string().min(1, "Enter a budget ID"), + tpm_limit: z.number().nullish(), + rpm_limit: z.number().nullish(), + tpd_limit: z.number().nullish(), + max_budget: z.number().nullish(), + budget_duration: z.string().nullish(), +}; + +const budgetSchema = z.object(budgetShape); + +export type BudgetFormValues = z.output; + +export const useBudgetForm = (defaultValues: z.input) => + useZodForm(budgetSchema, { defaultValues }); + +type BudgetForm = ReturnType; +type NumberFieldName = "tpm_limit" | "rpm_limit" | "tpd_limit" | "max_budget"; + +const NO_RESET = "never"; + +const RESET_OPTIONS = [ + { value: NO_RESET, label: "in total" }, + { value: "24h", label: "per day" }, + { value: "7d", label: "per week" }, + { value: "30d", label: "per month" }, +]; + +const numberProps = void }>({ + value, + onChange, + ...field +}: T) => ({ + ...field, + type: "number", + value: value ?? "", + placeholder: "No limit", + onChange: (event: React.ChangeEvent) => + onChange(event.target.value === "" ? null : event.target.valueAsNumber), +}); + +const OptionalBadge = () => Optional; + +function SpendLimitField({ form }: { form: BudgetForm }) { + const amountId = useId(); + return ( + + + Spend limit + + + + ( + + + $ + + + + )} + /> + ( + + )} + /> + + + ); +} + +function RateLimitField({ form, name, label }: { form: BudgetForm; name: NumberFieldName; label: string }) { + return ( + + {(control) => } + + ); +} + +export function BudgetFormFields({ form, isEditing }: { form: BudgetForm; isEditing: boolean }) { + return ( + + + {({ ref, ...field }) => } + + + +
+ + Rate limits + + + + + + + Batches count against the daily limit instead of the per-minute ones. Provider rate limits still apply. + +
+ + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.integration.test.tsx new file mode 100644 index 00000000000..0770150ab91 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.integration.test.tsx @@ -0,0 +1,92 @@ +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import BudgetModal from "./BudgetModal"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; + +const { createMock } = vi.hoisted(() => ({ createMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/budgets/useBudgets", () => ({ + useCreateBudget: () => ({ mutateAsync: createMock }), +})); + +const renderModal = () => render(); + +const create = async (user: ReturnType) => + user.click(screen.getByRole("button", { name: "Create budget" })); + +const resetPeriod = () => screen.getByRole("combobox", { name: "Reset period" }); + +describe("BudgetModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + createMock.mockResolvedValue(undefined); + }); + + it("submits the spend limit and reset period alongside the rate limits", async () => { + const user = userEvent.setup(); + renderModal(); + + expect(screen.getByLabelText(/^Spend limit/)).toBeVisible(); + expect(resetPeriod()).toBeVisible(); + fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } }); + fireEvent.change(screen.getByLabelText(/^Spend limit/), { target: { value: "42.567" } }); + await chooseSelectOption(user, resetPeriod(), "per month"); + fireEvent.change(screen.getByLabelText("Tokens per minute"), { target: { value: "500" } }); + fireEvent.change(screen.getByLabelText("Requests per minute"), { target: { value: "7" } }); + fireEvent.change(screen.getByLabelText("Batch tokens per day"), { target: { value: "90000" } }); + await create(user); + + await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); + const expected = { + budget_id: "budget-alpha", + max_budget: 42.57, + budget_duration: "30d", + tpm_limit: 500, + rpm_limit: 7, + tpd_limit: 90000, + }; + expect(createMock.mock.calls[0][0]).toEqual(expected); + }); + + it("submits a spend limit that never resets as a null reset period", async () => { + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } }); + fireEvent.change(screen.getByLabelText(/^Spend limit/), { target: { value: "10" } }); + expect(resetPeriod()).toHaveTextContent("in total"); + await chooseSelectOption(user, resetPeriod(), "per week"); + await chooseSelectOption(user, resetPeriod(), "in total"); + await create(user); + + await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); + expect(createMock.mock.calls[0][0]).toEqual({ budget_id: "budget-alpha", max_budget: 10, budget_duration: null }); + }); + + it("submits a cleared number field as null", async () => { + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } }); + fireEvent.change(screen.getByLabelText("Tokens per minute"), { target: { value: "5" } }); + await user.clear(screen.getByLabelText("Tokens per minute")); + await create(user); + + await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); + expect(createMock.mock.calls[0][0]).toEqual({ budget_id: "budget-alpha", tpm_limit: null }); + }); + + it("blocks submit while Budget ID is empty", async () => { + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("Tokens per minute"), { target: { value: "5" } }); + await create(user); + + await waitFor(() => expect(screen.getByLabelText("Budget ID")).toHaveAttribute("aria-invalid", "true")); + expect(createMock).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx new file mode 100644 index 00000000000..bd10837640d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetModal.tsx @@ -0,0 +1,56 @@ +import React, { useId } from "react"; +import { useCreateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { applyBudgetPrecision } from "./budgetPrecision"; +import { BudgetFormFields, type BudgetFormValues, useBudgetForm } from "./BudgetFormFields"; +import { toast } from "@/lib/toast"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogClose, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; + +interface BudgetModalProps { + isModalVisible: boolean; + setIsModalVisible: React.Dispatch>; +} +const BudgetModal: React.FC = ({ isModalVisible, setIsModalVisible }) => { + const formId = useId(); + const form = useBudgetForm({ budget_id: "" }); + const createBudget = useCreateBudget(); + + const handleCancel = () => { + setIsModalVisible(false); + form.reset(); + }; + + const handleCreate = async (formValues: BudgetFormValues) => { + try { + toast.info("Making API Call"); + await createBudget.mutateAsync(applyBudgetPrecision(formValues)); + toast.success("Budget Created"); + form.reset(); + setIsModalVisible(false); + } catch (error) { + console.error("Error creating the budget:", error); + toast.fromError(`Error creating the budget: ${error}`); + } + }; + + return ( + !open && handleCancel()}> + + + Create budget + +
+ + + + }>Cancel + + + + + ); +}; + +export default BudgetModal; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.test.tsx index f2608c1221f..334bad8d0d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.test.tsx @@ -6,7 +6,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { ApiError } from "@/lib/http/client"; -import BudgetPanel from "./budget_panel"; +import BudgetPanel from "./BudgetPanel"; const { getMock, budgetDeleteMock } = vi.hoisted(() => ({ getMock: vi.fn(), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx similarity index 98% rename from ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx index 630d90e91b9..785b09252af 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetPanel.tsx @@ -17,9 +17,9 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import { toast } from "@/lib/toast"; import { useBudgetList, useDeleteBudget, budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; -import BudgetModal from "./budget_modal"; +import BudgetModal from "./BudgetModal"; import BudgetTable from "./BudgetTable"; -import EditBudgetModal from "./edit_budget_modal"; +import EditBudgetModal from "./EditBudgetModal"; import { CREATE_END_USER_CURL_COMMAND, CHAT_COMPLETIONS_CURL_COMMAND, OPENAI_SDK_PYTHON_CODE } from "./constants"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { isProxyAdminRole } from "@/utils/roles"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.integration.test.tsx new file mode 100644 index 00000000000..ea84d57ac30 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.integration.test.tsx @@ -0,0 +1,111 @@ +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import type { components } from "@/lib/http/schema"; + +import EditBudgetModal from "./EditBudgetModal"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; + +const { updateMock } = vi.hoisted(() => ({ updateMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/budgets/useBudgets", () => ({ + useUpdateBudget: () => ({ mutateAsync: updateMock }), +})); + +type BudgetItem = components["schemas"]["BudgetListItem"]; + +const EXISTING_BUDGET: BudgetItem = { + budget_id: "budget-alpha", + max_budget: 12.345, + budget_duration: "7d", + tpm_limit: 1000, + rpm_limit: 10, + tpd_limit: 50000, + soft_budget: 25, + budget_reset_at: "2026-02-01T00:00:00Z", + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-02T00:00:00Z", +}; + +const renderModal = () => + render(); + +const save = async (user: ReturnType) => + user.click(screen.getByRole("button", { name: "Save" })); + +const resetPeriod = () => screen.getByRole("combobox", { name: "Reset period" }); + +describe("EditBudgetModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + updateMock.mockResolvedValue(undefined); + }); + + it("prefills every field from the existing budget", () => { + renderModal(); + + expect(screen.getByLabelText(/^Spend limit/)).toHaveValue(12.345); + expect(resetPeriod()).toHaveTextContent("per week"); + expect(screen.getByLabelText("Tokens per minute")).toHaveValue(1000); + expect(screen.getByLabelText("Requests per minute")).toHaveValue(10); + expect(screen.getByLabelText("Batch tokens per day")).toHaveValue(50000); + }); + + it("sends only the changed rate limit so the stored spend limit and reset period are left as they are", async () => { + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("Tokens per minute"), { target: { value: "500" } }); + await save(user); + + await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); + expect(updateMock.mock.calls[0][0]).toEqual({ budget_id: "budget-alpha", tpm_limit: 500 }); + }); + + it("sends a later change after another field was changed and reverted", async () => { + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("Tokens per minute"), { target: { value: "500" } }); + fireEvent.change(screen.getByLabelText("Tokens per minute"), { target: { value: "1000" } }); + fireEvent.change(screen.getByLabelText("Requests per minute"), { target: { value: "20" } }); + await save(user); + + await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); + expect(updateMock.mock.calls[0][0]).toEqual({ budget_id: "budget-alpha", rpm_limit: 20 }); + }); + + it("submits a changed spend limit and reset period", async () => { + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText(/^Spend limit/), { target: { value: "42.567" } }); + await chooseSelectOption(user, resetPeriod(), "per month"); + await save(user); + + await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); + const expected = { budget_id: "budget-alpha", max_budget: 42.57, budget_duration: "30d" }; + expect(updateMock.mock.calls[0][0]).toEqual(expected); + }); + + it("submits a cleared limit and an 'in total' reset period as null", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.clear(screen.getByLabelText("Tokens per minute")); + await chooseSelectOption(user, resetPeriod(), "in total"); + await save(user); + + await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); + const expected = { budget_id: "budget-alpha", tpm_limit: null, budget_duration: null }; + expect(updateMock.mock.calls[0][0]).toEqual(expected); + }); + + it("does not let the budget ID be edited", () => { + renderModal(); + + expect(screen.getByLabelText("Budget ID")).toBeDisabled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx new file mode 100644 index 00000000000..4aa69c5c15e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/EditBudgetModal.tsx @@ -0,0 +1,74 @@ +import React, { useEffect, useId } from "react"; +import { useUpdateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { applyBudgetPrecision } from "./budgetPrecision"; +import { BudgetFormFields, type BudgetFormValues, useBudgetForm } from "./BudgetFormFields"; +import { pickDirty } from "@/lib/forms/pickDirty"; +import { toast } from "@/lib/toast"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogClose, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; + +const toFormValues = (budget: budgetItem) => ({ + budget_id: budget.budget_id, + tpm_limit: budget.tpm_limit, + rpm_limit: budget.rpm_limit, + tpd_limit: budget.tpd_limit, + max_budget: budget.max_budget, + budget_duration: budget.budget_duration, +}); + +interface EditBudgetModalProps { + isModalVisible: boolean; + setIsModalVisible: React.Dispatch>; + existingBudget: budgetItem; +} +const EditBudgetModal: React.FC = ({ isModalVisible, setIsModalVisible, existingBudget }) => { + const formId = useId(); + const form = useBudgetForm(toFormValues(existingBudget)); + const { dirtyFields } = form.formState; + const updateBudget = useUpdateBudget(); + + useEffect(() => { + form.reset(toFormValues(existingBudget)); + }, [existingBudget, form]); + + const handleCancel = () => { + setIsModalVisible(false); + form.reset(); + }; + + const handleUpdate = async (formValues: BudgetFormValues) => { + try { + toast.info("Making API Call"); + const changedValues = pickDirty(formValues, dirtyFields); + await updateBudget.mutateAsync(applyBudgetPrecision({ budget_id: formValues.budget_id, ...changedValues })); + toast.success("Budget Updated"); + form.reset(); + setIsModalVisible(false); + } catch (error) { + console.error("Error updating the budget:", error); + toast.fromError(`Error updating the budget: ${error}`); + } + }; + + return ( + !open && handleCancel()}> + + + Edit budget + +
+ + + + }>Cancel + + + + + ); +}; + +export default EditBudgetModal; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx deleted file mode 100644 index 17a297d203d..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx +++ /dev/null @@ -1,138 +0,0 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import React from "react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; - -import BudgetModal from "./budget_modal"; -import { chooseSelectOption } from "../../../../../tests/test-utils"; - -const { createMock } = vi.hoisted(() => ({ createMock: vi.fn() })); - -vi.mock("@/app/(dashboard)/hooks/budgets/useBudgets", () => ({ - useCreateBudget: () => ({ mutateAsync: createMock }), -})); - -const FULL_PAYLOAD = { - budget_id: "budget-alpha", - tpm_limit: 500.57, - rpm_limit: 7, - max_budget: 42.57, - budget_duration: "30d", -}; - -const renderModal = () => render(); - -const create = async (user: ReturnType) => - user.click(screen.getByRole("button", { name: "Create Budget" })); - -const openOptionalSettings = async (user: ReturnType) => { - await user.click(screen.getByText("Optional Settings")); - await screen.findByLabelText("Max Budget (USD)"); -}; - -describe("BudgetModal", () => { - beforeEach(() => { - vi.clearAllMocks(); - createMock.mockResolvedValue(undefined); - }); - - it("submits only the mounted fields when Optional Settings stays collapsed", async () => { - const user = userEvent.setup(); - renderModal(); - - fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } }); - fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "500.567" } }); - fireEvent.change(screen.getByLabelText("Max Requests per minute"), { target: { value: "7" } }); - await create(user); - - await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); - expect(createMock.mock.calls[0][0]).toEqual({ - budget_id: "budget-alpha", - tpm_limit: 500.57, - rpm_limit: 7, - }); - }); - - it("submits every field once Optional Settings is expanded", async () => { - const user = userEvent.setup(); - renderModal(); - - fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } }); - fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "500.567" } }); - fireEvent.change(screen.getByLabelText("Max Requests per minute"), { target: { value: "7" } }); - - await openOptionalSettings(user); - fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - - await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); - - await create(user); - - await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); - expect(createMock.mock.calls[0][0]).toEqual(FULL_PAYLOAD); - }); - - it("drops Optional Settings values again when the section is collapsed before submit", async () => { - const user = userEvent.setup(); - renderModal(); - - fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } }); - - await openOptionalSettings(user); - fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); - - await user.click(screen.getByText("Optional Settings")); - await waitFor(() => expect(screen.queryByLabelText("Max Budget (USD)")).not.toBeInTheDocument()); - await create(user); - - await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); - expect(createMock.mock.calls[0][0]).toEqual({ budget_id: "budget-alpha" }); - }); - - it("submits a cleared number field as null", async () => { - const user = userEvent.setup(); - renderModal(); - - fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } }); - fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "5" } }); - await user.clear(screen.getByLabelText("Max Tokens per minute")); - await create(user); - - await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); - expect(createMock.mock.calls[0][0]).toEqual({ - budget_id: "budget-alpha", - tpm_limit: null, - }); - }); - - it("blocks submit while Budget ID is empty", async () => { - const user = userEvent.setup(); - renderModal(); - - fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "5" } }); - await create(user); - - await waitFor(() => expect(screen.getByLabelText("Budget ID")).toHaveAttribute("aria-invalid", "true")); - expect(createMock).not.toHaveBeenCalled(); - }); - - it("keeps a typed Optional Setting when the section is collapsed and reopened, as antd's store did", async () => { - const user = userEvent.setup(); - renderModal(); - fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "probe-budget" } }); - - await openOptionalSettings(user); - fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.5" } }); - - await user.click(screen.getByText("Optional Settings")); - await user.click(screen.getByText("Optional Settings")); - - expect(await screen.findByLabelText("Max Budget (USD)")).toHaveValue(42.5); - - await create(user); - - await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); - expect(createMock.mock.calls[0][0]).toMatchObject({ budget_id: "probe-budget", max_budget: 42.5 }); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx deleted file mode 100644 index 01f1595b365..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx +++ /dev/null @@ -1,181 +0,0 @@ -import { ChevronRight } from "lucide-react"; -import React from "react"; -import { z } from "zod"; -import { useCreateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; -import { applyBudgetPrecision } from "./budgetPrecision"; -import { toast } from "@/lib/toast"; -import { FieldGroup } from "@/components/ui/field"; -import { FormField } from "@/components/shared/form/FormField"; -import { Button } from "@/components/ui/button"; -import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; -import { Input } from "@/components/ui/input"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { useZodForm } from "@/lib/forms/useZodForm"; -import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; - -const budgetShape = { - budget_id: z.string().min(1, "Please input a human-friendly name for the budget"), - tpm_limit: z.number().nullish(), - rpm_limit: z.number().nullish(), - tpd_limit: z.number().nullish(), - max_budget: z.number().nullish(), - budget_duration: z.string().nullish(), -}; - -const budgetSchema = z.object(budgetShape); - -type BudgetFormValues = z.output; - -const BUDGET_DURATION_OPTIONS = [ - { value: "24h", label: "daily" }, - { value: "7d", label: "weekly" }, - { value: "30d", label: "monthly" }, -]; - -interface BudgetModalProps { - isModalVisible: boolean; - setIsModalVisible: React.Dispatch>; -} -const BudgetModal: React.FC = ({ isModalVisible, setIsModalVisible }) => { - const [optionalSettingsOpen, setOptionalSettingsOpen] = React.useState(false); - const form = useZodForm(budgetSchema, { defaultValues: { budget_id: "" } }); - const createBudget = useCreateBudget(); - - const handleCancel = () => { - setIsModalVisible(false); - form.reset(); - }; - - const handleCreate = async (formValues: BudgetFormValues) => { - try { - toast.info("Making API Call"); - await createBudget.mutateAsync( - applyBudgetPrecision( - optionalSettingsOpen ? formValues : { ...formValues, max_budget: undefined, budget_duration: undefined }, - ), - ); - toast.success("Budget Created"); - form.reset(); - setIsModalVisible(false); - } catch (error) { - console.error("Error creating the budget:", error); - toast.fromError(`Error creating the budget: ${error}`); - } - }; - - return ( - !open && handleCancel()}> - - - Create Budget - -
- - - {({ ref, ...field }) => } - - - {({ ref, value, onChange, ...field }) => ( - onChange(event.target.value === "" ? null : event.target.valueAsNumber)} - /> - )} - - - {({ ref, value, onChange, ...field }) => ( - onChange(event.target.value === "" ? null : event.target.valueAsNumber)} - /> - )} - - - {({ ref, value, onChange, ...field }) => ( - onChange(event.target.value === "" ? null : event.target.valueAsNumber)} - /> - )} - - - - - Optional Settings - - - - - {({ ref, value, onChange, ...field }) => ( - onChange(event.target.value === "" ? null : event.target.valueAsNumber)} - /> - )} - - - {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( - - )} - - - - - -
- -
- - - - ); -}; - -export default BudgetModal; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx deleted file mode 100644 index fe0ecfc10dd..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx +++ /dev/null @@ -1,112 +0,0 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import React from "react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; - -import type { components } from "@/lib/http/schema"; - -import EditBudgetModal from "./edit_budget_modal"; -import { chooseSelectOption } from "../../../../../tests/test-utils"; - -const { updateMock } = vi.hoisted(() => ({ updateMock: vi.fn() })); - -vi.mock("@/app/(dashboard)/hooks/budgets/useBudgets", () => ({ - useUpdateBudget: () => ({ mutateAsync: updateMock }), -})); - -type BudgetItem = components["schemas"]["BudgetListItem"]; - -const EXISTING_BUDGET: BudgetItem = { - budget_id: "budget-alpha", - max_budget: 100, - budget_duration: "7d", - tpm_limit: 1000, - rpm_limit: 10, - soft_budget: 25, - budget_reset_at: "2026-02-01T00:00:00Z", - created_at: "2026-01-01T00:00:00Z", - updated_at: "2026-01-02T00:00:00Z", -}; - -const renderModal = () => - render(); - -const save = async (user: ReturnType) => - user.click(screen.getByRole("button", { name: "Save" })); - -const openOptionalSettings = async (user: ReturnType) => { - await user.click(screen.getByText("Optional Settings")); - await screen.findByLabelText("Max Budget (USD)"); -}; - -describe("EditBudgetModal", () => { - beforeEach(() => { - vi.clearAllMocks(); - updateMock.mockResolvedValue(undefined); - }); - - it("submits only the mounted fields when Optional Settings stays collapsed", async () => { - const user = userEvent.setup(); - renderModal(); - - await user.clear(screen.getByLabelText("Max Tokens per minute")); - fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "500.567" } }); - await save(user); - - await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); - expect(updateMock.mock.calls[0][0]).toEqual({ - budget_id: "budget-alpha", - tpm_limit: 500.57, - rpm_limit: 10, - }); - }); - - it("submits every field once Optional Settings is expanded", async () => { - const user = userEvent.setup(); - renderModal(); - - await user.clear(screen.getByLabelText("Max Tokens per minute")); - fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "500.567" } }); - await user.clear(screen.getByLabelText("Max Requests per minute")); - fireEvent.change(screen.getByLabelText("Max Requests per minute"), { target: { value: "7" } }); - - await openOptionalSettings(user); - await user.clear(screen.getByLabelText("Max Budget (USD)")); - fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - - await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); - - await save(user); - - await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); - const expected = { - budget_id: "budget-alpha", - tpm_limit: 500.57, - rpm_limit: 7, - max_budget: 42.57, - budget_duration: "30d", - }; - - expect(updateMock.mock.calls[0][0]).toEqual(expected); - }); - - it("keeps a typed Optional Setting when the section is collapsed and reopened, as antd's store did", async () => { - const user = userEvent.setup(); - renderModal(); - - await openOptionalSettings(user); - const maxBudget = screen.getByLabelText("Max Budget (USD)"); - await user.clear(maxBudget); - fireEvent.change(maxBudget, { target: { value: "99.25" } }); - - await user.click(screen.getByText("Optional Settings")); - await user.click(screen.getByText("Optional Settings")); - - expect(await screen.findByLabelText("Max Budget (USD)")).toHaveValue(99.25); - - await save(user); - - await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); - expect(updateMock.mock.calls[0][0]).toMatchObject({ max_budget: 99.25 }); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx deleted file mode 100644 index 1931a88f096..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx +++ /dev/null @@ -1,187 +0,0 @@ -import { ChevronRight } from "lucide-react"; -import React, { useEffect } from "react"; -import { useForm } from "react-hook-form"; -import { useUpdateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; -import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; -import { applyBudgetPrecision } from "./budgetPrecision"; -import { toast } from "@/lib/toast"; -import { FieldGroup } from "@/components/ui/field"; -import { FormField } from "@/components/shared/form/FormField"; -import { Button } from "@/components/ui/button"; -import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; -import { Input } from "@/components/ui/input"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; - -type EditBudgetFormValues = Pick< - budgetItem, - "budget_id" | "tpm_limit" | "rpm_limit" | "tpd_limit" | "max_budget" | "budget_duration" ->; - -const toFormValues = (budget: budgetItem): EditBudgetFormValues => ({ - budget_id: budget.budget_id, - tpm_limit: budget.tpm_limit, - rpm_limit: budget.rpm_limit, - tpd_limit: budget.tpd_limit, - max_budget: budget.max_budget, - budget_duration: budget.budget_duration, -}); - -const BUDGET_DURATION_OPTIONS = [ - { value: "24h", label: "daily" }, - { value: "7d", label: "weekly" }, - { value: "30d", label: "monthly" }, -]; - -interface EditBudgetModalProps { - isModalVisible: boolean; - setIsModalVisible: React.Dispatch>; - existingBudget: budgetItem; -} -const EditBudgetModal: React.FC = ({ isModalVisible, setIsModalVisible, existingBudget }) => { - const [optionalSettingsOpen, setOptionalSettingsOpen] = React.useState(false); - const form = useForm({ defaultValues: toFormValues(existingBudget) }); - const updateBudget = useUpdateBudget(); - - useEffect(() => { - form.reset(toFormValues(existingBudget)); - }, [existingBudget, form]); - - const handleCancel = () => { - setIsModalVisible(false); - form.reset(); - }; - - const handleUpdate = async (formValues: EditBudgetFormValues) => { - try { - toast.info("Making API Call"); - await updateBudget.mutateAsync( - applyBudgetPrecision( - optionalSettingsOpen ? formValues : { ...formValues, max_budget: undefined, budget_duration: undefined }, - ), - ); - toast.success("Budget Updated"); - form.reset(); - setIsModalVisible(false); - } catch (error) { - console.error("Error updating the budget:", error); - toast.fromError(`Error updating the budget: ${error}`); - } - }; - - return ( - !open && handleCancel()}> - - - Edit Budget - -
- - - {({ ref, ...field }) => } - - - {({ ref, value, onChange, ...field }) => ( - onChange(event.target.value === "" ? null : event.target.valueAsNumber)} - /> - )} - - - {({ ref, value, onChange, ...field }) => ( - onChange(event.target.value === "" ? null : event.target.valueAsNumber)} - /> - )} - - - {({ ref, value, onChange, ...field }) => ( - onChange(event.target.value === "" ? null : event.target.valueAsNumber)} - /> - )} - - - - - Optional Settings - - - - - {({ ref, value, onChange, ...field }) => ( - onChange(event.target.value === "" ? null : event.target.valueAsNumber)} - /> - )} - - - {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( - - )} - - - - - -
- -
- - - - ); -}; - -export default EditBudgetModal; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/page.tsx index ca34589a679..7cac0788f73 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/page.tsx @@ -1,6 +1,6 @@ "use client"; -import BudgetPanel from "./_components/budget_panel"; +import BudgetPanel from "./_components/BudgetPanel"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; export default function Budgets() { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx index 3447602dd79..ece32449041 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx @@ -277,4 +277,25 @@ describe("CacheDashboard cache analytics charts", () => { expect(compactTicks(requestsCard).length).toBeGreaterThan(0); expect(compactTicks(tokensCard)).toContain("60K"); }); + + it("shows the failed state with the server message and Retry when the activity query errors", async () => { + const refetch = vi.fn(); + const failedActivity = { + data: undefined, + isError: true, + error: { error: { message: "Request timed out after 30s" } }, + refetch, + }; + useCacheActivity.mockReturnValue(failedActivity); + renderDashboard(); + + await screen.findByText(REQUESTS_CHART_TITLE); + const card = cardTitled(REQUESTS_CHART_TITLE); + const alert = within(card).getByRole("alert"); + expect(alert).toHaveTextContent("Failed to load cache activity"); + expect(alert).toHaveTextContent("Request timed out after 30s"); + expect(within(card).getByText(/Click a red failed-requests segment/)).toBeInTheDocument(); + fireEvent.click(within(alert).getByRole("button", { name: "Retry" })); + expect(refetch).toHaveBeenCalledTimes(1); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx index 29819adf524..9e730be4df3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx @@ -21,6 +21,8 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { RefreshCw } from "lucide-react"; import { cachingHealthCheckCall } from "@/components/networking"; +import { deriveErrorMessage } from "@/lib/http/client"; +import { ErrorState } from "@/components/shared/ErrorState"; import { useCacheActivity, type CacheActivityGroup } from "@/app/(dashboard)/hooks/caching/useCacheActivity"; // Import the new component @@ -94,7 +96,12 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole const [lastRefreshed, setLastRefreshed] = useState(""); const [healthCheckResponse, setHealthCheckResponse] = useState(""); - const { data: activity, refetch } = useCacheActivity({ + const { + data: activity, + isError: activityFailed, + error: activityError, + refetch, + } = useCacheActivity({ startDate: formatDateWithoutTZ(dateValue.from), endDate: formatDateWithoutTZ(dateValue.to), keyAliases: selectedApiKeys, @@ -295,19 +302,27 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole Click a red failed-requests segment to see which error codes caused those failures.

{hasUnknownGroup &&

{UNKNOWN_CALL_TYPE_NOTE}

} - { - if (item.categoryClicked === REQUEST_SERIES.failed) setErrorDrilldownCallType(item.name); - }} - /> + {activityFailed ? ( + void refetch()} + /> + ) : ( + { + if (item.categoryClicked === REQUEST_SERIES.failed) setErrorDrilldownCallType(item.name); + }} + /> + )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx new file mode 100644 index 00000000000..0d74b714ee6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx @@ -0,0 +1,78 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { GuardrailDeleteModal } from "./GuardrailDeleteModal"; +import { deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils"; + +vi.mock("@/components/networking", () => ({ deleteGuardrailCall: vi.fn() })); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + +const GUARDRAIL = { + guardrail_id: "gr-1", + guardrail_name: "word-block", + litellm_params: { guardrail: "litellm_content_filter", mode: "pre_call", default_on: true }, +}; + +const renderModal = () => { + const onClose = vi.fn(); + const onDeleted = vi.fn(); + renderWithProviders( + , + ); + return { onClose, onDeleted }; +}; + +const confirmDelete = () => { + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { + target: { value: GUARDRAIL.guardrail_name }, + }); + fireEvent.click(screen.getByRole("button", { name: "Delete" })); +}; + +describe("GuardrailDeleteModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("warns about name references and only enables Delete once the exact name is typed", () => { + renderModal(); + + const warning = screen.getByText(/silently stop running it/); + expect(warning).toHaveTextContent("Keys, teams, and policies that list this guardrail"); + expect(warning).toHaveTextContent("Policy pipeline steps that use it will block every request with a 400 error"); + expect(warning).toHaveTextContent("unless the step's on_error is next or allow"); + const confirmButton = screen.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { target: { value: "word" } }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { + target: { value: GUARDRAIL.guardrail_name }, + }); + expect(confirmButton).toBeEnabled(); + }); + + it("shows a success toast naming the guardrail and calls onDeleted after a successful delete", async () => { + vi.mocked(deleteGuardrailCall).mockResolvedValue({}); + const { onClose, onDeleted } = renderModal(); + + confirmDelete(); + + await waitFor(() => expect(onDeleted).toHaveBeenCalledTimes(1)); + expect(deleteGuardrailCall).toHaveBeenCalledWith("tok", GUARDRAIL.guardrail_id); + expect(toast.success).toHaveBeenCalledWith('Guardrail "word-block" deleted successfully'); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(onClose).toHaveBeenCalledTimes(1); + }); + + it("shows an error toast and does not call onDeleted when the delete fails", async () => { + vi.mocked(deleteGuardrailCall).mockRejectedValue(new Error("boom")); + const { onClose, onDeleted } = renderModal(); + + confirmDelete(); + + await waitFor(() => expect(onClose).toHaveBeenCalledTimes(1)); + expect(toast.fromError).toHaveBeenCalledWith("Failed to delete guardrail"); + expect(toast.success).not.toHaveBeenCalled(); + expect(onDeleted).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx new file mode 100644 index 00000000000..a3ed15b630f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx @@ -0,0 +1,68 @@ +import React, { useState } from "react"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import { deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; + +const DANGLING_REFERENCE_WARNING = + "Keys, teams, and policies that list this guardrail will silently stop running it. Policy pipeline steps that use it will block every request with a 400 error, unless the step's on_error is next or allow (or its on_fail, when on_error is not set)."; + +export interface DeletableGuardrail { + guardrail_id: string; + guardrail_name?: string | null; + litellm_params?: { + guardrail?: string; + mode?: unknown; + default_on?: boolean | null; + } | null; +} + +interface GuardrailDeleteModalProps { + guardrail: DeletableGuardrail | null; + accessToken: string | null; + onClose: () => void; + onDeleted: () => void; +} + +export function GuardrailDeleteModal({ guardrail, accessToken, onClose, onDeleted }: GuardrailDeleteModalProps) { + const [isDeleting, setIsDeleting] = useState(false); + const guardrailName = guardrail ? guardrail.guardrail_name || guardrail.guardrail_id : ""; + const providerDisplayName = getGuardrailLogoAndName(guardrail?.litellm_params?.guardrail || "").displayName; + + const handleDelete = async () => { + if (!guardrail || !accessToken) return; + setIsDeleting(true); + try { + await deleteGuardrailCall(accessToken, guardrail.guardrail_id); + toast.success(`Guardrail "${guardrailName}" deleted successfully`); + onClose(); + onDeleted(); + } catch { + toast.fromError("Failed to delete guardrail"); + onClose(); + } finally { + setIsDeleting(false); + } + }; + + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx new file mode 100644 index 00000000000..46bc09e7383 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx @@ -0,0 +1,51 @@ +import { Code, Trash2 } from "lucide-react"; +import React, { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { DeletableGuardrail, GuardrailDeleteModal } from "./GuardrailDeleteModal"; + +interface GuardrailSettingsActionsProps { + guardrail: DeletableGuardrail & { guardrail_definition_location?: string }; + isEditing: boolean; + onEdit: () => void; + onEditCode: () => void; + accessToken: string | null; + onDeleted: () => void; +} + +export function GuardrailSettingsActions({ + guardrail, + isEditing, + onEdit, + onEditCode, + accessToken, + onDeleted, +}: GuardrailSettingsActionsProps) { + const [isModalOpen, setIsModalOpen] = useState(false); + + if (isEditing || guardrail.guardrail_definition_location === "config") return null; + + return ( + <> + {guardrail.litellm_params?.guardrail === "custom_code" ? ( + + ) : ( + + )} + + setIsModalOpen(false)} + onDeleted={onDeleted} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx new file mode 100644 index 00000000000..19eb838da6b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx @@ -0,0 +1,135 @@ +import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import GuardrailsPanel from "./GuardrailsPanel"; +import { + deleteGuardrailCall, + getGuardrailInfo, + getGuardrailProviderSpecificParams, + getGuardrailsList, + getGuardrailUISettings, +} from "@/components/networking"; +import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils"; + +vi.mock("@/components/networking", () => ({ + getGuardrailsList: vi.fn(), + deleteGuardrailCall: vi.fn(), + getGuardrailInfo: vi.fn(), + getGuardrailUISettings: vi.fn(), + getGuardrailProviderSpecificParams: vi.fn(), + updateGuardrailCall: vi.fn(), +})); + +vi.mock("./add_guardrail_form", () => ({ __esModule: true, default: () => null })); +vi.mock("./guardrail_garden", () => ({ __esModule: true, default: () => null })); +vi.mock("./GuardrailTestPlayground", () => ({ __esModule: true, default: () => null })); +vi.mock("./TeamGuardrailsTab", () => ({ TeamGuardrailsTab: () => null })); + +const DB_GUARDRAIL = { + guardrail_id: "db-guardrail-1", + guardrail_name: "db-word-block", + litellm_params: { guardrail: "litellm_content_filter", mode: "pre_call", default_on: true }, + guardrail_info: null, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "db", +}; + +const CONFIG_GUARDRAIL = { + ...DB_GUARDRAIL, + guardrail_id: "config-guardrail-1", + guardrail_name: "cfg-word-block", + guardrail_definition_location: "config", +}; + +beforeAll(() => { + Object.defineProperty(window, "matchMedia", { + writable: true, + value: vi.fn().mockImplementation((query: string) => ({ + matches: false, + media: query, + onchange: null, + addListener: vi.fn(), + removeListener: vi.fn(), + addEventListener: vi.fn(), + removeEventListener: vi.fn(), + dispatchEvent: vi.fn(), + })), + }); +}); + +const UI_SETTINGS = { + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], +}; + +describe("GuardrailsPanel delete from the guardrail detail Settings card", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [DB_GUARDRAIL, CONFIG_GUARDRAIL] }); + vi.mocked(getGuardrailInfo).mockImplementation(async (_token: string, id: string) => + id === CONFIG_GUARDRAIL.guardrail_id ? CONFIG_GUARDRAIL : DB_GUARDRAIL, + ); + vi.mocked(getGuardrailUISettings).mockResolvedValue(UI_SETTINGS); + vi.mocked(getGuardrailProviderSpecificParams).mockResolvedValue({}); + }); + + const renderDetail = (guardrailId: string) => + renderWithProviders(, { + searchParams: `?guardrail=${guardrailId}`, + }); + + it("deletes a DB guardrail only after the admin types its name, then returns to the refreshed table", async () => { + vi.mocked(deleteGuardrailCall).mockResolvedValue({}); + renderDetail(DB_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + + const modal = within(await screen.findByRole("dialog")); + expect( + modal.getByText(/pipeline steps that use it will block every request with a 400 error/i), + ).toBeInTheDocument(); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + + const confirmationInput = modal.getByPlaceholderText(DB_GUARDRAIL.guardrail_name); + fireEvent.change(confirmationInput, { target: { value: "db-word" } }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(confirmationInput, { target: { value: DB_GUARDRAIL.guardrail_name } }); + expect(confirmButton).toBeEnabled(); + + fireEvent.click(confirmButton); + + await waitFor(() => expect(deleteGuardrailCall).toHaveBeenCalledWith("test-token", DB_GUARDRAIL.guardrail_id)); + expect(await screen.findByRole("table")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + expect(getGuardrailsList).toHaveBeenCalledTimes(2); + }); + + it("keeps the admin on the detail page when the delete request fails", async () => { + vi.mocked(deleteGuardrailCall).mockRejectedValue(new Error("boom")); + renderDetail(DB_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + const modal = within(await screen.findByRole("dialog")); + fireEvent.change(modal.getByPlaceholderText(DB_GUARDRAIL.guardrail_name), { + target: { value: DB_GUARDRAIL.guardrail_name }, + }); + fireEvent.click(modal.getByRole("button", { name: "Delete" })); + + await waitFor(() => expect(deleteGuardrailCall).toHaveBeenCalledTimes(1)); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.getByRole("button", { name: "Delete" })).toBeInTheDocument(); + expect(screen.queryByRole("table")).not.toBeInTheDocument(); + }); + + it("does not offer Delete on a config guardrail's detail page", async () => { + renderDetail(CONFIG_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + expect(await screen.findByRole("img", { name: "Config guardrail details" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Delete" })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx index 7be120a4cb7..ac824f63be9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx @@ -2,6 +2,7 @@ import { type UrlUpdateEvent } from "nuqs/adapters/testing"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import GuardrailsPanel from "./GuardrailsPanel"; import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils"; vi.mock("@/components/networking", () => ({ @@ -9,6 +10,8 @@ vi.mock("@/components/networking", () => ({ deleteGuardrailCall: vi.fn(), })); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + vi.mock("./add_guardrail_form", () => ({ __esModule: true, default: () =>
Mock Add Guardrail Form
, @@ -140,13 +143,36 @@ describe("GuardrailsPanel", () => { expect(modal.getByText("Delete Guardrail")).toBeInTheDocument(); expect(modal.getByText("test-guardrail-1")).toBeInTheDocument(); expect(modal.getByText("Test Provider")).toBeInTheDocument(); + expect( + modal.getByText(/pipeline steps that use it will block every request with a 400 error/i), + ).toBeInTheDocument(); - fireEvent.click(modal.getByRole("button", { name: "Delete" })); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(confirmButton); await waitFor(() => { expect(mockDeleteGuardrailCall).toHaveBeenCalledWith("test-token", "test-guardrail-1"); }); - expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2); + await waitFor(() => expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2)); + expect(toast.success).toHaveBeenCalledWith('Guardrail "Test Guardrail" deleted successfully'); + }); + + it("should show an error toast and not refetch when the row delete fails", async () => { + mockDeleteGuardrailCall.mockRejectedValue(new Error("boom")); + renderWithProviders(); + fireEvent.click(screen.getByText("Guardrails")); + + fireEvent.click(await screen.findByTestId("delete-button")); + const modal = within(await screen.findByRole("dialog")); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(modal.getByRole("button", { name: "Delete" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(toast.fromError).toHaveBeenCalledWith("Failed to delete guardrail"); + expect(toast.success).not.toHaveBeenCalled(); + expect(mockGetGuardrailsList).toHaveBeenCalledTimes(1); }); it("should mount every tab panel up front so panel state survives tab switches", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index 901e39004f3..b53d016e950 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -2,7 +2,7 @@ import { parseAsString, useQueryState } from "nuqs"; import React, { useState, useEffect } from "react"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { ChevronDown, Code, Plus } from "lucide-react"; -import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking"; +import { getGuardrailsList } from "@/components/networking"; import { buttonVariants } from "@/components/ui/button"; import { DropdownMenu, @@ -16,10 +16,8 @@ import GuardrailTable from "./guardrail_table"; import { isAdminRole } from "@/utils/roles"; import GuardrailInfoView from "./guardrail_info"; import GuardrailTestPlayground from "./GuardrailTestPlayground"; -import { toast } from "@/lib/toast"; import { Guardrail } from "@/components/guardrails/types"; -import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { GuardrailDeleteModal } from "./GuardrailDeleteModal"; import { CustomCodeModal } from "./custom_code"; import GuardrailGarden from "./guardrail_garden"; import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; @@ -38,9 +36,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isCustomCodeModalVisible, setIsCustomCodeModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(false); - const [isDeleting, setIsDeleting] = useState(false); const [guardrailToDelete, setGuardrailToDelete] = useState(null); - const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [selectedGuardrailId, setSelectedGuardrailId] = useQueryState( "guardrail", parseAsString.withOptions({ history: "push" }), @@ -97,40 +93,10 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole fetchGuardrails(); }; - const handleDeleteClick = (guardrailId: string, guardrailName: string) => { - const guardrail = guardrailsList.find((g) => g.guardrail_id === guardrailId) || null; - setGuardrailToDelete(guardrail); - setIsDeleteModalOpen(true); + const handleDeleteClick = (guardrailId: string) => { + setGuardrailToDelete(guardrailsList.find((g) => g.guardrail_id === guardrailId) || null); }; - const handleDeleteConfirm = async () => { - if (!guardrailToDelete || !accessToken) return; - - setIsDeleting(true); - try { - await deleteGuardrailCall(accessToken, guardrailToDelete.guardrail_id); - toast.success(`Guardrail "${guardrailToDelete.guardrail_name}" deleted successfully`); - await fetchGuardrails(); - } catch (error) { - console.error("Error deleting guardrail:", error); - toast.fromError("Failed to delete guardrail"); - } finally { - setIsDeleting(false); - setIsDeleteModalOpen(false); - setGuardrailToDelete(null); - } - }; - - const handleDeleteCancel = () => { - setIsDeleteModalOpen(false); - setGuardrailToDelete(null); - }; - - const providerDisplayName = - guardrailToDelete && guardrailToDelete.litellm_params - ? getGuardrailLogoAndName(guardrailToDelete.litellm_params.guardrail).displayName - : undefined; - return (
@@ -184,6 +150,10 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole { + closeGuardrailDetail(); + fetchGuardrails(); + }} accessToken={accessToken} isAdmin={isAdmin} /> @@ -210,24 +180,11 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole onSuccess={handleSuccess} /> - setGuardrailToDelete(null)} + onDeleted={fetchGuardrails} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx index 1689cf93289..732c50647b1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx @@ -84,7 +84,8 @@ const saveChanges = async (user: ReturnType) => { await user.click(screen.getByText("Save Changes")); }; -const renderView = () => render( {}} accessToken="123" isAdmin />); +const renderView = () => + render( {}} onDeleted={() => {}} accessToken="123" isAdmin />); const lastPayload = () => vi.mocked(networking.updateGuardrailCall).mock.calls.at(-1)?.[2]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 836921104ff..ae057689b37 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -10,6 +10,7 @@ vi.mock("@/components/networking", () => ({ getGuardrailUISettings: vi.fn(), getGuardrailProviderSpecificParams: vi.fn(), updateGuardrailCall: vi.fn(), + deleteGuardrailCall: vi.fn(), })); // Mock ContentFilterManager @@ -65,7 +66,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); // Wait for the loading to complete and data to be rendered await waitFor(() => { @@ -103,7 +106,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); expect(await screen.findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); }); @@ -131,7 +136,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); const logo = await screen.findByAltText("Presidio PII logo"); expect(logo).toHaveAttribute("src", expect.stringContaining("microsoft_azure.svg")); @@ -162,7 +169,7 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); const { container } = render( - {}} accessToken="123" isAdmin={true} />, + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, ); await waitFor(() => { @@ -212,7 +219,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); await waitFor(() => { expect(screen.getByText("PII Entity Configuration")).toBeInTheDocument(); @@ -245,7 +254,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); vi.mocked(networking.updateGuardrailCall).mockResolvedValue({ status: "success" }); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); await waitFor(() => { expect(screen.getByText("Settings")).toBeInTheDocument(); @@ -331,7 +342,9 @@ describe("Guardrail Info", () => { }); vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); expect(await screen.findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); @@ -355,10 +368,90 @@ describe("Guardrail Info when the guardrail cannot be loaded", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); const onClose = vi.fn(); - render(); + render( + {}} + accessToken="123" + isAdmin={true} + />, + ); expect(await screen.findByText("Guardrail not found")).toBeInTheDocument(); fireEvent.click(screen.getByRole("button", { name: /back to guardrails/i })); expect(onClose).toHaveBeenCalledTimes(1); }); + + describe("Delete from the Settings card", () => { + const dbGuardrail = { + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { guardrail: "presidio", mode: "pre_call", default_on: true }, + guardrail_definition_location: "db", + }; + const uiSettings = { + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call"], + }; + const mockDbGuardrail = () => { + vi.mocked(networking.getGuardrailInfo).mockResolvedValue(dbGuardrail); + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue(uiSettings); + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + }; + + const openDeleteModalAndConfirm = async () => { + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + const modal = within(await screen.findByRole("dialog")); + expect(modal.getByText(/silently stop running it/)).toBeInTheDocument(); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(confirmButton); + }; + + it("deletes the guardrail and calls onDeleted once the admin types its name", async () => { + mockDbGuardrail(); + vi.mocked(networking.deleteGuardrailCall).mockResolvedValue({}); + const onDeleted = vi.fn(); + render( + {}} onDeleted={onDeleted} accessToken="tok" isAdmin />, + ); + + await openDeleteModalAndConfirm(); + + await waitFor(() => expect(onDeleted).toHaveBeenCalledTimes(1)); + expect(networking.deleteGuardrailCall).toHaveBeenCalledWith("tok", "123"); + }); + + it("hides Delete while the settings are being edited", async () => { + mockDbGuardrail(); + render( {}} onDeleted={() => {}} accessToken="tok" isAdmin />); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Edit Settings" })); + + expect(await screen.findByRole("button", { name: "Save Changes" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Delete" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + }); + + it("does not call onDeleted when the delete request fails", async () => { + mockDbGuardrail(); + vi.mocked(networking.deleteGuardrailCall).mockRejectedValue(new Error("boom")); + const onDeleted = vi.fn(); + render( + {}} onDeleted={onDeleted} accessToken="tok" isAdmin />, + ); + + await openDeleteModalAndConfirm(); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(networking.deleteGuardrailCall).toHaveBeenCalledTimes(1); + expect(onDeleted).not.toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index 2cfe2b79886..eacacf5e30e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -33,6 +33,7 @@ import { } from "./GuardrailFormField"; import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal"; +import { GuardrailSettingsActions } from "./GuardrailSettingsActions"; import { GuardrailModeCard } from "./GuardrailModeDisplay"; import { GuardrailReadOnlyDetails } from "./GuardrailReadOnlyDetails"; import { @@ -72,11 +73,12 @@ const SectionHeading: React.FC<{ children: React.ReactNode }> = ({ children }) = export interface GuardrailInfoProps { guardrailId: string; onClose: () => void; + onDeleted: () => void; accessToken: string | null; isAdmin: boolean; } -const GuardrailInfoView: React.FC = ({ guardrailId, onClose, accessToken, isAdmin }) => { +const GuardrailInfoView: React.FC = ({ guardrailId, onClose, onDeleted, accessToken, isAdmin }) => { const [guardrailData, setGuardrailData] = useState(null); const [guardrailProviderSpecificParams, setGuardrailProviderSpecificParams] = useState(null); const [loading, setLoading] = useState(true); @@ -686,25 +688,21 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, {isAdmin && ( -
-

Guardrail Settings

+
+

Guardrail Settings

{isConfigGuardrail && ( )} - {!isEditing && - !isConfigGuardrail && - (guardrailData.litellm_params?.guardrail === "custom_code" ? ( - - ) : ( - - ))} + setIsEditing(true)} + onEditCode={() => setCustomCodeModalVisible(true)} + accessToken={accessToken} + onDeleted={onDeleted} + />
{isEditing ? ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index cdee86ed559..04b4dcf9be1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -1356,9 +1356,10 @@ describe("EntityUsage", () => { render(); - await waitFor(() => { - expect(screen.getAllByText(/Fetching spend data failed/).length).toBeGreaterThan(0); - }); + const alert = await screen.findByRole("alert"); + expect(alert).toHaveTextContent("Failed to load spend data"); + expect(alert).toHaveTextContent("aggregated unavailable"); + expect(within(alert).getByRole("button", { name: "Retry" })).toBeInTheDocument(); expect(mockDailyActivityAggregatedCall.mock.calls.filter((c) => c[0] === "team")).toHaveLength(1); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 8518f52e165..a75afd63633 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -14,7 +14,7 @@ import { hasCapability, type Capability } from "@/utils/capabilities"; import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import { Bot, Boxes, ChevronDown, ChevronRight, ExternalLink, Info, KeyRound, Layers, Server } from "lucide-react"; import type { ColumnDef } from "@tanstack/react-table"; -import { Alert, AlertDescription } from "@/components/shared/Alert"; +import { ErrorState } from "@/components/shared/ErrorState"; import { ChartLoader } from "@/components/shared/chart_loader"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; @@ -189,6 +189,8 @@ const EntityUsage: React.FC = ({ data: spendDataRaw, loading, failed, + errorMessage, + reload, } = useAggregatedDailyActivity({ fetch: () => api.aggregated(request as DailyActivityRequest), enabled: enabled && request !== null, @@ -208,6 +210,8 @@ const EntityUsage: React.FC = ({ data: agentSpendDataRaw, loading: agentLoading, failed: agentFailed, + errorMessage: agentErrorMessage, + reload: reloadAgent, } = useAggregatedDailyActivity({ fetch: () => ENTITY_API.agent.aggregated(agentRequest as DailyActivityRequest), enabled: enabled && showAgentBreakdown && agentRequest !== null, @@ -663,21 +667,9 @@ const EntityUsage: React.FC = ({ return (
- {failed && ( - - - Fetching spend data failed, so the totals below may be empty rather than final. Reload the page to try - again. - - - )} + {failed && } {showAgentBreakdown && agentFailed && ( - - - Fetching agent data failed, so the totals below may be empty rather than final. Reload the page to try - again. - - + )}
{ expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); }); expect(mockDailyActivityAggregatedCall.mock.calls.filter((c) => c[0] === "user")).toHaveLength(1); - expect(await screen.findByText(/Fetching spend data failed/)).toBeInTheDocument(); + const alert = await screen.findByRole("alert"); + expect(alert).toHaveTextContent("Failed to load spend data"); + expect(alert).toHaveTextContent("Aggregated endpoint not available"); + expect(within(alert).getByRole("button", { name: "Retry" })).toBeInTheDocument(); + }); + + it("re-issues the aggregated call when Retry is clicked", async () => { + mockUserDailyActivityAggregatedCall.mockRejectedValue(new Error("Aggregated endpoint not available")); + + renderWithProviders(); + + const alert = await screen.findByRole("alert"); + const callsBefore = mockUserDailyActivityAggregatedCall.mock.calls.length; + fireEvent.click(within(alert).getByRole("button", { name: "Retry" })); + + await waitFor(() => expect(mockUserDailyActivityAggregatedCall.mock.calls.length).toBe(callsBefore + 1)); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index 84cc9359ded..feb1adb0f1d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -13,6 +13,7 @@ import React, { useCallback, useEffect, useMemo, useRef, useState } from "react" import { StackedUsageChart } from "@/components/shared/charts"; import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { ErrorState } from "@/components/shared/ErrorState"; import { Button } from "@/components/ui/button"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -218,6 +219,8 @@ const UsagePage: React.FC = ({ teams, organizations }) => { data: aggregatedRaw, loading: aggregatedLoading, failed: aggregatedFailed, + errorMessage: aggregatedError, + reload: reloadAggregated, } = useAggregatedDailyActivity({ fetch: () => ENTITY_API.user.aggregated(dailyActivityRequest as DailyActivityRequest), enabled: dailyActivityRequest !== null, @@ -415,12 +418,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => {
{aggregatedFailed && ( - - - Fetching spend data failed, so the totals below may be empty rather than final. Reload the page to try - again. - - + )}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useAggregatedDailyActivity.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useAggregatedDailyActivity.test.ts index 16a45d03db9..7e667230cb0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useAggregatedDailyActivity.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useAggregatedDailyActivity.test.ts @@ -44,6 +44,19 @@ describe("useAggregatedDailyActivity", () => { await waitFor(() => expect(result.current.failed).toBe(true)); expect(result.current.loading).toBe(false); + expect(result.current.errorMessage).toBe("boom"); + }); + + it("reload fetches again after a failure and clears failed on success", async () => { + const fetch = vi.fn().mockRejectedValueOnce(new Error("boom")).mockResolvedValueOnce({ results: [], metadata: {} }); + const { result } = renderHook(() => useAggregatedDailyActivity({ fetch, enabled: true, deps: ["a"] })); + + await waitFor(() => expect(result.current.failed).toBe(true)); + act(() => result.current.reload()); + + await waitFor(() => expect(result.current.failed).toBe(false)); + expect(fetch).toHaveBeenCalledTimes(2); + expect(result.current.errorMessage).toBeNull(); }); it("ignores an out-of-order resolution from the previous deps", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useAggregatedDailyActivity.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useAggregatedDailyActivity.ts index 08aec2e77f2..482041c900d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useAggregatedDailyActivity.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/useAggregatedDailyActivity.ts @@ -14,16 +14,31 @@ interface Result { data: DailyActivityAggregatedResponse; loading: boolean; failed: boolean; + errorMessage: string | null; + reload: () => void; } interface SettledFetch { key: string; data: DailyActivityAggregatedResponse; failed: boolean; + errorMessage: string | null; } +const settledFetch = ( + key: string, + data: DailyActivityAggregatedResponse, + errorMessage: string | null, +): SettledFetch => ({ + key, + data, + failed: errorMessage !== null, + errorMessage, +}); + export function useAggregatedDailyActivity({ fetch, enabled, deps }: Options): Result { const [settled, setSettled] = useState(null); + const [attempt, setAttempt] = useState(0); const requestIdRef = useRef(0); const fetchRef = useRef(fetch); @@ -31,7 +46,7 @@ export function useAggregatedDailyActivity({ fetch, enabled, deps }: Options): R fetchRef.current = fetch; }); - const depsKey = JSON.stringify(deps); + const depsKey = JSON.stringify([attempt, ...deps]); useEffect(() => { if (!enabled) return; @@ -43,12 +58,13 @@ export function useAggregatedDailyActivity({ fetch, enabled, deps }: Options): R .current() .then((response) => { if (isStale()) return; - setSettled({ key: depsKey, data: response, failed: false }); + setSettled(settledFetch(depsKey, response, null)); }) .catch((error) => { if (isStale()) return; console.error("Error fetching daily activity:", error); - setSettled({ key: depsKey, data: EMPTY_DAILY_ACTIVITY_RESPONSE, failed: true }); + const message = error instanceof Error ? error.message : String(error); + setSettled(settledFetch(depsKey, EMPTY_DAILY_ACTIVITY_RESPONSE, message)); }); return () => { @@ -61,5 +77,7 @@ export function useAggregatedDailyActivity({ fetch, enabled, deps }: Options): R data: current?.data ?? EMPTY_DAILY_ACTIVITY_RESPONSE, loading: enabled && current === null, failed: current?.failed ?? false, + errorMessage: current?.errorMessage ?? null, + reload: () => setAttempt((n) => n + 1), }; } diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelCustomMetadata.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelCustomMetadata.test.tsx new file mode 100644 index 00000000000..49681154de4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/ModelCustomMetadata.test.tsx @@ -0,0 +1,29 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import { ModelCustomMetadata } from "./ModelCustomMetadata"; + +describe("ModelCustomMetadata", () => { + it("renders each tag and every metadata entry with its value", () => { + render( + , + ); + expect(screen.getByText("Tags:")).toBeInTheDocument(); + expect(screen.getByText("deprecated")).toBeInTheDocument(); + expect(screen.getByText("ptu")).toBeInTheDocument(); + const terms = screen.getAllByRole("term").map((term) => term.textContent); + const definitions = screen.getAllByRole("definition").map((definition) => definition.textContent); + expect(terms.map((term, index) => [term, definitions[index]])).toEqual([ + ["retirement_date", "2027-03-01"], + ["provisioned_ptu", "100"], + ["gsu", "false"], + ]); + }); + + it("renders nothing when the model has no tags or metadata", () => { + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelCustomMetadata.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelCustomMetadata.tsx new file mode 100644 index 00000000000..be2f509b646 --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/ModelCustomMetadata.tsx @@ -0,0 +1,48 @@ +import { Badge } from "@/components/ui/badge"; + +export type ModelCustomMetadataValue = string | number | boolean; + +interface ModelCustomMetadataProps { + customTags?: string[] | null; + customMetadata?: Record | null; +} + +export function ModelCustomTags({ tags, className = "" }: { tags: string[]; className?: string }) { + return ( +
+ {tags.map((tag) => ( + + {tag} + + ))} +
+ ); +} + +export function ModelCustomMetadata({ customTags, customMetadata }: ModelCustomMetadataProps) { + const tags = customTags ?? []; + const entries = Object.entries(customMetadata ?? {}); + return ( + <> + {tags.length > 0 && ( +
+

Tags:

+ +
+ )} + {entries.length > 0 && ( +
+

Custom Metadata:

+
+ {entries.map(([key, value]) => ( +
+
{key}
+
{String(value)}
+
+ ))} +
+
+ )} + + ); +} diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index fb665eb69c2..b5c2a3e3b05 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -54,6 +54,7 @@ import { isDecisionMode, } from "@/lib/decisionModels"; import DecisionModelGarden from "@/components/AIHub/DecisionModelGarden"; +import { ModelCustomMetadata } from "./ModelCustomMetadata"; interface ModelHubTableProps { accessToken: string | null; @@ -719,6 +720,10 @@ const ModelHubTable: React.FC = ({

{selectedModel.description}

)} +
diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx index dbe0371241a..f91eea4e301 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx @@ -61,6 +61,17 @@ describe("getModelHubTableColumns", () => { expect(screen.getByRole("cell", { name: "gpt-4o" })).not.toHaveClass("text-right"); }); + it("shows each custom tag as a badge and a dash when the model has none", () => { + renderTable([ + { ...mockModel, custom_tags: ["deprecated", "ptu"] }, + { ...mockModel, model_group: "untagged", custom_tags: null }, + ]); + expect(screen.getByRole("columnheader", { name: "Tags" })).toBeInTheDocument(); + const [taggedRow, untaggedRow] = screen.getAllByRole("row").slice(1); + expect(taggedRow).toHaveTextContent("deprecatedptu"); + expect(untaggedRow).not.toHaveTextContent("deprecated"); + }); + it("shows capability badges only for supported features", () => { renderTable([mockModel]); expect(screen.getByText("Vision")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx index 1a5f8889c93..e2dfdc67aab 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx @@ -17,6 +17,8 @@ import { cn } from "@/lib/cva.config"; import { displayMode } from "@/lib/decisionModels"; import { copyToClipboard } from "@/utils/dataUtils"; +import { ModelCustomTags, type ModelCustomMetadataValue } from "./ModelCustomMetadata"; + export interface ModelHubData { model_group: string; providers: string[]; @@ -33,6 +35,8 @@ export interface ModelHubData { supported_openai_params?: string[]; is_public_model_group: boolean; description?: string; + custom_tags?: string[] | null; + custom_metadata?: Record | null; [key: string]: any; } @@ -210,6 +214,20 @@ export const getModelHubTableColumns = ({ onModelClick }: ModelHubTableColumnsDe ); }, }, + { + id: "custom_tags", + meta: { title: "Tags", skeleton: "chips" }, + header: "Tags", + size: 160, + enableSorting: false, + cell: ({ row }) => { + const tags = row.original.custom_tags ?? []; + if (tags.length === 0) { + return -; + } + return ; + }, + }, { id: "is_public_model_group", accessorKey: "is_public_model_group", diff --git a/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.test.ts b/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.test.ts index 25d1b83ebe9..8591e6f82f8 100644 --- a/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.test.ts +++ b/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.test.ts @@ -21,6 +21,16 @@ describe("modelUsageExample", () => { ); }); + it("posts an extraction model to /v1/extraction with typed fields", () => { + const example = modelUsageExample("extraction", BASE_URL, "clairic", "X-Litellm-Key"); + expect(example).toContain('"https://proxy.example.com/v1/extraction"'); + expect(example).toContain('headers={"X-Litellm-Key": "Bearer your_api_key"}'); + expect(example).toContain('"model": "clairic"'); + expect(example).toContain('"fields": {'); + expect(example).not.toContain("chat.completions"); + expect(example).not.toContain("/v1/systemone"); + }); + it("keeps the chat completions example for every other mode", () => { for (const mode of ["chat", null, undefined]) { const example = modelUsageExample(mode, BASE_URL, "gpt-5.5", "Authorization"); @@ -28,6 +38,7 @@ describe("modelUsageExample", () => { expect(example).toContain('base_url="https://proxy.example.com"'); expect(example).toContain('model="gpt-5.5"'); expect(example).not.toContain("/v1/systemone"); + expect(example).not.toContain("/v1/extraction"); } }); diff --git a/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.ts b/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.ts index 2426ec09942..83f2af1e8c6 100644 --- a/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.ts +++ b/ui/litellm-dashboard/src/components/AIHub/modelUsageExample.ts @@ -38,10 +38,33 @@ response = requests.post( print(response.json()["answers"])`; +const extractionExample = (baseUrl: string, model: string, keyHeaderName: string): string => `import requests + +response = requests.post( + ${JSON.stringify(`${baseUrl}/v1/extraction`)}, # Your LiteLLM Proxy URL + headers={${JSON.stringify(keyHeaderName)}: "Bearer your_api_key"}, + json={ + "model": ${JSON.stringify(model)}, + "state": "Invoice #4471 from Acme Tools. Total $174.00, due in 30 days.", + "fields": { + "vendor": {"type": "text", "instructions": "Who sent the invoice?"}, + "total": {"type": "number", "unit": "USD"} + } + } +) + +print(response.json()["answers"])`; + export const modelUsageExample = ( mode: string | null | undefined, baseUrl: string, model: string, keyHeaderName: string, -): string => - isDecisionMode(mode) ? systemOneExample(baseUrl, model, keyHeaderName) : chatCompletionsExample(baseUrl, model); +): string => { + if (mode === "extraction") { + return extractionExample(baseUrl, model, keyHeaderName); + } + return isDecisionMode(mode) + ? systemOneExample(baseUrl, model, keyHeaderName) + : chatCompletionsExample(baseUrl, model); +}; diff --git a/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx b/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx index 0288c58e31e..94eafdea86c 100644 --- a/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx @@ -8,6 +8,7 @@ import { Badge } from "@/components/ui/badge"; import { getProviderLogoAndName } from "@/components/provider_info_helpers"; import { PUBLIC_MODEL_HUB_SORTABLE_FIELDS } from "@/components/publicModelHub/publicModelHubFilters"; import { displayMode } from "@/lib/decisionModels"; +import type { ModelCustomMetadataValue } from "@/components/AIHub/ModelCustomMetadata"; export interface ModelGroupInfo { model_group: string; @@ -27,6 +28,8 @@ export interface ModelGroupInfo { health_response_time?: number; health_checked_at?: string; description?: string; + custom_tags?: string[] | null; + custom_metadata?: Record | null; [key: string]: any; } @@ -260,6 +263,13 @@ export const getPublicModelHubColumns = ({ onModelClick }: PublicModelHubColumns return ; }, }, + { + id: "custom_tags", + meta: { title: "Tags", skeleton: "chips" }, + header: "Tags", + size: 140, + cell: ({ row }) => , + }, { id: "health_status", accessorKey: "health_status", diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx index cdc681a1c91..a95d9ffbfb8 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx @@ -465,7 +465,7 @@ describe("AddModelForm", () => { expect(await screen.findByRole("button", { name: "Add Model" })).toBeInTheDocument(); }); - it("offers the Decisions mode", async () => { + it.each(["Decisions - /v1/decisions", "Extraction - /v1/extraction"])("offers the %s mode", async (option) => { const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); @@ -474,7 +474,7 @@ describe("AddModelForm", () => { await screen.findByText("Provider"); await userEvent.click(screen.getByRole("combobox", { name: "Mode" })); - expect(await screen.findByRole("option", { name: "Decisions - /v1/decisions", exact: true })).toBeInTheDocument(); + expect(await screen.findByRole("option", { name: option, exact: true })).toBeInTheDocument(); }); it("shows only the Close button in the connection test dialog footer", async () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx index 45ffe969a21..63def76db8f 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx @@ -14,6 +14,7 @@ export const TEST_MODES = [ { value: "anthropic_messages", label: "Anthropic Messages - /v1/messages" }, { value: "ocr", label: "OCR - /ocr" }, { value: "decisions", label: "Decisions - /v1/decisions" }, + { value: "extraction", label: "Extraction - /v1/extraction" }, ]; // Define the available auto router routing strategies diff --git a/ui/litellm-dashboard/src/components/add_model/provider_auth_types.test.ts b/ui/litellm-dashboard/src/components/add_model/provider_auth_types.test.ts index 5467ff0085a..f0b739b877d 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_auth_types.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/provider_auth_types.test.ts @@ -28,6 +28,7 @@ describe("provider auth types", () => { "token_exchange_audience", "api_key", ]); + expect(authTypeFieldKeys("GITHUB_COPILOT")).toContain("github_copilot_host"); }); it("infers auth types from required fields and falls back when only optional defaults are present", () => { @@ -85,6 +86,9 @@ describe("provider auth types", () => { expect(inferAuthTypeId(authTypes, { api_key: "" })).toBe("shared_device_login"); expect(inferAuthTypeId(authTypes, { github_copilot_auth_type: "per_user_oauth" })).toBe("per_user_oauth"); expect(hiddenAuthFieldKeys(authTypes, "per_user_oauth")).toEqual(["api_base", "api_key"]); - expect(hiddenAuthFieldKeys(authTypes, "shared_device_login")).toEqual(["github_copilot_auth_type"]); + expect(hiddenAuthFieldKeys(authTypes, "shared_device_login")).toEqual([ + "github_copilot_auth_type", + "github_copilot_host", + ]); }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/provider_auth_types.ts b/ui/litellm-dashboard/src/components/add_model/provider_auth_types.ts index 6fe54955133..0bf98407bfa 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_auth_types.ts +++ b/ui/litellm-dashboard/src/components/add_model/provider_auth_types.ts @@ -13,6 +13,7 @@ export interface ProviderAuthType { const EMPTY_PROVIDER_AUTH_TYPES: readonly ProviderAuthType[] = []; export const GITHUB_COPILOT_AUTH_TYPE_KEY = "github_copilot_auth_type"; +export const GITHUB_COPILOT_HOST_KEY = "github_copilot_host"; export const GITHUB_COPILOT_PER_USER_AUTH_TYPE = "per_user_oauth"; export const PROVIDER_AUTH_TYPES: Partial> = { @@ -31,7 +32,7 @@ export const PROVIDER_AUTH_TYPES: Partial { expect(screen.queryByText("Refreshing every 15s")).not.toBeInTheDocument(); }); }); + + it("a failed logs request says so instead of showing an empty table", async () => { + vi.mocked(uiSpendLogsCall).mockRejectedValue(new Error("Request timed out after 30s")); + renderPanel(); + const alert = await screen.findByRole("alert", {}, { timeout: 3000 }); + expect(alert).toHaveTextContent("Failed to load"); + expect(alert).toHaveTextContent("Request timed out after 30s"); + expect(screen.getByRole("button", { name: "Retry" })).toBeEnabled(); + }); }); diff --git a/ui/litellm-dashboard/src/components/logs/request/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/logs/request/RequestLogsPanel.tsx index 605ff373c7c..e29c48b94c0 100644 --- a/ui/litellm-dashboard/src/components/logs/request/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/logs/request/RequestLogsPanel.tsx @@ -301,6 +301,8 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, data={rows} rowCount={rowCount} isLoading={logsQuery.isLoading} + isError={logsQuery.isError} + errorMessage={logsQuery.error instanceof Error ? logsQuery.error.message : undefined} isRefreshing={logsQuery.isFetching} pagination={pagination} onPaginationChange={handlePaginationChange} diff --git a/ui/litellm-dashboard/src/components/logs/request/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/logs/request/RequestLogsTable.tsx index 822bf73302b..d3d6b43e2e3 100644 --- a/ui/litellm-dashboard/src/components/logs/request/RequestLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/logs/request/RequestLogsTable.tsx @@ -18,6 +18,8 @@ interface RequestLogsTableProps { data: LogEntry[]; rowCount: number; isLoading: boolean; + isError?: boolean; + errorMessage?: ReactNode; isRefreshing: boolean; pagination: PaginationState; onPaginationChange: OnChangeFn; @@ -66,6 +68,8 @@ export function RequestLogsTable({ data, rowCount, isLoading, + isError = false, + errorMessage, isRefreshing, pagination, onPaginationChange, @@ -112,6 +116,9 @@ export function RequestLogsTable({ columnFilters={columnFilters} onColumnFiltersChange={onColumnFiltersChange} isLoading={isLoading} + isError={isError} + errorMessage={errorMessage} + onRetry={onRefresh} loadingMessage="Loading request logs…" noDataMessage={} size="compact" diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialModal.integration.test.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialModal.integration.test.tsx index e1e259063bf..c7d939333ec 100644 --- a/ui/litellm-dashboard/src/components/model_add/CredentialModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialModal.integration.test.tsx @@ -67,6 +67,13 @@ vi.mock("../networking", async () => { credential_fields: [ { key: "api_base", label: "API Base", field_type: "text" }, { key: "api_key", label: "API Key", field_type: "password" }, + { + key: "github_copilot_host", + label: "GitHub host", + placeholder: "github.com", + tooltip: "Leave empty for github.com. For GHE.com data residency enter your SUBDOMAIN.ghe.com host", + field_type: "text", + }, ], }, ]), @@ -532,7 +539,7 @@ describe("CredentialModal with Anthropic workload identity federation", () => { describe("CredentialModal with GitHub Copilot auth types", () => { const perUserCopilotCredential: CredentialItem = { credential_name: "copilot-per-user", - credential_values: { github_copilot_auth_type: "per_user_oauth" }, + credential_values: { github_copilot_auth_type: "per_user_oauth", github_copilot_host: "octo.ghe.com" }, credential_info: { custom_llm_provider: "GITHUB_COPILOT" }, }; @@ -577,6 +584,7 @@ describe("CredentialModal with GitHub Copilot auth types", () => { const onSubmit = renderModal({ mode: "edit", existingCredential: perUserCopilotCredential }); expect(await screen.findByRole("combobox", { name: "Auth Type:" })).toHaveTextContent("Per-user GitHub OAuth"); + expect(screen.getByLabelText("GitHub host")).toHaveValue("octo.ghe.com"); await chooseOption(user, /^Auth Type:/, "Shared device login"); await user.click(screen.getByRole("button", { name: "Update Credential" })); @@ -586,7 +594,7 @@ describe("CredentialModal with GitHub Copilot auth types", () => { custom_llm_provider: "GITHUB_COPILOT", credential_values: {}, }); - expect(valuesToDelete).toEqual(["github_copilot_auth_type"]); + expect([...valuesToDelete].sort()).toEqual(["github_copilot_auth_type", "github_copilot_host"]); }); }); diff --git a/ui/litellm-dashboard/src/components/model_add/GithubCopilotConnectModal.integration.test.tsx b/ui/litellm-dashboard/src/components/model_add/GithubCopilotConnectModal.integration.test.tsx index 963ebab81d3..42f30db08a8 100644 --- a/ui/litellm-dashboard/src/components/model_add/GithubCopilotConnectModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/model_add/GithubCopilotConnectModal.integration.test.tsx @@ -79,6 +79,18 @@ describe("GithubCopilotConnectModal", () => { expect(props.onDisconnect).toHaveBeenCalledTimes(1); }); + it("uses the GHE verification URI returned by the start response", async () => { + const verificationUri = "https://octo.ghe.com/login/device"; + networking.userConnectionStartCall.mockResolvedValue({ + ...START_RESPONSE, + verification_uri: verificationUri, + }); + renderModal(); + await advance(0); + + expect(screen.getByRole("link", { name: verificationUri })).toHaveAttribute("href", verificationUri); + }); + it("stops polling and offers a restart when the device code expires", async () => { networking.userConnectionStartCall.mockResolvedValue({ ...START_RESPONSE, expires_in: 12 }); networking.userConnectionPollCall.mockResolvedValue({ status: "pending" }); diff --git a/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.test.ts b/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.test.ts index 72c2e783adf..9645839bbb4 100644 --- a/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.test.ts +++ b/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.test.ts @@ -30,11 +30,29 @@ describe("connectFlowReducer", () => { expect(started()).toEqual(expected); }); - it("never links anywhere but github.com", () => { + it("falls back for an untrusted verification host", () => { const state = started("https://evil.example/login/device"); expect(state.kind === "awaiting" && state.verificationUri).toBe("https://github.com/login/device"); }); + it("preserves a single-label GHE.com verification URL", () => { + const uri = "https://octo.ghe.com/login/device"; + const state = started(uri); + expect(state.kind === "awaiting" && state.verificationUri).toBe(uri); + }); + + it.each([ + "http://octo.ghe.com/login/device", + "https://a.b.ghe.com/login/device", + "https://octo.ghe.com.evil.example/login/device", + "https://user@octo.ghe.com/login/device", + "https://octo.ghe.com:8443/login/device", + "not a url", + ])("falls back for unsafe verification URL %s", (uri) => { + const state = started(uri); + expect(state.kind === "awaiting" && state.verificationUri).toBe("https://github.com/login/device"); + }); + it("keeps waiting on pending and slows down by the server interval or by five seconds", () => { const pending = connectFlowReducer(started(), { type: "polled", response: { status: "pending" } }); expect(pending).toMatchObject({ kind: "awaiting", pollCount: 1, intervalSeconds: 5 }); diff --git a/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.ts b/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.ts index 0e38fc91bb1..6f0e76f9aaf 100644 --- a/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.ts +++ b/ui/litellm-dashboard/src/components/model_add/github_copilot_connect_flow.ts @@ -34,8 +34,22 @@ export type ConnectFlowEvent = export const initialConnectFlowState: ConnectFlowState = { kind: "starting", attempt: 0 }; -const safeVerificationUri = (uri: string): string => - uri.startsWith("https://github.com/") ? uri : GITHUB_DEVICE_LOGIN_URL; +const safeVerificationUri = (uri: string): string => { + try { + const url = new URL(uri); + const usesHttps = url.protocol === "https:"; + const hasCredentials = url.username !== "" || url.password !== ""; + const hasNoPort = url.port === ""; + const hostIsAllowed = + url.hostname === "github.com" || /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.ghe\.com$/.test(url.hostname); + const urlHasAllowedAuthority = !hasCredentials && hasNoPort && hostIsAllowed; + const isSafeUrl = usesHttps && urlHasAllowedAuthority; + + return isSafeUrl ? uri : GITHUB_DEVICE_LOGIN_URL; + } catch { + return GITHUB_DEVICE_LOGIN_URL; + } +}; const afterPoll = ( state: Extract, diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index 78adb7e4d40..23879a1d23e 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -53,6 +53,7 @@ import { getProviderLogoAndName } from "./provider_info_helpers"; import { displayMode, isDecisionMode } from "@/lib/decisionModels"; import { modelUsageExample } from "@/components/AIHub/modelUsageExample"; import { filterBySearchTerm, rankBySearchRelevance } from "@/utils/searchUtils"; +import { ModelCustomMetadata } from "@/components/AIHub/ModelCustomMetadata"; interface PublicModelHubProps { accessToken?: string | null; @@ -749,6 +750,10 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded

{selectedModel.description}

)} +
{/* Wildcard Routing Note */} diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 05ba7e04c39..959b5ea2375 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -254,6 +254,34 @@ describe("DataTable sorting", () => { }); }); +describe("DataTable error state", () => { + it("shows the failure and a Retry button instead of the empty state when the fetch failed", async () => { + const onRetry = vi.fn(); + render( + , + ); + const alert = screen.getByRole("alert"); + expect(alert).toHaveTextContent("Failed to load"); + expect(alert).toHaveTextContent("Request timed out after 30s"); + expect(screen.queryByText("No people")).not.toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: "Retry" })); + expect(onRetry).toHaveBeenCalledTimes(1); + }); + + it("keeps showing rows it already has when a refetch fails", () => { + render(); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + expect(screen.getByText("Alice")).toBeInTheDocument(); + }); +}); + describe("DataTable layout", () => { it("stretches the table to fill the container when resizing is on, so hidden columns leave no right-side gap", () => { render(); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index a94c4e34611..0c484de8e73 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -26,6 +26,7 @@ import { type VisibilityState, } from "@tanstack/react-table"; import { SearchX } from "lucide-react"; +import { ErrorState } from "@/components/shared/ErrorState"; import * as React from "react"; import { Fragment, useEffect, useState } from "react"; @@ -599,6 +600,9 @@ export function DataTable(props: DataTableProps(props: DataTableProps ); } + if (isError && rows.length === 0) { + return ( + + + + ); + } if (rows.length === 0) { return ( diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts index 4529e3df164..787ea633492 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts @@ -28,6 +28,8 @@ export interface DataTableResolvedProps { isLoading?: boolean; isError?: boolean; + errorMessage?: React.ReactNode; + onRetry?: () => void; loadingMessage?: string; skeletonRowCount?: number; noDataMessage?: React.ReactNode; diff --git a/ui/litellm-dashboard/src/components/shared/ErrorState.tsx b/ui/litellm-dashboard/src/components/shared/ErrorState.tsx new file mode 100644 index 00000000000..140aaf211c9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/ErrorState.tsx @@ -0,0 +1,26 @@ +import { Button } from "@/components/ui/button"; +import { TriangleAlert } from "lucide-react"; +import * as React from "react"; + +interface ErrorStateProps { + title?: string; + message?: React.ReactNode; + onRetry?: () => void; +} + +export function ErrorState({ title = "Failed to load", message, onRetry }: ErrorStateProps) { + return ( +
+
+ +
+
{title}
+ {message ?
{message}
: null} + {onRetry ? ( + + ) : null} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx index 7f567a7f936..fc054e51e5b 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx @@ -816,3 +816,44 @@ describe("TeamInfoView - member budget reset prompt", () => { expect(screen.getByRole("option", { name: "claude-haiku-4-5" })).toBeInTheDocument(); }); }); + +describe("TeamInfoView - model rate limit picker", () => { + const props = { + teamId: "123", + onUpdate: vi.fn(), + onClose: vi.fn(), + accessToken: "test-token", + is_team_admin: true, + is_proxy_admin: true, + editTeam: false, + }; + + beforeEach(() => { + seedDefaultMocks(); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("offers access-group models, not no-default-models, in the model rate limit picker", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + models: ["no-default-models"], + access_group_models: ["claude-haiku-4-5"], + }) as TeamData, + ); + + renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(screen.getByRole("button", { name: /add model limit/i })); + await user.click(screen.getByPlaceholderText("Select model")); + + expect(await screen.findByRole("option", { name: "claude-haiku-4-5" })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: "no-default-models" })).not.toBeInTheDocument(); + expect(screen.queryByRole("option", { name: "gpt-4o" })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 29c048034f4..0d63f5d0833 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -75,6 +75,7 @@ import { computeTeamModelBadges, normalizeTeamModelSelection, TeamAccessGroupModelGrant, + teamPerModelLimitOptions, TeamModelBadge, TeamModelBadgeKind, } from "./teamModelAccess"; @@ -89,7 +90,6 @@ import ModelAliasManager from "../common_components/ModelAliasManager"; import AgentSelector from "../agent_management/AgentSelector"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; -import { unfurlWildcardModelsInList } from "../key_team_helpers/fetch_available_models_team_key"; import GuardrailSettingsView from "../GuardrailSettingsView"; import LoggingSettingsView from "../logging_settings_view"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; @@ -666,20 +666,9 @@ const TeamInfoView: React.FC = ({ [mcpServersLoading || mcpToolsetsLoading || accessGroupsLoading, "the MCP server inventory is still loading"], ] as const ).find(([failed]) => failed)?.[1] ?? null; - const availableRateLimitModels = useMemo(() => { + const perModelLimitModels = useMemo(() => { const selected = (isEditing ? watchedModels : undefined) ?? teamData?.team_info?.models ?? []; - if (selected.includes("all-proxy-models") || selected.includes("all-team-models")) { - return userModels; - } - return unfurlWildcardModelsInList(selected, userModels); - }, [isEditing, watchedModels, teamData, userModels]); - const availableModelBudgetModels = useMemo(() => { - const selected = (isEditing ? watchedModels : undefined) ?? teamData?.team_info?.models ?? []; - if (selected.length === 0 || selected.includes("all-proxy-models") || selected.includes("all-team-models")) { - return userModels; - } - const granted = [...selected, ...(teamData?.team_info?.access_group_models ?? [])]; - return unfurlWildcardModelsInList(Array.from(new Set(granted)), userModels); + return teamPerModelLimitOptions(selected, teamData?.team_info?.access_group_models ?? [], userModels); }, [isEditing, watchedModels, teamData, userModels]); const teamEditAccess = useMemo(() => parseTeamEditAccess(teamData?.team_info?.caller_edit_access), [teamData]); @@ -1724,7 +1713,7 @@ const TeamInfoView: React.FC = ({ premiumUser={premiumUser} value={teamMemberModelMaxBudget} onChange={setTeamMemberModelMaxBudget} - availableModels={availableModelBudgetModels} + availableModels={perModelLimitModels} hint="Set a default per-model spend cap for each team member, with an independent reset window." /> = ({ premiumUser={premiumUser} value={teamModelMaxBudget} onChange={setTeamModelMaxBudget} - availableModels={availableRateLimitModels} + availableModels={perModelLimitModels} usage={info.model_max_budget_usage} hint="Cap this team's spend on individual models, each with its own reset window. Every key on the team shares the cap unless the key sets its own budget for that model." /> @@ -1906,7 +1895,7 @@ const TeamInfoView: React.FC = ({ inputId={id} value={value ?? ""} onValueChange={onChange} - options={availableRateLimitModels.map((model) => ({ + options={perModelLimitModels.map((model) => ({ label: model, value: model, }))} @@ -2644,7 +2633,7 @@ const TeamInfoView: React.FC = ({ ), type: "model-max-budget" as const, - availableModels: availableModelBudgetModels, + availableModels: perModelLimitModels, premiumUser, }, { diff --git a/ui/litellm-dashboard/src/components/team/teamModelAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamModelAccess.test.ts index 3f0b4e20407..5ca42a9391b 100644 --- a/ui/litellm-dashboard/src/components/team/teamModelAccess.test.ts +++ b/ui/litellm-dashboard/src/components/team/teamModelAccess.test.ts @@ -1,5 +1,10 @@ import { describe, expect, it } from "vitest"; -import { computeTeamModelBadges, normalizeTeamModelSelection, TeamAccessGroupModelGrant } from "./teamModelAccess"; +import { + computeTeamModelBadges, + normalizeTeamModelSelection, + TeamAccessGroupModelGrant, + teamPerModelLimitOptions, +} from "./teamModelAccess"; const GRANTS: TeamAccessGroupModelGrant[] = [ { access_group_id: "ag-1", access_group_name: "shared", models: ["haiku", "gpt-4o-mini"] }, @@ -84,3 +89,35 @@ describe("computeTeamModelBadges", () => { ]); }); }); + +describe("teamPerModelLimitOptions", () => { + const USER_MODELS = ["gpt-4o", "haiku", "openai/gpt-4o-mini", "openai/gpt-4.1"]; + + it("offers access-group models, not the sentinel, for a no-default-models team", () => { + expect(teamPerModelLimitOptions(["no-default-models"], ["gpt-4o", "haiku"], USER_MODELS)).toEqual([ + "gpt-4o", + "haiku", + ]); + }); + + it("merges direct and access-group models without duplicates", () => { + expect(teamPerModelLimitOptions(["gpt-4o"], ["gpt-4o", "haiku"], USER_MODELS)).toEqual(["gpt-4o", "haiku"]); + }); + + it("offers every user model when the team covers all models", () => { + expect(teamPerModelLimitOptions([], [], USER_MODELS)).toEqual(USER_MODELS); + expect(teamPerModelLimitOptions(["all-proxy-models"], [], USER_MODELS)).toEqual(USER_MODELS); + expect(teamPerModelLimitOptions(["all-team-models"], [], USER_MODELS)).toEqual(USER_MODELS); + }); + + it("expands wildcards to concrete names and drops the pattern itself", () => { + expect(teamPerModelLimitOptions(["no-default-models"], ["openai/*"], USER_MODELS)).toEqual([ + "openai/gpt-4o-mini", + "openai/gpt-4.1", + ]); + }); + + it("offers nothing when the team has no models it can call", () => { + expect(teamPerModelLimitOptions(["no-default-models"], [], USER_MODELS)).toEqual([]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/teamModelAccess.ts b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts index d5d9245ed51..3027fbb8686 100644 --- a/ui/litellm-dashboard/src/components/team/teamModelAccess.ts +++ b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts @@ -1,3 +1,5 @@ +import { unfurlWildcardModelsInList } from "../key_team_helpers/fetch_available_models_team_key"; + export const ALL_PROXY_MODELS = "all-proxy-models"; export const NO_DEFAULT_MODELS = "no-default-models"; @@ -18,6 +20,24 @@ export interface TeamModelBadge { tooltip: string; } +const ALL_TEAM_MODELS = "all-team-models"; +const MODEL_SENTINELS: ReadonlySet = new Set([ALL_PROXY_MODELS, ALL_TEAM_MODELS, NO_DEFAULT_MODELS]); + +// Per-model rate limits and budgets are enforced on the model name a request sends, +// so only concrete names are offered: sentinels and wildcard patterns never match one. +export function teamPerModelLimitOptions( + teamModels: string[], + accessGroupModels: string[], + userModels: string[], +): string[] { + const coversAllModels = + teamModels.length === 0 || teamModels.includes(ALL_PROXY_MODELS) || teamModels.includes(ALL_TEAM_MODELS); + const granted = coversAllModels ? userModels : [...new Set([...teamModels, ...accessGroupModels])]; + return unfurlWildcardModelsInList(granted, userModels).filter( + (model) => !MODEL_SENTINELS.has(model) && !model.endsWith("*"), + ); +} + export function normalizeTeamModelSelection(models: string[] | undefined): string[] { return models && models.length > 0 ? models : [NO_DEFAULT_MODELS]; } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4bd945b13ca..2b174a1141a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -5333,6 +5333,39 @@ export interface paths { patch: operations["assemblyai_proxy_route_eu_assemblyai__endpoint__patch"]; trace?: never; }; + "/extraction": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Extract typed fields from a state + * @description Read a `state` text and up to 4 images and get one typed answer per named field in `fields`. + * + * Field types are `text`, `number`, `date`, `time`, `datetime`, `duration`, `list` (set `items`), `estimate` + * (set `min` and `max`) and `forecast` (set `min`, `max` and `horizon`). Each answer carries its `value`, which is + * null when the input does not determine it, plus `confidence`, `p_null`, text `span`s into the state and, for + * estimates, `mean`, `median`, `p10` and `p90`. Images are data URLs or `{content_type, base64}` objects, remote + * URLs are not accepted. An unknown field type, a field name that is not 1 to 64 letters, digits, `_`, `.` or `-` + * starting with a letter or `_`, a remote image URL or more than 4 images is rejected with a 400 before any + * provider is called. Any other top-level key, `stream: true` and OpenAI generation settings such as `temperature` + * included, is rejected with a 400 naming it unless `drop_params` or the deployment's `additional_drop_params` + * drops it. Settings inside `fields` and `options` go to the provider as sent, so a field missing a setting its + * type needs or an empty `state` with no images comes back with the provider's own error. `model` is any + * extraction model in the proxy model list, such as Cloudflare Clairic. The request and the answers keep the + * provider's format. + */ + post: operations["extraction_extraction_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/fal_ai/{endpoint}": { parameters: { query?: never; @@ -20204,6 +20237,39 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/extraction": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Extract typed fields from a state + * @description Read a `state` text and up to 4 images and get one typed answer per named field in `fields`. + * + * Field types are `text`, `number`, `date`, `time`, `datetime`, `duration`, `list` (set `items`), `estimate` + * (set `min` and `max`) and `forecast` (set `min`, `max` and `horizon`). Each answer carries its `value`, which is + * null when the input does not determine it, plus `confidence`, `p_null`, text `span`s into the state and, for + * estimates, `mean`, `median`, `p10` and `p90`. Images are data URLs or `{content_type, base64}` objects, remote + * URLs are not accepted. An unknown field type, a field name that is not 1 to 64 letters, digits, `_`, `.` or `-` + * starting with a letter or `_`, a remote image URL or more than 4 images is rejected with a 400 before any + * provider is called. Any other top-level key, `stream: true` and OpenAI generation settings such as `temperature` + * included, is rejected with a 400 naming it unless `drop_params` or the deployment's `additional_drop_params` + * drops it. Settings inside `fields` and `options` go to the provider as sent, so a field missing a setting its + * type needs or an empty `state` with no images comes back with the provider's own error. `model` is any + * extraction model in the proxy model list, such as Cloudflare Clairic. The request and the answers keep the + * provider's format. + */ + post: operations["extraction_v1_extraction_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/files": { parameters: { query?: never; @@ -27515,7 +27581,7 @@ export interface components { * Mode * @description The mode to test the model with. If not provided, resolved the way /health does: the deployment's model_info.mode (only while the request tests the deployment's own model), then the mode the provider requires for that model, then the model cost map. */ - mode?: ("chat" | "completion" | "embedding" | "audio_speech" | "audio_transcription" | "image_generation" | "image_edit" | "video_generation" | "batch" | "rerank" | "realtime" | "responses" | "anthropic_messages" | "ocr" | "evaluation" | "decisions") | null; + mode?: ("chat" | "completion" | "embedding" | "audio_speech" | "audio_transcription" | "image_generation" | "image_edit" | "video_generation" | "batch" | "rerank" | "realtime" | "responses" | "anthropic_messages" | "ocr" | "evaluation" | "decisions" | "extraction") | null; /** * Model Info * @description Model info for the health check @@ -28399,7 +28465,7 @@ export interface components { * CallTypes * @enum {string} */ - CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "decisions" | "adecisions" | "systemone" | "asystemone" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "decisions" | "adecisions" | "systemone" | "asystemone" | "extraction" | "aextraction" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; /** CallbackDelete */ CallbackDelete: { /** Callback Name */ @@ -32203,6 +32269,69 @@ export interface components { * @enum {string} */ ExportType: "daily" | "daily_with_keys" | "daily_with_models" | "daily_with_users"; + /** ExtractionAnswer */ + ExtractionAnswer: { + /** Confidence */ + confidence?: number | null; + guess?: components["schemas"]["JsonValue"]; + mean?: components["schemas"]["JsonValue"]; + median?: components["schemas"]["JsonValue"]; + p10?: components["schemas"]["JsonValue"]; + p90?: components["schemas"]["JsonValue"]; + /** P Above */ + p_above?: number | null; + /** P Below */ + p_below?: number | null; + /** P Null */ + p_null?: number | null; + /** Quantiles */ + quantiles?: components["schemas"]["JsonValue"][] | null; + /** Span */ + span?: number[] | null; + /** Spans */ + spans?: number[][] | null; + /** Steps */ + steps?: components["schemas"]["JsonValue"][] | null; + /** Type */ + type: string; + /** @description The answer, null when the input does not determine it */ + value: components["schemas"]["JsonValue"]; + } & { + [key: string]: unknown; + }; + /** ExtractionResponse */ + ExtractionResponse: { + /** Answers */ + answers: { + [key: string]: components["schemas"]["ExtractionAnswer"]; + }; + /** Model */ + model: string; + /** State Text */ + state_text?: string | null; + /** Truncated */ + truncated: boolean; + usage: components["schemas"]["ExtractionUsage"]; + } & { + [key: string]: unknown; + }; + /** ExtractionUsage */ + ExtractionUsage: { + /** Image Tokens */ + image_tokens?: number | null; + /** Input Tokens */ + input_tokens: number; + /** Output Tokens */ + output_tokens: number; + /** Passes */ + passes?: number | null; + /** Schema Tokens */ + schema_tokens?: number | null; + /** State Tokens */ + state_tokens?: number | null; + } & { + [key: string]: unknown; + }; /** * FacetListResponse * @description The distinct values one column takes over a filtered query. `data` holds bare values, not entity rows. @@ -38245,6 +38374,12 @@ export interface components { ModelGroupInfoProxy: { /** Configurable Clientside Auth Params */ configurable_clientside_auth_params?: (string | components["schemas"]["ConfigurableClientsideParamsCustomAuth-Output"])[] | null; + /** Custom Metadata */ + custom_metadata?: { + [key: string]: string | number | boolean; + } | null; + /** Custom Tags */ + custom_tags?: string[] | null; /** Description */ description?: string | null; /** Health Checked At */ @@ -59011,6 +59146,180 @@ export interface operations { }; }; }; + extraction_extraction_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + /** + * @example { + * "fields": { + * "days_to_pay": { + * "max": 90, + * "min": 0, + * "type": "estimate" + * }, + * "invoice_date": { + * "type": "date" + * }, + * "items": { + * "instructions": "What was bought?", + * "items": "text", + * "type": "list" + * }, + * "total": { + * "type": "number", + * "unit": "USD" + * }, + * "vendor": { + * "instructions": "Who sent the invoice?", + * "type": "text" + * } + * }, + * "model": "clairic", + * "state": "Invoice #4471 from Acme Tools, dated 2026-03-03. 12 hammers. Total $174. Payment due in 30 days." + * } + */ + "application/json": { + /** + * Fields + * @description Field name to typed field spec. Each name becomes a key in the answers + */ + fields: { + [key: string]: { + /** + * Decimals + * @description Numbers, durations and number lists only, 0 to 6 + */ + decimals?: number | null; + /** + * Horizon + * @description Forecast only, number of steps + */ + horizon?: number | null; + /** + * Instructions + * @description What the field should contain + */ + instructions?: string | null; + /** + * Integer + * @description Numbers and durations only + */ + integer?: boolean | null; + /** Interval */ + interval?: ("step" | "minute" | "hour" | "day" | "week" | "month" | "quarter" | "year") | null; + /** + * Items + * @description Item type, required for list fields + */ + items?: ("text" | "number" | "date") | null; + /** + * Max + * @description Estimate and forecast upper bound + */ + max?: number | string | null; + /** + * Min + * @description Estimate and forecast lower bound + */ + min?: number | string | null; + /** + * Null Threshold + * @description Probability of null above which the value is null + */ + null_threshold?: number | null; + /** Scale */ + scale?: ("linear" | "log") | null; + /** + * Step + * @description Estimate grid step, or month or year + */ + step?: number | string | null; + /** + * Type + * @description The answer type for this field + * @enum {string} + */ + type: "text" | "number" | "date" | "time" | "datetime" | "duration" | "list" | "estimate" | "forecast"; + /** + * Unit + * @description Unit label. For durations, seconds through years + */ + unit?: string | null; + } & { + [key: string]: unknown; + }; + }; + /** + * Images + * @description Up to 4 images as data URLs or content_type plus base64, no remote URLs + */ + images?: (string | ({ + /** + * Base64 + * @description Base64-encoded image bytes + */ + base64: string; + /** + * Content Type + * @description The image media type, such as image/png + */ + content_type: string; + } & { + [key: string]: unknown; + }))[] | null; + /** + * Model + * @description An extraction model from the proxy model_list + */ + model: string; + options?: ({ + /** + * Include Quantiles + * @description Add 99 percentiles to estimates + */ + include_quantiles?: boolean | null; + /** + * Return State Text + * @description Return the text that spans index into + */ + return_state_text?: boolean | null; + } & { + [key: string]: unknown; + }) | null; + /** + * Reference Date + * @description ISO date that anchors dates written without a year + */ + reference_date?: string | null; + /** + * State + * @description The text to read. Must be non-empty unless images are sent. Text answers are spans of it + * @default + */ + state?: string; + } & { + [key: string]: unknown; + }; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ExtractionResponse"]; + }; + }; + }; + }; fal_ai_proxy_route_fal_ai__endpoint__get: { parameters: { query?: never; @@ -78336,6 +78645,180 @@ export interface operations { }; }; }; + extraction_v1_extraction_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + /** + * @example { + * "fields": { + * "days_to_pay": { + * "max": 90, + * "min": 0, + * "type": "estimate" + * }, + * "invoice_date": { + * "type": "date" + * }, + * "items": { + * "instructions": "What was bought?", + * "items": "text", + * "type": "list" + * }, + * "total": { + * "type": "number", + * "unit": "USD" + * }, + * "vendor": { + * "instructions": "Who sent the invoice?", + * "type": "text" + * } + * }, + * "model": "clairic", + * "state": "Invoice #4471 from Acme Tools, dated 2026-03-03. 12 hammers. Total $174. Payment due in 30 days." + * } + */ + "application/json": { + /** + * Fields + * @description Field name to typed field spec. Each name becomes a key in the answers + */ + fields: { + [key: string]: { + /** + * Decimals + * @description Numbers, durations and number lists only, 0 to 6 + */ + decimals?: number | null; + /** + * Horizon + * @description Forecast only, number of steps + */ + horizon?: number | null; + /** + * Instructions + * @description What the field should contain + */ + instructions?: string | null; + /** + * Integer + * @description Numbers and durations only + */ + integer?: boolean | null; + /** Interval */ + interval?: ("step" | "minute" | "hour" | "day" | "week" | "month" | "quarter" | "year") | null; + /** + * Items + * @description Item type, required for list fields + */ + items?: ("text" | "number" | "date") | null; + /** + * Max + * @description Estimate and forecast upper bound + */ + max?: number | string | null; + /** + * Min + * @description Estimate and forecast lower bound + */ + min?: number | string | null; + /** + * Null Threshold + * @description Probability of null above which the value is null + */ + null_threshold?: number | null; + /** Scale */ + scale?: ("linear" | "log") | null; + /** + * Step + * @description Estimate grid step, or month or year + */ + step?: number | string | null; + /** + * Type + * @description The answer type for this field + * @enum {string} + */ + type: "text" | "number" | "date" | "time" | "datetime" | "duration" | "list" | "estimate" | "forecast"; + /** + * Unit + * @description Unit label. For durations, seconds through years + */ + unit?: string | null; + } & { + [key: string]: unknown; + }; + }; + /** + * Images + * @description Up to 4 images as data URLs or content_type plus base64, no remote URLs + */ + images?: (string | ({ + /** + * Base64 + * @description Base64-encoded image bytes + */ + base64: string; + /** + * Content Type + * @description The image media type, such as image/png + */ + content_type: string; + } & { + [key: string]: unknown; + }))[] | null; + /** + * Model + * @description An extraction model from the proxy model_list + */ + model: string; + options?: ({ + /** + * Include Quantiles + * @description Add 99 percentiles to estimates + */ + include_quantiles?: boolean | null; + /** + * Return State Text + * @description Return the text that spans index into + */ + return_state_text?: boolean | null; + } & { + [key: string]: unknown; + }) | null; + /** + * Reference Date + * @description ISO date that anchors dates written without a year + */ + reference_date?: string | null; + /** + * State + * @description The text to read. Must be non-empty unless images are sent. Text answers are spans of it + * @default + */ + state?: string; + } & { + [key: string]: unknown; + }; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ExtractionResponse"]; + }; + }; + }; + }; list_files_v1_files_get: { parameters: { query?: { diff --git a/uv.lock b/uv.lock index 9b490cf95f4..03403828225 100644 --- a/uv.lock +++ b/uv.lock @@ -29,6 +29,7 @@ constraints = [ ] overrides = [ { name = "cryptography", specifier = ">=50.0.0,<51.0" }, + { name = "mcp", specifier = ">=2.3.0,<3" }, { name = "packaging", specifier = ">=24.0" }, ] @@ -4603,7 +4604,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.107.0" +version = "1.108.0" source = { editable = "." } dependencies = [ { name = "aiohttp" }, @@ -4778,7 +4779,8 @@ ci = [ { name = "logfire" }, { name = "lunary", version = "1.4.36", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "lunary", version = "1.4.37", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "pillow" }, + { name = "pillow", version = "12.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pillow", version = "13.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "psutil" }, { name = "psycopg2-binary" }, { name = "pyarrow" }, @@ -4902,8 +4904,8 @@ requires-dist = [ { name = "gunicorn", marker = "extra == 'proxy'", specifier = ">=23.0.0,<24.0" }, { name = "hiredis", marker = "extra == 'proxy'", specifier = ">=3.0.0,<4.0" }, { name = "httpx", extras = ["http2"], specifier = ">=0.28.0,<1.0" }, - { name = "httpx2", marker = "extra == 'mcp'", specifier = ">=2.5.0,<3" }, - { name = "httpx2", marker = "extra == 'proxy'", specifier = ">=2.5.0,<3" }, + { name = "httpx2", marker = "extra == 'mcp'", specifier = ">=2.10.0,<3" }, + { name = "httpx2", marker = "extra == 'proxy'", specifier = ">=2.10.0,<3" }, { name = "huggingface-hub", specifier = ">=0.34.0,<2.0" }, { name = "importlib-metadata", specifier = ">=8.0.0,<9.0" }, { name = "inquirerpy", marker = "extra == 'cli'", specifier = ">=0.3.4,<1.0" }, @@ -4916,8 +4918,8 @@ requires-dist = [ { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, { name = "mangum", marker = "extra == 'proxy-runtime'", specifier = ">=0.17.0,<1.0" }, - { name = "mcp", marker = "extra == 'mcp'", specifier = ">=2.2.0,<3" }, - { name = "mcp", marker = "extra == 'proxy'", specifier = ">=2.2.0,<3" }, + { name = "mcp", marker = "extra == 'mcp'", specifier = ">=2.3.0,<3" }, + { name = "mcp", marker = "extra == 'proxy'", specifier = ">=2.3.0,<3" }, { name = "mlflow", marker = "extra == 'mlflow'", specifier = ">=3.11.1,<4.0" }, { name = "numpy", marker = "extra == 'stt-nvidia-riva'", specifier = ">=1.26.0" }, { name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" }, @@ -4977,7 +4979,7 @@ provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "sa admin-mcp = [{ name = "litellm-admin-mcp", marker = "python_full_version >= '3.12'", url = "https://github.com/BerriAI/liteadmin-mcp/archive/d35ec9c19c117d4c50cc1eccf6ce1296aac25a1a.tar.gz" }] benchmarks = [ { name = "a2a-sdk", specifier = "==1.1.0" }, - { name = "mcp", specifier = ">=2.2.0,<3" }, + { name = "mcp", specifier = ">=2.3.0,<3" }, { name = "pytest", specifier = "==9.0.3" }, { name = "pytest-codspeed", specifier = "==4.3.0" }, ] @@ -5000,7 +5002,8 @@ ci = [ { name = "logfire", specifier = "==4.6.0" }, { name = "lunary", marker = "python_full_version == '3.10.*'", specifier = "==1.4.36" }, { name = "lunary", marker = "python_full_version >= '3.11'", specifier = "==1.4.37" }, - { name = "pillow", specifier = "==12.3.0" }, + { name = "pillow", marker = "python_full_version == '3.10.*'", specifier = "==12.3.0" }, + { name = "pillow", marker = "python_full_version >= '3.11'", specifier = "==13.0.0" }, { name = "psutil", specifier = "==7.2.2" }, { name = "psycopg2-binary", specifier = "==2.9.11" }, { name = "pyarrow", specifier = "==23.0.1" }, @@ -5055,7 +5058,7 @@ dev = [ e2e-dev = [ { name = "anthropic", specifier = "==0.84.0" }, { name = "locust", specifier = "==2.45.0" }, - { name = "mcp", specifier = ">=2.2.0,<3" }, + { name = "mcp", specifier = ">=2.3.0,<3" }, { name = "playwright", specifier = "==1.61.0" }, { name = "psutil", specifier = "==7.2.2" }, { name = "websockets", specifier = ">=15.0.1,<16.0" }, @@ -5105,12 +5108,12 @@ provides-extras = ["test"] [[package]] name = "litellm-enterprise" -version = "0.1.78" +version = "0.1.79" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.110" +version = "0.4.111" source = { editable = "litellm-proxy-extras" } dependencies = [ { name = "psycopg" }, @@ -5497,7 +5500,8 @@ dependencies = [ { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "packaging" }, - { name = "pillow" }, + { name = "pillow", version = "12.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pillow", version = "13.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pyparsing" }, { name = "python-dateutil" }, ] @@ -5570,7 +5574,7 @@ wheels = [ [[package]] name = "mcp" -version = "2.2.0" +version = "2.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -5588,22 +5592,22 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/76/31/ac54fb0fdd5b37de704486e288bba4fbbb463f24cfcfedbede407b854513/mcp-2.2.0.tar.gz", hash = "sha256:2dc37ecb1974becdcebdbf7561e7c15a07dbbf20ba21ba16c3593b3038b3afbd", size = 4084129, upload-time = "2026-09-07T16:06:23.439Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9d/8d/e0d339616f4810e9051d4aba6887afab289ab1f81875fe908b606cdfd0e3/mcp-2.3.0.tar.gz", hash = "sha256:8b147a50441cf059dc88c684e0aeed3687f0aa0f39c6cde7b90330effd2b34d8", size = 4171955, upload-time = "2026-10-02T22:06:56.091Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1b/ff/8e7eade68b8a28f7da0ed1085544341b51f9c935dbf6b95c76b7edfea6a0/mcp-2.2.0-py3-none-any.whl", hash = "sha256:bde982589473a060ae145e3406e9a5333fe538c97229ba841f5a7f92be004f81", size = 365656, upload-time = "2026-09-07T16:06:19.711Z" }, + { url = "https://files.pythonhosted.org/packages/20/d7/4a2878f128df729db3ca04e71bfce26405cc2e0c9ebe19c592feeb4c61b9/mcp-2.3.0-py3-none-any.whl", hash = "sha256:dd0c44c089d16453e8ae31a3877a0054d7a2314caaa81f5e0541b9b1734b2377", size = 368664, upload-time = "2026-10-02T22:06:52.119Z" }, ] [[package]] name = "mcp-types" -version = "2.2.0" +version = "2.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pydantic" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ae/91/762d7755d971aff8a28d75f7961656148edf27875c8026e6385aaab08ae7/mcp_types-2.2.0.tar.gz", hash = "sha256:d3ed53703ddd10d9c6399f29d322bb66f3f67ab41348ac8556ba23e07fedefad", size = 65892, upload-time = "2026-09-07T16:06:25.187Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9e/2d/7c251e34207f6c51000312fc8839111ac45cfe02023f90b44e7f1051dd8e/mcp_types-2.3.0.tar.gz", hash = "sha256:d1e46549edb35ee19a94940fcee6d1addd7e589ab7ea92dda83f5d84781fc362", size = 65898, upload-time = "2026-10-02T22:06:57.863Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8f/d7/6ffba5d8cd5dd9b8a19478875c50e04945314ba5074e84d749283f27f62d/mcp_types-2.2.0-py3-none-any.whl", hash = "sha256:ea476b73ee86709ab5abc9452385ed36cc05907e582355622e294595c9a04f13", size = 69106, upload-time = "2026-09-07T16:06:21.461Z" }, + { url = "https://files.pythonhosted.org/packages/ce/48/32583066b155bb5d4507e7887dc88ee23118cc4de8e2a12870c49719e2a8/mcp_types-2.3.0-py3-none-any.whl", hash = "sha256:968efdbdaedfab06adae40d378a34395f1090c5921d4be3c9cde283aaf76d91d", size = 69108, upload-time = "2026-10-02T22:06:54.099Z" }, ] [[package]] @@ -7173,6 +7177,9 @@ wheels = [ name = "pillow" version = "12.3.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/25/c2/669d88644cddb1485bd9534e63e8cf476c8e51cb3c3a1297677023505c0e/pillow-12.3.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:6c0016e7b354317c4e9e525b937ac8596c38d2d232b419529b9cd7a1cd46e39a", size = 5392418, upload-time = "2026-07-01T11:53:27.808Z" }, @@ -7242,6 +7249,81 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/54/0169bc772ec491108b62f644f8ecf1fe5d8ae5ebafde2ee2142210166903/pillow-12.3.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:04f01d28a6aaff387bf842a13be313df23ba0597a44f1a976c9feb3c6ff4711a", size = 7231786, upload-time = "2026-07-01T11:56:35.046Z" }, ] +[[package]] +name = "pillow" +version = "13.0.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.*'", +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/a4/9cc30690a394a8305e8be4c9036aa16d0dfddadd09079860b582f562496d/pillow-13.0.0.tar.gz", hash = "sha256:cdcd1ba258178f7356d74170c9a81ba4eab6a5d55748bc402a5a2a6fe6be3510", size = 47058631, upload-time = "2026-10-16T22:18:09.718Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/d7/af6af1b8cc55e07cb0489d5b0950bf5a9bb1461eeba825700b6d8d0e26ea/pillow-13.0.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:6f38f3f885a09cafc93fd667b1173a92f379e43785d635e7983c893f0da319fb", size = 5592324, upload-time = "2026-10-16T22:15:24.672Z" }, + { url = "https://files.pythonhosted.org/packages/0b/d2/2c12cb4fbb5d95f38d1001ff3d95fafcb830077018cba40e56a8f932dfbc/pillow-13.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:37290fbef0c73e2d51d594eb78fded7dc7170d1c3ad0a890269e6fd8c8c4e061", size = 4856298, upload-time = "2026-10-16T22:15:27.669Z" }, + { url = "https://files.pythonhosted.org/packages/5c/ee/229b27472483a4801ea94990815e710391836afdf987b4b722b62bd75db6/pillow-13.0.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7c88c10dd432f3b39ccdafc9ed81d0c21b8fd634502574da76d1bfa4fe2c76d0", size = 6517304, upload-time = "2026-10-16T22:15:29.43Z" }, + { url = "https://files.pythonhosted.org/packages/36/46/f509b74787f40ce0627501089d630b2c741b63b244c1f40b640b3b01e181/pillow-13.0.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a625df9a2a5597fa5e5266e01aefc82604bb85d291e3547981e863560d75da6c", size = 7140667, upload-time = "2026-10-16T22:15:31.368Z" }, + { url = "https://files.pythonhosted.org/packages/aa/a4/5a398b8ec2c37145397d6461c0b551c23ec03072aa64f357aed25e582b00/pillow-13.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:10b5019883c98f1dd819ca48ffc6f33a91cee04948fe637c605f4345c7cafd6f", size = 6610583, upload-time = "2026-10-16T22:15:34.025Z" }, + { url = "https://files.pythonhosted.org/packages/4f/7b/f2f1109b230fd4e2d57ab2305de31adc37ea9b4ff8719436c623cb484415/pillow-13.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f391dc92a015c478745192fe4f2e120c2af2aba18d9fe0616341cbf537599806", size = 7218241, upload-time = "2026-10-16T22:15:36.214Z" }, + { url = "https://files.pythonhosted.org/packages/ab/d3/069edb9083b267746baf59756282cacef904d3475c5288c7216fb5195a98/pillow-13.0.0-cp311-cp311-win32.whl", hash = "sha256:6a0920cb55df5a4d3f6862964908078c0bd59332bdd1bcd20d14e8c70ef27cc5", size = 6515278, upload-time = "2026-10-16T22:15:38.129Z" }, + { url = "https://files.pythonhosted.org/packages/7d/f6/6e2d96413e4c9dbec12ea584d9e646d602877d4cce3ea669099f959deac6/pillow-13.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:d7710eb676e9bb68965cb045be0d333b4c0d7f9dfbaa7fbf121385b1588683d3", size = 7278471, upload-time = "2026-10-16T22:15:40.286Z" }, + { url = "https://files.pythonhosted.org/packages/70/e7/baff21f8cb60496c8c250609501b88919eeedb4088d3da106f6c62f3f3c0/pillow-13.0.0-cp311-cp311-win_arm64.whl", hash = "sha256:12fbb851cf323cc19aac3eea1df5ad17b57e91e9fe3cfd06100dba7bb8208c13", size = 2698981, upload-time = "2026-10-16T22:15:42.637Z" }, + { url = "https://files.pythonhosted.org/packages/7a/b8/58d79332229482eb0d27c82973850808c19d37e5bdc14443d3417bd373f5/pillow-13.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:67c8d44ad89fa7eb88de5d6ef6199b1e03dd24974f85a67639fc8badb801fe79", size = 5568643, upload-time = "2026-10-16T22:15:44.474Z" }, + { url = "https://files.pythonhosted.org/packages/39/4b/f2ed53d83aa44ac19599ac0a54897f867fafb95f406980d1e24083f8cbe3/pillow-13.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc481f742ccda20e9ed760b8f978dff50f2a4e9bc253ac167fb0a05b0783ef52", size = 4852329, upload-time = "2026-10-16T22:15:46.284Z" }, + { url = "https://files.pythonhosted.org/packages/3b/bb/9e30e8ed18c20800ee327a12107d5249ad8dc444867b3adfcc21654e3c71/pillow-13.0.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8bf295fe459343bc6093c8e5e0df60bfa189e0fcb775a22e2c2488fdeccc6c0d", size = 6519309, upload-time = "2026-10-16T22:15:48.14Z" }, + { url = "https://files.pythonhosted.org/packages/e6/99/fdf65cc539e266ab93546c261bdf98411fe165c6e4e13f9dbf597097192b/pillow-13.0.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:42d63be3af4f9b951b199b0a9d23365d37ab14687e6d77b74940432c8515c4ca", size = 7144542, upload-time = "2026-10-16T22:15:50.093Z" }, + { url = "https://files.pythonhosted.org/packages/26/9b/f38b7c4d8f599537e2d1a6d546cec9beef7536ff8f5c49a3f9c78179d8ac/pillow-13.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4bd024d567eca931b668978d672dba914d4b8e2628e2c3073a2ad178de8a2620", size = 6618042, upload-time = "2026-10-16T22:15:52.054Z" }, + { url = "https://files.pythonhosted.org/packages/a8/7c/45ee28531f89ee4582967d8a17f7a26d71ca64cd902b93c269590cd60a00/pillow-13.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:23edf4ec4b6f4f240e8df6e4985e9ee2ced86474632433e55fb1e2548744bf40", size = 7228137, upload-time = "2026-10-16T22:15:54.009Z" }, + { url = "https://files.pythonhosted.org/packages/01/bc/9c2a969cb1b70201d04fa17fd5aedf02a732870b8da86bceaaa30b0ca385/pillow-13.0.0-cp312-cp312-win32.whl", hash = "sha256:02464da86f561f410e13b061ec37c3068a9cb557a095494157742d07b3c24a27", size = 6515423, upload-time = "2026-10-16T22:15:56.25Z" }, + { url = "https://files.pythonhosted.org/packages/ff/cd/73a79ecdc822076002eb24b5e1e4a530b0bb5310fcb6d239abbec877110b/pillow-13.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:ff3fe9eb621fbb9db4c4be33ab946f6bac5f98ab50ca8d3260503653656c961b", size = 7272166, upload-time = "2026-10-16T22:15:58.597Z" }, + { url = "https://files.pythonhosted.org/packages/4a/a3/be5a1236efd085b3984e4305da8f179ae9dab80c991d2ad78f2bd5c4c843/pillow-13.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:c3bc8ee0271387e0b1e3ad0ff893645990a656e79bf9c42f4f734a3642b11d51", size = 2698943, upload-time = "2026-10-16T22:16:00.637Z" }, + { url = "https://files.pythonhosted.org/packages/82/e2/9132d38df320e8099e2642d9fa538f18a7dfebf68a875a31725dc2a115aa/pillow-13.0.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:e98b4f91f99b24b1da72e10621ba181f7281922384a7403094debb4199de6c00", size = 4190620, upload-time = "2026-10-16T22:16:02.527Z" }, + { url = "https://files.pythonhosted.org/packages/cb/be/436dac0996876f65ef1916ffb711c5ef4af2ebc57414b6f16642e07f6b47/pillow-13.0.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:34fb3db1149c879343861831728c243c432d4a5afc8f1878ca7612ef8de2a033", size = 4290286, upload-time = "2026-10-16T22:16:04.861Z" }, + { url = "https://files.pythonhosted.org/packages/60/16/7c1007162e4169ea0d5bec47e2e0b2174f56c50cf5b8c02a1dc23e41037e/pillow-13.0.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:2d50500a4a3bc6469adcdb319cf05d688b08ca3bf0f22a05339f3b25b433a599", size = 3723169, upload-time = "2026-10-16T22:16:06.956Z" }, + { url = "https://files.pythonhosted.org/packages/08/bc/578ab2eaa04705a8978bd9be0fe15c3d43a8fc95097c4f7e39540e79e4fa/pillow-13.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:ef5fe6b7db3f501c56000ba8fe7c7e1a35ca1d45824fbeb95a67daa85cff4dbe", size = 5568469, upload-time = "2026-10-16T22:16:08.664Z" }, + { url = "https://files.pythonhosted.org/packages/98/c8/ccede3dcc4f274b80fb49199eed1ef34ba05566c8b71976b9d1f111da2e3/pillow-13.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f11a3c1358bb3c68012a7b1f41b223c3db9f7bfde90b481df028cfe1c3363313", size = 4852187, upload-time = "2026-10-16T22:16:10.693Z" }, + { url = "https://files.pythonhosted.org/packages/fb/3c/cd0770e96d20bfd990db024bcbd43897ff3e88da1f4782cf90ea231b1d4f/pillow-13.0.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:19b66ca0a4d25cc0200b381ce824c7ea090db61427315e73cfb05dbf1b35c369", size = 6515908, upload-time = "2026-10-16T22:16:12.896Z" }, + { url = "https://files.pythonhosted.org/packages/ef/fe/f5cdaa2d1c4086e757430f368c271a58ee29cfe7378e7ae7b9d964015c25/pillow-13.0.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5f56499c1a74d2b862beb0c33b8016f99d095858455df9d21583bbbc1c15b582", size = 7142544, upload-time = "2026-10-16T22:16:14.844Z" }, + { url = "https://files.pythonhosted.org/packages/fb/b5/a54f6d44adbcd23737c8c12953f42be107e4cca64e247c0ca3a5b563061c/pillow-13.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a2a14cd73d5e6c927f644b0110ee81f59c4ef506259ec9183904747f2caacc81", size = 6613097, upload-time = "2026-10-16T22:16:16.893Z" }, + { url = "https://files.pythonhosted.org/packages/e6/b4/9e6c1207c21b2e91db0d15c16d30f3701ce5add1bd4da0e787761e66cb6e/pillow-13.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:08b826a9bf1265c76de04a42061d61ea6b109418b7826f28ff1ecdfaa26c5502", size = 7224846, upload-time = "2026-10-16T22:16:19.195Z" }, + { url = "https://files.pythonhosted.org/packages/78/9f/d3c3cb0518f3dd1eee0ef501535f53d15597a53c5c4118d178086143397a/pillow-13.0.0-cp313-cp313-win32.whl", hash = "sha256:12b0fa15879fe5a9d9eebe66e4e2d74128a773b1a4448010c9c714f45e04b9d8", size = 6516013, upload-time = "2026-10-16T22:16:21.389Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ed/1b16fb66a37e1b9e41370b93b2c614c7da761a5b2a41dc333f7af4639962/pillow-13.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:9dda0788a9932776a75d8bcc020076edc308ce2508b51f6739a45e3158b3fc42", size = 7284777, upload-time = "2026-10-16T22:16:23.393Z" }, + { url = "https://files.pythonhosted.org/packages/87/7d/8d02eba9f7ecc87f8e90f731e032cc4773802f2a843763afd0c66a535ff4/pillow-13.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:40a266371ece581e5de13b6635ef78d7487824af80561b6ca27694a88c19bc0f", size = 2698826, upload-time = "2026-10-16T22:16:25.382Z" }, + { url = "https://files.pythonhosted.org/packages/8e/45/fbad9c7f5e2f1745c373a434d46072bc84eddfefaa28b852425ac957fda6/pillow-13.0.0-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:44629d4498c9cb6b2051806155f22141df0d8a929c56134b166655232dfb2f8c", size = 4190721, upload-time = "2026-10-16T22:16:27.524Z" }, + { url = "https://files.pythonhosted.org/packages/a2/bf/2074902a259a90692c5dbf132885d21671644cbbd3f690047516243a55db/pillow-13.0.0-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:bf97ba1b5542050e9b0c74e632323b6d8136cc99d64e193abebee22fe6013c76", size = 4290236, upload-time = "2026-10-16T22:16:29.562Z" }, + { url = "https://files.pythonhosted.org/packages/ff/a5/3d1bb93b8ea38d15d21356c1b8e9f50971ac592aaa248f7fb71739094c0e/pillow-13.0.0-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:aa694c49787040943570d61c9b8fc01271bd2a2cfdd55f0d896f8aecdce1a981", size = 3723093, upload-time = "2026-10-16T22:16:31.249Z" }, + { url = "https://files.pythonhosted.org/packages/84/a1/be060135d85187b25fa8b2573786c0e45c62ae27527fff5257675eb97f4a/pillow-13.0.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:49c12af515073e9bf511e4ef736743bf70d531a643668722279314b937742ba4", size = 5572253, upload-time = "2026-10-16T22:16:33.103Z" }, + { url = "https://files.pythonhosted.org/packages/b5/f4/0bf6f1ac9ba7ee2ffe9914aeb410a1035a0c1a3c42f122388d7fc950bc4b/pillow-13.0.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d9b0a2bee3f3e03e9f29057aedb807dbf1e1d96ee1b9d0f422d9db6c8583271d", size = 4852269, upload-time = "2026-10-16T22:16:35.102Z" }, + { url = "https://files.pythonhosted.org/packages/83/16/b603274617b3e6728da3a0133d4fba1ea8d7fbe0c4ea5fef9cbffa398b8b/pillow-13.0.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5933fcea5a5ba64c686f55c886a115cbfd92f2900460784a1e79909fdbafd7b6", size = 6516058, upload-time = "2026-10-16T22:16:37.228Z" }, + { url = "https://files.pythonhosted.org/packages/1d/f2/1371707ce438ba9a43f9d24089828494ec6f9396e28cd7903aa704dd6306/pillow-13.0.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e91c2b1cbeae768c98752e67e717e283d68dae45717b41f2d91eaee648e37b7d", size = 7142574, upload-time = "2026-10-16T22:16:39.172Z" }, + { url = "https://files.pythonhosted.org/packages/b9/a9/62cf68587f910bc87c5951bcb6e9430c7643a65076adf5f18b4bf350e2bd/pillow-13.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7f6299e494e46253d40c34653cac369412f44ca53b5f85d848d3f95b1ed0546c", size = 6614295, upload-time = "2026-10-16T22:16:41.117Z" }, + { url = "https://files.pythonhosted.org/packages/6c/c9/9e9b48178649730299be0bafbf5f810a5b6977538f5a92d47e4ba5f5d608/pillow-13.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f9bd7d260d986f3553518c5d723544346fa48031f8aef3265d786aeb59845bc0", size = 7224841, upload-time = "2026-10-16T22:16:43.118Z" }, + { url = "https://files.pythonhosted.org/packages/8f/48/c44e91384b1b6c072a714a63d00608acc16a6a968735e478e4be4384e5ac/pillow-13.0.0-cp314-cp314-win32.whl", hash = "sha256:2e31c4e68cffd7b1f116127626813b82fcba19edf6b940a7b0c266d6484d0055", size = 6508694, upload-time = "2026-10-16T22:16:45.883Z" }, + { url = "https://files.pythonhosted.org/packages/50/e2/7ad17747517e4c813dca78257d6bb593466f8231baf1384bc98d92745edb/pillow-13.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:f7cbb7a64e7ec664775195e4f81923bc8a93d9451b963f844e59118f63147a7a", size = 7281830, upload-time = "2026-10-16T22:16:47.725Z" }, + { url = "https://files.pythonhosted.org/packages/41/9f/e5b51f2ce8ef4563215334035403ce61a244f93a2e65c24a940131618f87/pillow-13.0.0-cp314-cp314-win_arm64.whl", hash = "sha256:a6aaa7871aebca081eb5c45f83e845bbd9712ae99fa9286dbfdcee835e63f4b4", size = 2695468, upload-time = "2026-10-16T22:16:49.515Z" }, + { url = "https://files.pythonhosted.org/packages/97/73/a53502f731f6c84805a4ce87a8881e86d07a65aee64728b900d112dc703e/pillow-13.0.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:7b999816b5850a35ada08443e6e03945f1f6488144a83afe2d37b240a552aafa", size = 5578196, upload-time = "2026-10-16T22:16:51.291Z" }, + { url = "https://files.pythonhosted.org/packages/21/50/7e949959a56b5559c78ef172b70fc6025706acb289f3f648a15329e4fa6f/pillow-13.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:67871a13cf47275201553b536c025da777f997fa562a19eb798af4ad4b97c18f", size = 4857276, upload-time = "2026-10-16T22:16:53.399Z" }, + { url = "https://files.pythonhosted.org/packages/c9/52/43ded0ac08e77e39d17bd6d894162572201b5e8fdc4a022b14c55a1960a4/pillow-13.0.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e8ca8f69b7c77fa96dad99f8638c1fc4f1c0bed4acaa1baf7a3c31bdb80a0ab1", size = 6607008, upload-time = "2026-10-16T22:16:55.525Z" }, + { url = "https://files.pythonhosted.org/packages/0d/2e/4a57ed216d5bf877387cba4d44c02aa87abe51ad2e06404aa6b597c4da39/pillow-13.0.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6bce6a29c3898191f4c9998e4d4e8f69790caeabde46e54da4bc235ce5bde683", size = 7224281, upload-time = "2026-10-16T22:16:57.738Z" }, + { url = "https://files.pythonhosted.org/packages/5c/3d/2929b248e135752d63aaeabf53f6de797e905b7b6263f9c1e2adade20aba/pillow-13.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:816752a5ef76a80484c8c441c85593b7ea9304f6231fc5776ad1fbff7e773058", size = 6695600, upload-time = "2026-10-16T22:16:59.842Z" }, + { url = "https://files.pythonhosted.org/packages/bc/33/fb0da992cf084ed1ee6ee05e96b6191e35c351eb707ca2df874feb24ce47/pillow-13.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5bdb38ba220b61b4ec31dd1046dea6b016d3c6850a23f43ed8b58332b2fc668e", size = 7300589, upload-time = "2026-10-16T22:17:01.686Z" }, + { url = "https://files.pythonhosted.org/packages/bb/af/c30dd9ef6a3dce84ecb1b2393cac629047a32e66279056c6cb46c3ffd65c/pillow-13.0.0-cp314-cp314t-win32.whl", hash = "sha256:06903aac074d00c988e77a1fc435286dc498506d4ca8d6f69df9babddf4ecc15", size = 6517204, upload-time = "2026-10-16T22:17:03.722Z" }, + { url = "https://files.pythonhosted.org/packages/85/d5/ef3a8ae7609c8ab26570abf18194882534b480d245a9148b895a2262c004/pillow-13.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:9a635068d9793c7f0d929f39f18b58f1bd86c236e8e7b64167b86af9232ebca0", size = 7286695, upload-time = "2026-10-16T22:17:05.543Z" }, + { url = "https://files.pythonhosted.org/packages/84/3c/0861db0c69a03689ebe3e7599d35668a414deb5aaafd8255215e4e62e3d3/pillow-13.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0d1a9cf7bc6917538b20f8d06b9ab2147e542f68849e96fc58eb66bdf6acc692", size = 2698098, upload-time = "2026-10-16T22:17:07.674Z" }, + { url = "https://files.pythonhosted.org/packages/68/c0/fb7b1caf25ed5ac5d46116f09b8b0706b9559833c65c667cecb1a53153f2/pillow-13.0.0-pp311-pypy311_pp80-macosx_10_15_x86_64.whl", hash = "sha256:405be6058e206a86fedce2f715ecc2343f94072d7925251632dffa2c05f70dca", size = 5521630, upload-time = "2026-10-16T22:17:50.008Z" }, + { url = "https://files.pythonhosted.org/packages/1a/e6/4d11d5b1137f4fc080605cce410cfaae4f303ad400bdec75d27d3083fa29/pillow-13.0.0-pp311-pypy311_pp80-macosx_11_0_arm64.whl", hash = "sha256:b84be41387fc00db28c13c02e30d42b02d8d8f6a4a741a1aee6012ca7f1f6039", size = 4806023, upload-time = "2026-10-16T22:17:51.932Z" }, + { url = "https://files.pythonhosted.org/packages/f1/45/0923202a658e8ff6aafd72c1a4395ebadcdebaa182c437d8595d3dd3db32/pillow-13.0.0-pp311-pypy311_pp80-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a192360d512305fe92d3053e4f7627bbad1f7708f1f4ad38cccd9336acabe9ad", size = 5464831, upload-time = "2026-10-16T22:17:53.742Z" }, + { url = "https://files.pythonhosted.org/packages/07/52/d3e38f0cad1dadea2a658b9a7d6ba8a04048f5ae14f011b806e5c3bd6bcb/pillow-13.0.0-pp311-pypy311_pp80-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:83b37abd87b590acd8f47ca8a61f70893322b11834d36e07c2298dacad4d14d2", size = 6043133, upload-time = "2026-10-16T22:17:55.625Z" }, + { url = "https://files.pythonhosted.org/packages/c4/f9/2ed7685980c1d64f8131a7d87a19c369424747577152622b73ddeff163b7/pillow-13.0.0-pp311-pypy311_pp80-win_amd64.whl", hash = "sha256:73ee25cf8cf8b75d5a5aeed7cb6588f5e5dbfbb0fde51707b630bbb23897645c", size = 7365208, upload-time = "2026-10-16T22:17:57.509Z" }, + { url = "https://files.pythonhosted.org/packages/60/da/f6d644e631c5b370c5d0b899f3ccadd2780aab070b29a110527cf4e1fd2c/pillow-13.0.0-pp312-pypy312_pp80-macosx_10_15_x86_64.whl", hash = "sha256:8619685cad8efc4c38c66d3b36aad854523297c43d748b3b6f0cc03c90ba3463", size = 5523188, upload-time = "2026-10-16T22:17:59.334Z" }, + { url = "https://files.pythonhosted.org/packages/e4/b0/2594f202d743a35d4fe80682bfd96f277e5ca7aecb27cf64ef5bfc1d000b/pillow-13.0.0-pp312-pypy312_pp80-macosx_11_0_arm64.whl", hash = "sha256:7e91e97a929232cdb828889ee6ad0e48d7f40f4ec38818f4e4a8968b2b1af22e", size = 4806808, upload-time = "2026-10-16T22:18:01.188Z" }, + { url = "https://files.pythonhosted.org/packages/8e/ed/d504165ef79af6fafe25c76726cf30f544cc236002f7ed054f1d51adc4ff/pillow-13.0.0-pp312-pypy312_pp80-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:93ff9c9c3663ce23aaa5d0b6b6abac55359e2a0c06c7e30f3771829af816be63", size = 5464358, upload-time = "2026-10-16T22:18:02.779Z" }, + { url = "https://files.pythonhosted.org/packages/7b/fe/a10aba5503ca13e9e009b8bd0a1a4ce5584bdce246245fd6f3a83701369c/pillow-13.0.0-pp312-pypy312_pp80-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e65dd6e75de67b2362a4fa3e16aa627595f65054dc4ec73d64680c84d592fa61", size = 6043514, upload-time = "2026-10-16T22:18:04.538Z" }, + { url = "https://files.pythonhosted.org/packages/c0/37/66b12dc476041fe523bca29c9d58fbf1ea92a4eaf73acfb9e993948f1191/pillow-13.0.0-pp312-pypy312_pp80-win_amd64.whl", hash = "sha256:567f20f1c75897d7663830da6ea8af719e03d712b1d9ee7f813319f9ece42e5a", size = 7366030, upload-time = "2026-10-16T22:18:06.638Z" }, +] + [[package]] name = "platformdirs" version = "4.9.6" @@ -8697,7 +8779,8 @@ version = "5.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "charset-normalizer" }, - { name = "pillow" }, + { name = "pillow", version = "12.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pillow", version = "13.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/4a/51/dbe28534ae12c852f61be91f039f343305fd1f34f1c66b8de75afae7a525/reportlab-5.0.1.tar.gz", hash = "sha256:ebd13154be1c8515e665de70bd2d303ae9ddc3ef47e44afd5116441ca0283a26", size = 3945711, upload-time = "2026-08-20T13:48:16.461Z" } wheels = [ diff --git a/vscode-extension/package-lock.json b/vscode-extension/package-lock.json index 497c5807d5f..536d7a88f25 100644 --- a/vscode-extension/package-lock.json +++ b/vscode-extension/package-lock.json @@ -899,9 +899,9 @@ } }, "node_modules/@oxc-project/types": { - "version": "0.150.0", - "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.150.0.tgz", - "integrity": "sha512-rDS5/31E9HfPl/CIzGrn0DOlvBbXFseQ5URJ9sYMfstbKLD/c6Gm9vmRzRGDdAXyOIL4zmO37lc9RIwYqVruZw==", + "version": "0.154.0", + "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.154.0.tgz", + "integrity": "sha512-KYOUGQFlLWR+s4QMV2Q9ko38TZm7EdO87Qm7rvZLqk/webse/cDjtyxKhg/47l0byCbs/k46LT8qc4PKjAikzA==", "dev": true, "license": "MIT", "funding": { @@ -909,9 +909,9 @@ } }, "node_modules/@rolldown/binding-android-arm-eabi": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm-eabi/-/binding-android-arm-eabi-1.2.9.tgz", - "integrity": "sha512-tNISae1QEf/vkb3xkRcjV5SEdzPE97We5IVaa2Z8jSszQPZ8U60B/YCYpw4QI7VidYsBtKavczXf+DyDs9WGxw==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm-eabi/-/binding-android-arm-eabi-1.2.14.tgz", + "integrity": "sha512-aWR+XKH9+rW7m5fWq/5360J3mHycrx/o6xZM2DQFdJnTMdyP7/WzxPjhtLbVWdKo2y612/4U1bWqb3HQhiF7vg==", "cpu": [ "arm" ], @@ -926,9 +926,9 @@ } }, "node_modules/@rolldown/binding-android-arm64": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm64/-/binding-android-arm64-1.2.9.tgz", - "integrity": "sha512-YC8YsI30o606GTZi0VyzYlsDKFP8W61i/QzayHDkLbNEz/IShqAmTa+hsJRj13xTHA0H+6fk4b2UmGn+Q/cMlg==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm64/-/binding-android-arm64-1.2.14.tgz", + "integrity": "sha512-QiVxHofl6A8lQU7WAtOUY4QMVZXohxwz7WJJoeh8NBEUyKfIkMRKhYyq8az/1RypNiIXzNTtiS6KwHz0erxYRQ==", "cpu": [ "arm64" ], @@ -943,9 +943,9 @@ } }, "node_modules/@rolldown/binding-darwin-arm64": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.2.9.tgz", - "integrity": "sha512-IwhlH3qK5urrY8hZiEgGkHKEFN901p/p2bjxCxJlr4GyNnF7wYpUvK+Y43uaRYuC4hpfjzbR3SJC3arX1jGvmw==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.2.14.tgz", + "integrity": "sha512-lTQOzHuBIooh3jIr5XINQkjPGEuGy+Dsy0meXDr6Sjty0b99kYgGh4uBC/8mJX1T0kC8d6S0aPP1Y+bHkh6aBw==", "cpu": [ "arm64" ], @@ -960,9 +960,9 @@ } }, "node_modules/@rolldown/binding-darwin-x64": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.2.9.tgz", - "integrity": "sha512-XxpJfVzFh+jilRxIXUqcfYAYcunIc/XEzIizsOL1fcJee5Sf7H3mH8WlLmfHfluz5amqR88QQo9izKtmMlavAw==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.2.14.tgz", + "integrity": "sha512-hyMKEGatjmBYCs/Q4sbIcL/GqpFEqsX3wc+HYVQ/85A4rxc1o7rPm4DUI4GV4k/h0ITddwFSGpqQunI14GXjTA==", "cpu": [ "x64" ], @@ -977,9 +977,9 @@ } }, "node_modules/@rolldown/binding-freebsd-x64": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.2.9.tgz", - "integrity": "sha512-kSfvhmgeWyfkbT3p/1s5vSgboogoah2zkm9fX2zjg2hHxSV7T4KhMWRUUaRk4OXNqoD3QAUeRqLcs1aZOK4U1g==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.2.14.tgz", + "integrity": "sha512-u371GbXoVXrEgP20PwzoGNOdAOqAMhzo19JzzY+V4YmmnC+DTOoqdbMu+O+54emM59Mik/UdrtWgcuEDKWRPKA==", "cpu": [ "x64" ], @@ -994,9 +994,9 @@ } }, "node_modules/@rolldown/binding-linux-arm-gnueabihf": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.2.9.tgz", - "integrity": "sha512-1RVzG17pxqbTfYLC352JlLt6kKLG+6Hr30n8DlIJqsnV5luUDd2Qdx9Ayw1Cabfyb1K9k0jXEZ7evxkRoT+uiw==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.2.14.tgz", + "integrity": "sha512-BqwKsdZQf6EqVWr8fv2MrDYzAvOZ7+NYFSJESRGs87fkseEhs0TU9TXgRJTkZoRRkxTizyEg1LKUXfU3/78uqg==", "cpu": [ "arm" ], @@ -1011,9 +1011,9 @@ } }, "node_modules/@rolldown/binding-linux-arm64-gnu": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.2.9.tgz", - "integrity": "sha512-BXqPvZ2drqVD+/Z8UpKwcs4Mp7grM+eGFku4CAEKrEtcbAsUpzREphK1sogCRZGreVPiMkiiBtw0n3TPteuqvw==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.2.14.tgz", + "integrity": "sha512-SD9kY8M009x/7s57Z1dgUVlP5+uYBgKQK0Tkr1WIWngrtUJmRCW1RJFMY7QDla2khXhmg2i5Io7MUjcmZNZ/KQ==", "cpu": [ "arm64" ], @@ -1031,9 +1031,9 @@ } }, "node_modules/@rolldown/binding-linux-arm64-musl": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.2.9.tgz", - "integrity": "sha512-11vWvo8YDwLzukt27J3aYDWU+gg2P7J+ZOmiJ0hkF5BXZDW7pVya7r40MXDy6ya0i9KamoENSVKIugvJNgFXIA==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.2.14.tgz", + "integrity": "sha512-BJ2CL19M6uuWesMgnqJzgxoI+N+eV7b0A9O9vLV42y+gitC41fNa8KElau1tHVzuLWX5GzEuxMPhAO0tQXQMkg==", "cpu": [ "arm64" ], @@ -1051,9 +1051,9 @@ } }, "node_modules/@rolldown/binding-linux-ppc64-gnu": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.2.9.tgz", - "integrity": "sha512-a1tijMkdwsIARtc0F39ApURROkf3NwqinI6TOiSSWCTR7dT96dffNvMUtDHnq64wKNTIZOIlzKrFvvFUznJiyw==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.2.14.tgz", + "integrity": "sha512-4+Kl8daPV8NDHaD7mfCcc1pRM+Qg7ab8DyMm5FG5WoZgv4h2i3KFN1CybYBNbQqmuYm7ftJLvmMI1eHYJ22S/Q==", "cpu": [ "ppc64" ], @@ -1071,9 +1071,9 @@ } }, "node_modules/@rolldown/binding-linux-s390x-gnu": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.2.9.tgz", - "integrity": "sha512-x6SQNdAvv4c3hWqTMaWuawzMX9myaCs/yEmlGsxJzkdClnHW7FbrjQuSiRDhuSYzEYoEMhsaJy9qHG/XNemJPQ==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.2.14.tgz", + "integrity": "sha512-V/h8ePq1fKDrmtFMz6xcCN3DKvuvMRrwwD1q0PkBeeo4DpK64biaKvqWGVxqCeYKqeRk7gE7/iKPxLCn08O97A==", "cpu": [ "s390x" ], @@ -1091,9 +1091,9 @@ } }, "node_modules/@rolldown/binding-linux-x64-gnu": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.2.9.tgz", - "integrity": "sha512-9s0AZ8BFK5/n7B/TBoa2yJE3gI3KURrbXcPBlsAsvjU4VeJKgE90y1YtNxyEUIcHPQkg6/yfF3qihUrcM/Kf0Q==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.2.14.tgz", + "integrity": "sha512-WxalXzgrLVpKC93xfpaJ3gCXUWzeCr+FRCbDVD2mUAy4L6Vo3lNVz0epUEUpxPWJkLNYUwM4OSIikv36OJ9mLg==", "cpu": [ "x64" ], @@ -1111,9 +1111,9 @@ } }, "node_modules/@rolldown/binding-linux-x64-musl": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.2.9.tgz", - "integrity": "sha512-P7VWAmV+WdJluH7ovnRGoiv2i8To7GAZ+kGzfGup635cyL7SyYl3lSUaA3Gp5THf0n/Co5EyEqb2zbqq+nMOHQ==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.2.14.tgz", + "integrity": "sha512-3G3tnAmTpdk9jO3GQtyHXqeFD2SlecfTws/Y3M2mmj9nMBwwuLT4kur8sZKeu7ENUHSpVO3cIsm5A52LEWWwlg==", "cpu": [ "x64" ], @@ -1131,9 +1131,9 @@ } }, "node_modules/@rolldown/binding-openharmony-arm64": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.2.9.tgz", - "integrity": "sha512-1qixtsE4BK8h+yS3BfmZ09UhA7O/N4IACva6YBr7EBvCJraByTuRcgOTaiA62Tm0vey3UcKXLOaoGHtYmNGEVg==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.2.14.tgz", + "integrity": "sha512-qyr5XiNGxm2ytx17vvtcJ9nXn6v1gENETi503bJx3OGwZ1tSonj7BosxbMa/xeAfSl5iS9X6KjEWYXLu6v1L7Q==", "cpu": [ "arm64" ], @@ -1148,9 +1148,9 @@ } }, "node_modules/@rolldown/binding-win32-arm64-msvc": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.2.9.tgz", - "integrity": "sha512-ok8IQjcEPs1AKZfuEUznVBrJw+gK4soq+bx8b1X2XoMqVClarc1q5JDmVtWXY1xfr6ZuHTAsPXHTgTrqKTZeww==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.2.14.tgz", + "integrity": "sha512-oA4o5o0S2yB199YyYXDzKUv1Z56Rc+mgOjwm7Z/vJ9Z8igiDPYvoOuDyo6g7AK/ZaF09nPhDKVnGBcYHO11GXA==", "cpu": [ "arm64" ], @@ -1165,9 +1165,9 @@ } }, "node_modules/@rolldown/binding-win32-x64-msvc": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.2.9.tgz", - "integrity": "sha512-Ip2mXoU0hM0boq3Rf+ekuT653OROSo6aSYcPT1VHE4q52KvyxgFkQgrgb/IEsxOuvQ2fZZbs8khJAyCEPM24/g==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.2.14.tgz", + "integrity": "sha512-9mE94bf2GZUaR/fUSXCMpBoWkctu421Mi0ORpm4zqfbkjKep7RxzPFa5etcGs5h+2KApoWBOA/hAJa/7saGgEQ==", "cpu": [ "x64" ], @@ -2730,9 +2730,9 @@ "license": "ISC" }, "node_modules/nanoid": { - "version": "3.3.19", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.19.tgz", - "integrity": "sha512-Y2tUNy4ouw6tq5oDSKeQYGOyhkUBhNOcGV/02KC+6kd9eDGqdZd++mjMiIDilrBYvjEnCYvVtsuHCuP+okSfug==", + "version": "3.3.20", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.20.tgz", + "integrity": "sha512-uKdg2G3GNCKQn9byYOpxbGqrT2fGO5KRt5J/8b3pok8rT6qxGWF6hxMyJiEYtAf+FVyYuD9hRaDqX5uPFYJ4ZQ==", "dev": true, "funding": [ { @@ -2879,9 +2879,9 @@ } }, "node_modules/postcss": { - "version": "8.5.28", - "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.28.tgz", - "integrity": "sha512-RRuzqDtt5Y9h3quz5hWhK+TPnsmVs6WwSU6LkJMeY4HstUEDuYTG8UJSdawMRzmzAtV+KEoG8N3Qg2qLy5vM/A==", + "version": "8.5.30", + "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.30.tgz", + "integrity": "sha512-zkd+e9S1AQHhmaBKye9zlFHWjOwVgqXl1q3ECYy8561K7qxR2cREf/STLskulBFd78k6AOFeA+ff1TtEzk9bAA==", "dev": true, "funding": [ { @@ -2899,9 +2899,9 @@ ], "license": "MIT", "dependencies": { - "nanoid": "^3.3.18", + "nanoid": "^3.3.20", "picocolors": "^1.1.1", - "source-map-js": "^1.2.1" + "source-map-js": "^1.2.2" }, "engines": { "node": "^10 || ^12 || >=14" @@ -2960,13 +2960,13 @@ } }, "node_modules/rolldown": { - "version": "1.2.9", - "resolved": "https://registry.npmjs.org/rolldown/-/rolldown-1.2.9.tgz", - "integrity": "sha512-hx/Pv0N1haXRb11qkfnK5MXB/iqr7i0yjWQqmO9uHqZpBgQSqzc8UsSnEpalsh+j1I8qQ2CkXAkJC8Br3dKSlg==", + "version": "1.2.14", + "resolved": "https://registry.npmjs.org/rolldown/-/rolldown-1.2.14.tgz", + "integrity": "sha512-eTBTTqMR6uHKLgcUqpJP8ntzDI6V4LYA8negAMBEI7gkgVvOdrx6VsvS2Ym+83KhOPdLCHWJu0w7UQFJZhHKNw==", "dev": true, "license": "MIT", "dependencies": { - "@oxc-project/types": "=0.150.0", + "@oxc-project/types": "=0.154.0", "@rolldown/pluginutils": "^1.0.0" }, "bin": { @@ -2976,21 +2976,21 @@ "node": "^20.19.0 || >=22.12.0" }, "optionalDependencies": { - "@rolldown/binding-android-arm-eabi": "1.2.9", - "@rolldown/binding-android-arm64": "1.2.9", - "@rolldown/binding-darwin-arm64": "1.2.9", - "@rolldown/binding-darwin-x64": "1.2.9", - "@rolldown/binding-freebsd-x64": "1.2.9", - "@rolldown/binding-linux-arm-gnueabihf": "1.2.9", - "@rolldown/binding-linux-arm64-gnu": "1.2.9", - "@rolldown/binding-linux-arm64-musl": "1.2.9", - "@rolldown/binding-linux-ppc64-gnu": "1.2.9", - "@rolldown/binding-linux-s390x-gnu": "1.2.9", - "@rolldown/binding-linux-x64-gnu": "1.2.9", - "@rolldown/binding-linux-x64-musl": "1.2.9", - "@rolldown/binding-openharmony-arm64": "1.2.9", - "@rolldown/binding-win32-arm64-msvc": "1.2.9", - "@rolldown/binding-win32-x64-msvc": "1.2.9" + "@rolldown/binding-android-arm-eabi": "1.2.14", + "@rolldown/binding-android-arm64": "1.2.14", + "@rolldown/binding-darwin-arm64": "1.2.14", + "@rolldown/binding-darwin-x64": "1.2.14", + "@rolldown/binding-freebsd-x64": "1.2.14", + "@rolldown/binding-linux-arm-gnueabihf": "1.2.14", + "@rolldown/binding-linux-arm64-gnu": "1.2.14", + "@rolldown/binding-linux-arm64-musl": "1.2.14", + "@rolldown/binding-linux-ppc64-gnu": "1.2.14", + "@rolldown/binding-linux-s390x-gnu": "1.2.14", + "@rolldown/binding-linux-x64-gnu": "1.2.14", + "@rolldown/binding-linux-x64-musl": "1.2.14", + "@rolldown/binding-openharmony-arm64": "1.2.14", + "@rolldown/binding-win32-arm64-msvc": "1.2.14", + "@rolldown/binding-win32-x64-msvc": "1.2.14" } }, "node_modules/run-applescript": { @@ -3312,16 +3312,16 @@ } }, "node_modules/vite": { - "version": "8.3.0", - "resolved": "https://registry.npmjs.org/vite/-/vite-8.3.0.tgz", - "integrity": "sha512-lhZBVvEHefgE+HQZC9O7EBJgCU/nVzFNl7vkS4RE0APtWLP02/8QVIkQtzBxPquh7lq5/78NHipTj7ODQ6XuyQ==", + "version": "8.3.4", + "resolved": "https://registry.npmjs.org/vite/-/vite-8.3.4.tgz", + "integrity": "sha512-IftLE8phJhdNRIGNEV4JCCvfNScoZQVfcXaN8TQropWGJtOHcu+iAB9pHcuGvmEb/kddw0LBe/jS8ztdC6xdiQ==", "dev": true, "license": "MIT", "dependencies": { "lightningcss": "^1.33.0", "picomatch": "^4.0.7", - "postcss": "^8.5.28", - "rolldown": "~1.2.6", + "postcss": "^8.5.29", + "rolldown": "~1.2.12", "tinyglobby": "^0.2.17" }, "bin": {