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 2b2d0a82431..961795838a8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -8,21 +8,20 @@ from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast +from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( + ManagedObjectCreatorAttribution, + ManagedObjectRow, +) + from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( - CLI_SESSION_KEY_PREFIX, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, ) from litellm.repositories.table_repositories import ManagedObjectRepository -from litellm.repositories.team_repository import TeamRepository -from litellm.repositories.user_repository import UserRepository -from litellm.repositories.verification_token_repository import VerificationTokenRepository if TYPE_CHECKING: - from prisma import models as prisma_models - from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.prisma_protocols import TableActions @@ -47,29 +46,6 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = ( ) -class _ManagedObjectRow(Protocol): - @property - def id(self) -> str: ... - - @property - def unified_object_id(self) -> str: ... - - @property - def created_by(self) -> str | None: ... - - @property - def file_object(self) -> object: ... - - @property - def org_id(self) -> str | None: ... - - @property - def api_key(self) -> str | None: ... - - @property - def team_id(self) -> str | None: ... - - class _ReadableFileContent(Protocol): async def read(self) -> bytes: ... @@ -87,22 +63,10 @@ async def _file_content_bytes(file_content: object) -> bytes: return cast(bytes, file_content) -def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]": +def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]": return ManagedObjectRepository(prisma_client).table -def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]": - return UserRepository(prisma_client).table - - -def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]": - return VerificationTokenRepository(prisma_client).table - - -def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]": - return TeamRepository(prisma_client).table - - class CheckBatchCost: def __init__( self, @@ -118,6 +82,7 @@ class CheckBatchCost: self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router self._track_unmanaged_batch_cost = track_unmanaged_batch_cost + self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost") # Cached after the first poll cycle. Once we know the column is absent we skip # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True @@ -149,133 +114,6 @@ class CheckBatchCost: return self.batch_processed_support_confirmed = True - async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]: - """ - Look up user email and key alias by user_id for enriching the S3 callback metadata. - Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None). - Returns an empty dict when user_id is None: batches created by a team or service - account key carry no user id, and find_unique(where={"user_id": None}) raises. - """ - if not user_id: - return {} - try: - user_row: prisma_models.LiteLLM_UserTable | None = await _user_table(self.prisma_client).find_unique( - where={"user_id": user_id} - ) - if user_row is None: - return {} - return { - "user_api_key_user_email": getattr(user_row, "user_email", None), - "user_api_key_alias": getattr(user_row, "user_alias", None), - } - except Exception as e: - verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") - return {} - - async def _get_key_alias(self, batch_id: str, api_key: str | None, created_by: str | None) -> str | None: - """Resolve the creating virtual key's alias from its hashed token.""" - if not api_key: - return None - if created_by and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}": - return api_key - try: - key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table( - self.prisma_client - ).find_unique(where={"token": api_key}) - return getattr(key_row, "key_alias", None) if key_row is not None else None - except Exception as e: - verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}") - return None - - async def _get_team_alias(self, team_id: str | None) -> str | None: - """Resolve a team's alias from its id.""" - if not team_id: - return None - try: - team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique( - where={"team_id": team_id} - ) - return getattr(team_row, "team_alias", None) if team_row is not None else None - except Exception as e: - verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}") - return None - - async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | None: - org_id: Final = getattr(job, "org_id", None) - if org_id: - return org_id - api_key: Final = getattr(job, "api_key", None) - team_id: Final = getattr(job, "team_id", None) - if api_key: - try: - key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table( - self.prisma_client - ).find_unique(where={"token": api_key}) - key_org_id: Final = ( - cast(str | None, getattr(key_row, "organization_id", None)) - if key_row is not None - else None - ) - if key_org_id: - return key_org_id - except Exception as e: - verbose_proxy_logger.error( - f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, " - f"still trying the team's: {e}" - ) - if not team_id: - return None - try: - team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique( - where={"team_id": team_id} - ) - return cast(str | None, getattr(team_row, "organization_id", None)) if team_row is not None else None - except Exception as e: - verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}") - return None - - async def _build_creator_attribution_metadata( - self, job: "_ManagedObjectRow", batch_id: str - ) -> dict[str, object]: - """ - Rebuild the spend-tracking metadata for the key, team, and tags that created the - batch so the batch-cost spend log is attributed the same way a non-batch request - is. Rows created before api_key and request_tags were persisted carry only - created_by and team_id, and fall back to those. A named creating key owns - user_api_key_alias; when it has no alias, or the key has since been rotated or - deleted, the field keeps the creating user's alias that _get_user_info filled in, - because a resolvable name is more useful on the spend row than a null. - - user_api_key_org_id must be resolved here too: the spend update writer reads it - off this metadata to increment organization spend, so leaving it out silently - drops batch cost from org accounting for keys and teams that belong to one. - """ - api_key = getattr(job, "api_key", None) - team_id = getattr(job, "team_id", None) - request_tags = getattr(job, "request_tags", None) - - metadata: dict[str, object] = { - "user_api_key_user_id": job.created_by, - "user_api_key": api_key, - "user_api_key_hash": api_key, - "user_api_key_team_id": team_id, - **(await self._get_user_info(batch_id, job.created_by)), - } - - key_alias = await self._get_key_alias(batch_id, api_key, job.created_by) - if key_alias is not None: - metadata["user_api_key_alias"] = key_alias - team_alias = await self._get_team_alias(team_id) - if team_alias is not None: - metadata["user_api_key_team_alias"] = team_alias - org_id: Final = await self._get_org_id(job, batch_id) - if org_id is not None: - metadata["user_api_key_org_id"] = org_id - if isinstance(request_tags, list) and request_tags: - metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)] - - return metadata - async def _cleanup_stale_managed_objects(self) -> None: """ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days @@ -317,7 +155,7 @@ class CheckBatchCost: f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed" ) - async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]": + async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]": """Query batch jobs without the batch_processed filter (for older schemas).""" return await _managed_object_table(self.prisma_client).find_many( where={ @@ -337,7 +175,7 @@ class CheckBatchCost: order={"created_at": "asc"}, ) - async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None: + async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None: """ Take a row that can never be costed out of the poll page. Leaving it selectable would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and @@ -364,7 +202,7 @@ class CheckBatchCost: "so it will no longer be polled" ) - async def _claim_job_for_costing(self, job: "_ManagedObjectRow") -> bool: + async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool: """ Atomically flip batch_processed from false to true, returning whether this pod won the row. Every pod and uvicorn worker schedules its own poller against the shared @@ -390,7 +228,7 @@ class CheckBatchCost: return False return claimed > 0 - async def _release_job_claim(self, job: "_ManagedObjectRow") -> None: + async def _release_job_claim(self, job: "ManagedObjectRow") -> None: """Give a claimed row back once billing it failed, so a later poll cycle retries it. Safe to match on batch_processed=True: while this poller is active the retrieve @@ -411,7 +249,7 @@ class CheckBatchCost: ) @staticmethod - def _has_unified_id_without_model(job: "_ManagedObjectRow") -> bool: + def _has_unified_id_without_model(job: "ManagedObjectRow") -> bool: """A unified id that decodes but carries no model_id can never be routed.""" from litellm.proxy.openai_files_endpoints.common_utils import ( convert_b64_uid_to_unified_uid, @@ -460,7 +298,7 @@ class CheckBatchCost: return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error) async def _finalize_unbilled_terminal_job( - self, job: "_ManagedObjectRow", response: "LiteLLMBatch" + self, job: "ManagedObjectRow", response: "LiteLLMBatch" ) -> None: """Persist a terminal batch that has nothing billable, converting any raw provider file ids to managed ids, and take it out of the poll page.""" @@ -505,7 +343,7 @@ class CheckBatchCost: def _resolve_job_routing( self, - job: "_ManagedObjectRow", + job: "ManagedObjectRow", prom_logger: Optional["PrometheusLogger"], ) -> Optional[Tuple[str, str]]: """ @@ -582,7 +420,7 @@ class CheckBatchCost: def _resolve_unmanaged_provider_routing( self, - job: "_ManagedObjectRow", + job: "ManagedObjectRow", prom_logger: Optional["PrometheusLogger"], llm_provider: str, bare_model_name: str, @@ -678,7 +516,7 @@ class CheckBatchCost: @classmethod def _get_managed_file_model_name( cls, - job: "_ManagedObjectRow", + job: "ManagedObjectRow", deployment_info: "Deployment", ) -> Optional[str]: """ @@ -698,7 +536,7 @@ class CheckBatchCost: ) @staticmethod - def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]: + def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]: import json from litellm.types.utils import LiteLLMBatch @@ -719,7 +557,7 @@ class CheckBatchCost: async def _track_completed_batch_cost( self, - job: "_ManagedObjectRow", + job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, @@ -735,9 +573,9 @@ class CheckBatchCost: """ from litellm.batches.batch_utils import ( + calculate_batch_cost_and_usage, count_error_file_failed_requests, get_file_content_as_dictionary, - calculate_batch_cost_and_usage, ) from litellm.files.main import afile_content from litellm.files.types import FileContentCallOptions @@ -911,7 +749,7 @@ class CheckBatchCost: }, **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}), "metadata": { - **(await self._build_creator_attribution_metadata(job, batch_id)), + **(await self._creator_attribution.batch_metadata(job, batch_id)), # spend logs read the deployment identity off these metadata keys, so # without them the batch cost row carries no model_id or model_group "model_info": {"id": model_id}, 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 7dd1fd7a1ee..00779bfe8ec 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -1,12 +1,36 @@ """ Polls LiteLLM_ManagedObjectTable to check if the response is complete. -Cost tracking is handled by the get-responses call, which prices normally only because the -poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the -same route are non-inference and free. +A background response is billed exactly once: whichever client read or poll first sees it +finished moves its row out of a pending status, then logs the body that read already fetched as +one spend event stamped with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN and the creating key's +attribution, priced against the deployment that served it. The poll's status checks are stamped +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. +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 +the finished response through it before the client sees response.completed. +A read the caller retargets with client-side credentials, and a delete of a still-pending response, +settle the bill from a fetch through the configured deployment instead: the client controls what a +retargeted read returns, and a deleted response leaves the poll nothing to fetch. +A row stored before this billing model keeps the encrypted create id as its model_object_id. The +first lookup by that id rekeys it by the provider id, and the poll retires such a row without +billing it once a read has already tracked the response under the provider id. """ +import uuid +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Dict, Final, Protocol, cast +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, cast + +from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( + ManagedObjectCreatorAttribution, + ManagedObjectRow, +) +from pydantic import BaseModel, ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -14,36 +38,124 @@ from litellm.constants import ( INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, + PENDING_RESPONSE_STATUSES, STALE_OBJECT_CLEANUP_BATCH_SIZE, + TERMINAL_RESPONSE_STATUSES, ) +from litellm.proxy.hooks.responses_id_security import decrypted_response_id from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN +from litellm.types.router import Deployment +from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, +) if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.prisma_protocols import TableActions from litellm.router import Router -TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"}) +_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) +_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses" -class _ManagedObjectRow(Protocol): - @property - def id(self) -> str: ... - - @property - def unified_object_id(self) -> str: ... - - @property - def created_by(self) -> str | None: ... - - @property - def file_object(self) -> object: ... +@dataclass(frozen=True, slots=True) +class _PollTarget: + response_id: str + model_metadata: Mapping[str, object] + deployment: Deployment | None -def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]": +@dataclass(frozen=True, slots=True) +class BilledBackgroundResponse: + response: ResponsesAPIResponse + model: str + custom_llm_provider: str | None + litellm_params: Mapping[str, object] + metadata: Mapping[str, object] + + +@dataclass(frozen=True, slots=True) +class BilledCharge: + cost: float | None + logging_obj: "LiteLLMLogging" + + +BilledResponseLogger = Callable[[BilledBackgroundResponse], Awaitable[BilledCharge]] +UntrackedResponseTracker = Callable[[ResponsesAPIResponse], Awaitable[bool]] + + +def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse: + unpriced: Final = response.model_copy(deep=True) + unpriced.hidden_params = {key: value for key, value in response.hidden_params.items() if key != "response_cost"} + 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 + + logging_obj: Final = LiteLLMLogging( + model=billed.model, + messages=[{"role": "user", "content": ""}], + stream=False, + call_type=_RETRIEVE_RESPONSE_CALL_TYPE, + start_time=datetime.now(), + litellm_call_id=str(uuid.uuid4()), + function_id=str(uuid.uuid4()), + ) + logging_obj.update_environment_variables( + litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)}, + optional_params={}, + custom_llm_provider=billed.custom_llm_provider, + ) + unpriced: Final = _without_read_price(billed.response) + cost: Final = logging_obj.response_cost_calculator(result=unpriced) + end_time: Final = datetime.now() + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + async_coroutine=logging_obj.async_success_handler( + result=unpriced, start_time=logging_obj.start_time, end_time=end_time + ) + ) + logging_obj.handle_sync_success_callbacks_for_async_calls( + result=unpriced, start_time=logging_obj.start_time, end_time=end_time + ) + return BilledCharge(cost=cost, logging_obj=logging_obj) + + +class _DeploymentBaseModel(BaseModel): + model_config = ConfigDict(frozen=True) + + base_model: str | None = None + + +def _deployment_litellm_params(deployment: Deployment) -> Mapping[str, object]: + from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials + + api_base: Final = deployment.litellm_params.api_base + base_model: Final = _DeploymentBaseModel.model_validate(deployment.litellm_params, from_attributes=True).base_model + return MappingProxyType( + { + **({"api_base": mask_api_base_credentials(api_base)} if api_base else {}), + **({"base_model": base_model} if base_model else {}), + } + ) + + +def _deployment_metadata(deployment: Deployment) -> Mapping[str, object]: + return MappingProxyType( + { + "model_info": deployment.model_info.model_dump(exclude_none=True), + "model_group": deployment.model_name, + "deployment": deployment.litellm_params.model, + } + ) + + +def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]": return ManagedObjectRepository(prisma_client).table @@ -51,12 +163,17 @@ def _is_response_gone_at_provider(error: Exception, provider_response_id: str) - return getattr(error, "status_code", None) == 404 and provider_response_id in str(error) +def _keyed_by_issued_id(job: ManagedObjectRow) -> bool: + return decrypted_response_id(job.model_object_id) != job.model_object_id + + class CheckResponsesCost: def __init__( self, proxy_logging_obj: "ProxyLogging", prisma_client: "PrismaClient", llm_router: "Router", + log_billed_response: BilledResponseLogger = log_billed_background_response, ): from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -64,15 +181,17 @@ class CheckResponsesCost: self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router + self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost") + self._log_billed_response: Final = log_billed_response - def _resolve_deployment(self, response_id: str) -> bool: + def _resolve_deployment(self, response_id: str) -> Deployment | None: model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) - return model_id is not None and self.llm_router.get_deployment(model_id=model_id) is not None + return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None async def _get_response( self, response_id: str, - litellm_metadata: Dict[str, str], + litellm_metadata: dict[str, object], via_router: bool, ) -> ResponsesAPIResponse: """Fetch the upstream response, using deployment credentials when available. @@ -91,9 +210,17 @@ class CheckResponsesCost: ) return cast(ResponsesAPIResponse, router_response) - async def _expire_stale_rows( - self, cutoff: datetime, batch_size: int - ) -> int: + async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse: + return await self._get_response( + response_id=target.response_id, + litellm_metadata={ + **target.model_metadata, + INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, + }, + via_router=target.deployment is not None, + ) + + async def _expire_stale_rows(self, cutoff: datetime, batch_size: int) -> int: """Execute the bounded UPDATE that marks stale rows as 'stale_expired'. Isolated so it can be swapped / mocked in tests without touching the @@ -140,14 +267,154 @@ class CheckResponsesCost: f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" ) + def _target(self, response_id: str, model_name: str | None) -> _PollTarget: + return _PollTarget( + response_id=response_id, + model_metadata=MappingProxyType({"model": model_name, "model_group": model_name} if model_name else {}), + deployment=self._resolve_deployment(response_id), + ) + + def _poll_target(self, job: ManagedObjectRow) -> _PollTarget: + stored_model: Final = _STORED_RESPONSE_ADAPTER.validate_python(job.file_object).get("model") + return self._target( + decrypted_response_id(job.unified_object_id), stored_model if isinstance(stored_model, str) else None + ) + + async def _claim_for_billing(self, job: ManagedObjectRow) -> bool: + claimed: Final = await _managed_object_table(self.prisma_client).update_many( + # bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple + where={"id": job.id, "status": {"in": list(PENDING_RESPONSE_STATUSES)}}, + data={"status": "completed"}, + ) + return claimed > 0 + + async def _release_billing_claim(self, job: ManagedObjectRow) -> None: + await _managed_object_table(self.prisma_client).update_many( + where={"id": job.id, "status": "completed"}, + data={"status": job.status or PENDING_RESPONSE_STATUSES[0]}, + ) + + async def _billed_response( + self, + job: ManagedObjectRow, + target: _PollTarget, + response: ResponsesAPIResponse, + ) -> BilledBackgroundResponse: + deployment: Final = target.deployment + return BilledBackgroundResponse( + response=response.model_copy(update={"id": job.unified_object_id}), + model=deployment.litellm_params.model if deployment is not None else (response.model or ""), + custom_llm_provider=ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get( + "custom_llm_provider" + ), + 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)), + INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + **target.model_metadata, + **(_deployment_metadata(deployment) if deployment is not None else {}), + } + ), + ) + + async def _bill_once( + self, + job: ManagedObjectRow, + target: _PollTarget, + response: ResponsesAPIResponse, + ) -> BilledCharge | None: + if not await self._claim_for_billing(job): + 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)) + except BaseException: + await self._release_billing_claim(job) + raise + verbose_proxy_logger.info(f"Response {job.unified_object_id} reached a terminal status and was billed") + return charge + + async def _job_for_provider_id(self, provider_response_id: str) -> ManagedObjectRow | None: + return await _managed_object_table(self.prisma_client).find_first( + where={"model_object_id": provider_response_id, "file_purpose": "response"} + ) + + async def _rekey_by_provider_id(self, job: ManagedObjectRow) -> None: + await _managed_object_table(self.prisma_client).update_many( + where={"id": job.id}, data={"model_object_id": decrypted_response_id(job.model_object_id)} + ) + + async def _tracked_job(self, response_id: str) -> ManagedObjectRow | None: + if tracked := await self._job_for_provider_id(decrypted_response_id(response_id)): + return tracked + stored_by_issued_id: Final = await _managed_object_table(self.prisma_client).find_first( + where={"unified_object_id": response_id, "file_purpose": "response"} + ) + if stored_by_issued_id is not None and _keyed_by_issued_id(stored_by_issued_id): + await self._rekey_by_provider_id(stored_by_issued_id) + return stored_by_issued_id + + async def _superseded_by_a_provider_id_row(self, job: ManagedObjectRow) -> bool: + if not _keyed_by_issued_id(job): + return False + if await self._job_for_provider_id(decrypted_response_id(job.model_object_id)) is None: + await self._rekey_by_provider_id(job) + return False + await self._claim_for_billing(job) + return True + + async def _track_from_deployment( + self, response_id: str, track_untracked: UntrackedResponseTracker | None + ) -> BilledCharge | None: + target: Final = self._target(decrypted_response_id(response_id), None) + if track_untracked is None or target.deployment is None: + return None + return await self.bill_if_finished(response_id, await self._fetch_unbilled(target), track_untracked) + + async def _start_tracking( + self, response_id: str, response: ResponsesAPIResponse, track_untracked: UntrackedResponseTracker | None + ) -> ManagedObjectRow | None: + if track_untracked is None or not await track_untracked(response): + return None + return await self._tracked_job(response_id) + + async def bill_if_finished( + self, + response_id: str, + response: ResponsesAPIResponse, + track_untracked: UntrackedResponseTracker | None = None, + ) -> BilledCharge | None: + if response.status not in TERMINAL_RESPONSE_STATUSES: + return None + job: Final = await self._tracked_job(response_id) or await self._start_tracking( + response_id, response, track_untracked + ) + if job is None or job.status not in PENDING_RESPONSE_STATUSES: + return None + return await self._bill_once(job, self._poll_target(job), response) + + async def bill_from_deployment( + self, response_id: str, track_untracked: UntrackedResponseTracker | None = None + ) -> BilledCharge | None: + job: Final = await self._tracked_job(response_id) + if job is None: + return await self._track_from_deployment(response_id, track_untracked) + if job.status not in PENDING_RESPONSE_STATUSES: + return None + target: Final = self._poll_target(job) + response: Final = await self._fetch_unbilled(target) + if response.status not in TERMINAL_RESPONSE_STATUSES: + return None + return await self._bill_once(job, target, response) + async def check_responses_cost(self): """ Check if background responses are complete and track their cost. - Get all status="queued" or "in_progress" and file_purpose="response" jobs - - Query the provider to check if response is complete - - Cost is tracked by the get-responses call, billed because the poll is stamped - with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - - Mark responses in a terminal state as complete in the database + - Query the provider, unbilled, to check if the response is complete + - Claim each terminal job by moving its row to completed, then log the body that check fetched + as its one billed event, so two pollers racing over the same job bill it once - Mark responses the provider no longer has (404 through a resolved deployment) as stale_expired """ @@ -160,99 +427,53 @@ class CheckResponsesCost: jobs = await _managed_object_table(self.prisma_client).find_many( where={ - "status": {"in": ["queued", "in_progress"]}, + # bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple + "status": {"in": list(PENDING_RESPONSE_STATUSES)}, "file_purpose": "response", }, take=MAX_OBJECTS_PER_POLL_CYCLE, order={"created_at": "asc"}, ) - + verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check") - completed_jobs: Final[list[_ManagedObjectRow]] = [] - expired_jobs: Final[list[_ManagedObjectRow]] = [] + expired_jobs: Final[list[ManagedObjectRow]] = [] for job in jobs: unified_object_id = job.unified_object_id try: - from litellm.proxy.hooks.responses_id_security import ( - ResponsesIDSecurity, - ) - - # Get the stored response object to extract model information - stored_response = job.file_object - model_name = stored_response.get("model", None) - - # Decrypt the response ID - responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id) - - # Prepare metadata with model information for cost tracking - litellm_metadata = { - "user_api_key_user_id": job.created_by or "default-user-id", - INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, - } - - # Add model information if available - if model_name: - litellm_metadata["model"] = model_name - litellm_metadata["model_group"] = model_name # Use same value for model_group - via_router = self._resolve_deployment(responses_id_security) + if await self._superseded_by_a_provider_id_row(job): + verbose_proxy_logger.info(f"Response {unified_object_id} is billed under its provider id row") + continue + target = self._poll_target(job) except Exception as e: - verbose_proxy_logger.warning( - f"Skipping job {unified_object_id} due to error: {e}" - ) + verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}") continue - provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id( - responses_id_security - ).get("response_id", responses_id_security) + provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get( + "response_id", target.response_id + ) try: - response = await self._get_response( - response_id=responses_id_security, - litellm_metadata=litellm_metadata, - via_router=via_router, - ) + response = await self._fetch_unbilled(target) - verbose_proxy_logger.debug( - f"Response {unified_object_id} status: {response.status}, model: {model_name}" - ) + verbose_proxy_logger.debug(f"Response {unified_object_id} status: {response.status}") + if response.status in TERMINAL_RESPONSE_STATUSES: + await self._bill_once(job, target, response) except Exception as e: - if via_router and _is_response_gone_at_provider(e, provider_response_id): + if target.deployment is not None and _is_response_gone_at_provider(e, provider_response_id): verbose_proxy_logger.info( f"Response {unified_object_id} no longer available at provider (404), marking stale_expired: {e}" ) expired_jobs.append(job) continue - verbose_proxy_logger.warning( - f"Skipping job {unified_object_id} due to error: {e}" - ) + verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}") continue - if response.status in TERMINAL_RESPONSE_STATUSES: - verbose_proxy_logger.info( - f"Response {unified_object_id} has terminal status {response.status}, marking as complete" - ) - completed_jobs.append(job) - - # Mark completed jobs in the database - if len(completed_jobs) > 0: - await _managed_object_table(self.prisma_client).update_many( - # bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above - where={"id": {"in": [job.id for job in completed_jobs]}}, - data={"status": "completed"}, - ) - verbose_proxy_logger.info( - f"Marked {len(completed_jobs)} response jobs as completed" - ) - if len(expired_jobs) > 0: await _managed_object_table(self.prisma_client).update_many( # bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above where={"id": {"in": [job.id for job in expired_jobs]}}, data={"status": "stale_expired"}, ) - verbose_proxy_logger.info( - f"Marked {len(expired_jobs)} response jobs as stale_expired" - ) - + verbose_proxy_logger.info(f"Marked {len(expired_jobs)} response jobs as stale_expired") diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py new file mode 100644 index 00000000000..3c7735b53c0 --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -0,0 +1,218 @@ +"""Creator attribution for the spend row a managed-object cost poll writes. + +A batch or a background response is billed by a poll job that runs long after the creating +request returned, so the managed object row's creator columns are the only record of which +key, team, organization, user, end user, and tags carry the spend. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Protocol, cast + +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.constants import CLI_SESSION_KEY_PREFIX +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository + +if TYPE_CHECKING: + from prisma import models as prisma_models + + from litellm.proxy.utils import PrismaClient + from litellm.repositories.prisma_protocols import TableActions + + +class ManagedObjectRow(Protocol): + @property + def id(self) -> str: ... + + @property + def unified_object_id(self) -> str: ... + + @property + def model_object_id(self) -> str: ... + + @property + def created_by(self) -> str | None: ... + + @property + def status(self) -> str | None: ... + + @property + def file_object(self) -> object: ... + + @property + def org_id(self) -> str | None: ... + + @property + def api_key(self) -> str | None: ... + + @property + def team_id(self) -> str | None: ... + + @property + def request_tags(self) -> object: ... + + @property + def end_user_id(self) -> str | None: ... + + +_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object]) +_KEY_OWNER_SCOPE_PREFIX: Final = "key:" + + +def _string_tags(request_tags: object) -> list[str] | None: + if not isinstance(request_tags, list): + return None + return [tag for tag in _TAG_LIST_ADAPTER.validate_python(request_tags) if isinstance(tag, str)] or None + + +def _creator_user_id(created_by: str | None) -> str | None: + if created_by is None or created_by.startswith(_KEY_OWNER_SCOPE_PREFIX): + return None + return created_by + + +def _user_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_UserTable]: + return UserRepository(prisma_client).table + + +def _token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_VerificationToken]: + return VerificationTokenRepository(prisma_client).table + + +def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_TeamTable]: + return TeamRepository(prisma_client).table + + +@dataclass(frozen=True, slots=True) +class _CreatorRows: + key: prisma_models.LiteLLM_VerificationToken | None + team: prisma_models.LiteLLM_TeamTable | None + user: prisma_models.LiteLLM_UserTable | None + + +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, + user, aliases, and tags. + + Rows created before api_key and request_tags were persisted carry only created_by and + team_id, and fall back to those. A named creating key owns user_api_key_alias; when it + has no alias, or the key has since been rotated or deleted, the field keeps the creating + user's alias, because a resolvable name is more useful on the spend row than a null. + user_api_key_org_id is resolved too: the spend update writer reads it off this metadata + to increment organization spend, so leaving it out drops the cost from org accounting. + A background response also carries its end user and the key, team, and user budgets, the + way the read that billed it before did, and never names a key-scoped creator as its user. + """ + + def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None: + self._prisma_client: Final = prisma_client + self._poller_name: Final = poller_name + + async def batch_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]: + 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 {}), + **_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"), + **_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"), + **_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"), + } + + async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows: + return _CreatorRows( + key=await self._key_row(object_id, getattr(job, "api_key", None), job.created_by), + team=await self._team_row(object_id, getattr(job, "team_id", None)), + user=await self._user_row(object_id, _creator_user_id(job.created_by)), + ) + + async def _user_row(self, object_id: str, user_id: str | None) -> prisma_models.LiteLLM_UserTable | None: + """Objects created by a team or service account key carry no user id, and + find_unique(where={"user_id": None}) raises, so those resolve no row.""" + if not user_id: + return None + try: + return await _user_table(self._prisma_client).find_unique(where={"user_id": user_id}) + except Exception as e: + verbose_proxy_logger.error(f"{self._poller_name}: could not look up user {user_id} for {object_id}: {e}") + return None + + 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}) + 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: + return None + try: + return await _team_table(self._prisma_client).find_unique(where={"team_id": 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 + + +def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: + api_key: Final = getattr(job, "api_key", None) + team_id: Final = getattr(job, "team_id", None) + key_alias: Final = _key_alias(api_key, job.created_by, rows.key) + team_alias: Final = getattr(rows.team, "team_alias", None) + org_id: Final = getattr(job, "org_id", None) or _organization_id(rows.key) or _organization_id(rows.team) + tags: Final = _string_tags(getattr(job, "request_tags", None)) + return { + "user_api_key": api_key, + "user_api_key_hash": api_key, + "user_api_key_team_id": team_id, + **( + { + "user_api_key_user_email": getattr(rows.user, "user_email", None), + "user_api_key_alias": getattr(rows.user, "user_alias", None), + } + if rows.user is not None + else {} + ), + **({"user_api_key_alias": key_alias} if key_alias is not None else {}), + **({"user_api_key_team_alias": team_alias} if team_alias is not None else {}), + **({"user_api_key_org_id": org_id} if org_id is not None else {}), + **({"tags": tags} if tags else {}), + } + + +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}" + + +def _key_alias( + api_key: str | None, created_by: str | None, key_row: prisma_models.LiteLLM_VerificationToken | None +) -> str | None: + if api_key and _is_cli_session_key(api_key, created_by): + return api_key + return getattr(key_row, "key_alias", None) if key_row is not None else None + + +def _organization_id(row: object) -> str | None: + return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None + + +def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]: + if row is None: + return {} + return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)} diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 47fcf5d1b77..a33b33662ab 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -24,7 +24,9 @@ from uuid import NAMESPACE_URL, uuid5 import httpx from fastapi import HTTPException -from pydantic import BaseModel, ValidationError +from openai import APIConnectionError, AsyncOpenAI +from openai.types.file_deleted import FileDeleted +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, Unpack import litellm @@ -35,17 +37,17 @@ from litellm.caching.caching import DualCache from litellm.constants import ( BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS, BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, + MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, MAX_FILE_LIST_LIMIT, + PENDING_RESPONSE_STATUSES, ) from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR +from litellm.litellm_core_utils.internal_call_metadata import requests_background from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_metadata, ) -from openai import APIConnectionError, AsyncOpenAI -from openai.types.file_deleted import FileDeleted - from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.llms.base_llm.managed_resources.isolation import ( @@ -88,6 +90,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_metadata, ) from litellm.repositories.managed_file_repository import ManagedFileRepository +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue] AllMessageValues, AsyncCursorPage, @@ -100,6 +103,7 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess ResponsesAPIResponse, ) from litellm.types.utils import ( + ADDRESSED_RESPONSE_ID_FIELD, CallTypesLiteral, LiteLLMBatch, LiteLLMFineTuningJob, @@ -346,6 +350,8 @@ class _ManagedObjectTableActions(Protocol): async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + async def create(self, data: Mapping[str, object]) -> "PrismaManagedObjectRow": ... + class _ManagedResourceDatabase(Protocol): @property @@ -392,6 +398,39 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag return prisma_client.db.litellm_managedobjecttable +_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00" + + +def _storable_text(text: str) -> str: + return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "") + + +def _storable_json_value(value: JsonValue) -> JsonValue: + if isinstance(value, str): + return _storable_text(value) + if isinstance(value, list): + return [_storable_json_value(item) for item in value] + if isinstance(value, dict): + return {_storable_text(key): _storable_json_value(item) for key, item in value.items()} + return value + + +def _postgres_storable_json(file_object: BaseModel) -> str: + dumped: Final = file_object.model_dump_json() + if "\\u0000" not in dumped: + return dumped + return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode() + + +def _without_model_output( + file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse, +) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse: + return ( + file_object.model_copy(update={"output": []}) if isinstance(file_object, ResponsesAPIResponse) else file_object + ) + + def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]: hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {} @@ -405,6 +444,40 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s ) +_MANAGED_RESPONSE_READ_CALL_TYPES: Final = frozenset({"aget_responses", "acancel_responses"}) +_REQUEST_FIELDS: Final = TypeAdapter(dict[str, object]) + + +def _addresses_a_proxy_issued_response(data: Mapping[str, object]) -> bool: + from litellm.proxy.hooks.responses_id_security import decrypted_response_id + + addressed_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) or data.get("response_id") + return ( + isinstance(addressed_id, str) + and ResponsesAPIRequestUtils.get_model_id_from_response_id(decrypted_response_id(addressed_id)) is not None + ) + + +def _is_streamed_background_create(request: Mapping[str, object]) -> bool: + return ( + request.get("stream") is True + and isinstance(request.get("litellm_metadata"), dict) + and requests_background(request) + ) + + +def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str, object]: + metadata_field: Final = "litellm_metadata" if isinstance(request.get("litellm_metadata"), dict) else "metadata" + existing_metadata: Final = request.get(metadata_field) + return { + **request, + metadata_field: { + **(_REQUEST_FIELDS.validate_python(existing_metadata) if isinstance(existing_metadata, dict) else {}), + MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True, + }, + } + + _MANAGED_FILES_TARGET: Final = "managed_files" _PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0) @@ -529,11 +602,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) -> None: """Persist a managed object row, caching it and upserting it in the DB. - persist_attribution is set only by the batch create, which is the one caller - that can speak for the creator; it gates the api_key and request_tags columns - that CheckBatchCost bills against, so a later poll or retrieve of the same - batch cannot record itself as the paying key. Like created_by and team_id, - both are written only in the upsert create branch, never on update. + persist_attribution is set only by the batch and background response creates, + the callers that can speak for the creator; it gates the api_key, request_tags, and, + for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or + retrieve of the same object cannot record itself as the paying key. Like created_by and team_id, + they are written only in the upsert create branch, never on update. + + A stored response keeps no model output: billing reads only its status and model from + the row, so the row never holds generated content. create_if_missing is cleared by callers that observe a batch they did not create, such as a poll. They still refresh status and file_object, but a @@ -546,11 +622,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): It is written only in the upsert create branch. """ verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") + stored_object: Final = _without_model_output(file_object) litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, file_purpose=file_purpose, - file_object=file_object, + file_object=stored_object, ) await self.internal_usage_cache.async_set_cache( key=unified_object_id, @@ -561,32 +638,35 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): from prisma import Json api_key = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None + end_user_id: Final = user_api_key_dict.end_user_id if file_purpose == "response" else None attribution_columns = ( { **({"api_key": api_key} if api_key is not None else {}), - **({"request_tags": Json(list(request_tags))} if request_tags else {}), + **({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}), + **({"end_user_id": end_user_id} if end_user_id else {}), } if persist_attribution else {} ) # FIX: Update status and file_object on every operation to keep state in sync update_columns: Final = { - "file_object": file_object.model_dump_json(), + "file_object": _postgres_storable_json(stored_object), "status": file_object.status, "updated_by": user_api_key_dict.user_id, } + managed_object_table: Final = _managed_object_table(self.prisma_client) if not create_if_missing: - await _managed_object_table(self.prisma_client).update_many( + await managed_object_table.update_many( where={"unified_object_id": unified_object_id}, data=update_columns, ) return - await _managed_object_table(self.prisma_client).upsert( + await managed_object_table.upsert( where={"unified_object_id": unified_object_id}, data={ "create": { "unified_object_id": unified_object_id, - "file_object": file_object.model_dump_json(), + "file_object": _postgres_storable_json(stored_object), "model_object_id": model_object_id, "file_purpose": file_purpose, "created_by": resolve_resource_owner_id(user_api_key_dict), @@ -601,6 +681,42 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): }, ) + async def track_untracked_background_response( + self, + response_id: str, + response: ResponsesAPIResponse, + user_api_key_dict: UserAPIKeyAuth, + request_tags: Sequence[str] | None = None, + ) -> bool: + if not _addresses_a_proxy_issued_response({"response_id": response_id}): + return False + from prisma import Json + from prisma.errors import UniqueViolationError + + from litellm.proxy.hooks.responses_id_security import decrypted_response_id + + api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None + try: + await _managed_object_table(self.prisma_client).create( + data={ + "unified_object_id": response_id, + "model_object_id": decrypted_response_id(response_id), + "file_purpose": "response", + "file_object": _postgres_storable_json(_without_model_output(response)), + "status": PENDING_RESPONSE_STATUSES[0], + "created_by": resolve_resource_owner_id(user_api_key_dict), + "team_id": user_api_key_dict.team_id, + "org_id": await self._resolve_creator_org_id(user_api_key_dict), + "updated_by": user_api_key_dict.user_id, + **({"api_key": api_key} if api_key is not None else {}), + **({"end_user_id": user_api_key_dict.end_user_id} if user_api_key_dict.end_user_id else {}), + **({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}), + } + ) + except UniqueViolationError: + return False + return True + @with_service_target(_MANAGED_FILES_TARGET) async def get_unified_file_id( self, file_id: str, litellm_parent_otel_span: Optional[Span] = None @@ -1375,6 +1491,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools) if unified_vs_ids: await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict) + create_request: Final = _REQUEST_FIELDS.validate_python(data) + if _is_streamed_background_create(create_request): + return _with_managed_background_marker(create_request) elif call_type == CallTypes.afile_content.value: retrieve_file_id = cast(Optional[str], data.get("file_id")) potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False @@ -1432,6 +1551,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id) elif retrieve_object_id and accessor_key == "batch_id": await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict) + elif call_type in _MANAGED_RESPONSE_READ_CALL_TYPES: + request: Final = _REQUEST_FIELDS.validate_python(data) + if _addresses_a_proxy_issued_response(request): + return _with_managed_background_marker(request) elif call_type == CallTypes.acreate_fine_tuning_job.value: fine_tuning_file_ids: Final = tuple( file_id diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261015000000_managed_object_end_user_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261015000000_managed_object_end_user_id/migration.sql new file mode 100644 index 00000000000..35dc554953c --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261015000000_managed_object_end_user_id/migration.sql @@ -0,0 +1,4 @@ +SET LOCAL lock_timeout = '10s'; + +-- AlterTable +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "end_user_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 12438614755..d40bac380aa 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") + end_user_id String? updated_at DateTime @updatedAt updated_by String? diff --git a/litellm/constants.py b/litellm/constants.py index 6e1b5baad13..5b2bbaf95c1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1622,6 +1622,7 @@ CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" +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" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" @@ -1862,6 +1863,8 @@ RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900 PROXY_BATCH_POLLING_INTERVAL: Final = get_env_int("PROXY_BATCH_POLLING_INTERVAL", 3600) MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, get_env_int("MAX_OBJECTS_PER_POLL_CYCLE", 50)) MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)) +PENDING_RESPONSE_STATUSES: Final = ("queued", "in_progress") +TERMINAL_RESPONSE_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "incomplete"}) BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0 BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60 STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, get_env_int("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000)) diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index d844cbae367..c5e57bad050 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -21,9 +21,17 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES +from litellm.constants import ( + INTERNAL_CALL_ORIGIN_METADATA_KEY, + MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, + NON_INFERENCE_CALL_TYPES, +) from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params -from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin +from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, + InternalCallOrigin, +) BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) @@ -62,10 +70,32 @@ def is_background_response(response: object) -> bool: because the callers hold it in both shapes. """ if isinstance(response, Mapping): - return response.get("background") is True + return "background" in response and response["background"] is True return getattr(response, "background", None) is True +def requests_background(request_body: Mapping[str, object]) -> bool: + return bool(request_body.get("background")) or is_background_response(request_body.get("extra_body")) + + +_RESPONSE_CREATE_CALL_TYPES: Final = frozenset({"responses", "aresponses"}) + + +def _finished_response(response: object) -> object: + wrapped: Final = response.get("response") if isinstance(response, Mapping) else getattr(response, "response", None) + return response if wrapped is None else wrapped + + +def _may_be_unbilled(call_type: str | None, response: object) -> bool: + if call_type in _RESPONSE_CREATE_CALL_TYPES: + return is_background_response(_finished_response(response)) + return call_type in NON_INFERENCE_CALL_TYPES + + +def _is_managed_background_call(metadata: Mapping[str, object] | None) -> bool: + return metadata is not None and metadata.get(MANAGED_BACKGROUND_RESPONSE_METADATA_KEY) is True + + def is_unbilled_non_inference_call( call_type: str | None, metadata: Mapping[str, object] | None, @@ -74,18 +104,25 @@ def is_unbilled_non_inference_call( """A read/management route priced at zero, because the usage it reports belongs to the call that created the object it just read. - Retrieving a background response is the exception, and the enterprise cost poller's read - is the same exception seen from the other side: that job's create billed nothing, so its - retrieval is the only place the spend is ever visible. Pricing those at zero would lose - the spend rather than deduplicate it. + A background response is the exception, since its create billed nothing. When the proxy + tracks it as a managed object, the read is marked with + ``MANAGED_BACKGROUND_RESPONSE_METADATA_KEY`` and stays free, because the one billing event + the proxy logs from the finished response, stamped with + ``BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN``, bills the job. A streamed background create + the proxy tracks carries the same marker and is free for the same reason. Any other + read of a background response (the SDK, a proxy without a database) is the only place its + spend is ever visible, so it prices normally. """ - if call_type not in NON_INFERENCE_CALL_TYPES: + if not _may_be_unbilled(call_type, response): return False - if is_background_response(response): + if call_type in _RESPONSE_CREATE_CALL_TYPES: + return _is_managed_background_call(metadata) + origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) if metadata is not None else None + if origin == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: return False - if metadata is None: + if not is_background_response(response): return True - return metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + return origin == BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN or _is_managed_background_call(metadata) def is_unbilled_non_inference_call_from_params( @@ -95,10 +132,10 @@ def is_unbilled_non_inference_call_from_params( ) -> bool: """:func:`is_unbilled_non_inference_call` for callers holding raw ``litellm_params``. - The call-type membership test runs first so that inference traffic, which is every + The call-type and background tests run first so that inference traffic, which is every request in a normal workload, never pays for the metadata merge behind it. """ - if call_type not in NON_INFERENCE_CALL_TYPES: + if not _may_be_unbilled(call_type, response): return False from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9b10b95956f..3d6f8616805 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5856,7 +5856,7 @@ class StandardLoggingPayloadSetup: return messages @staticmethod - def merge_litellm_metadata(litellm_params: Mapping[str, object]) -> dict: + def merge_litellm_metadata(litellm_params: Mapping[str, object]) -> dict[str, object]: """ Merge both litellm_metadata and metadata from litellm_params. @@ -6483,7 +6483,7 @@ class StandardLoggingPayloadSetup: @staticmethod def get_request_tags( litellm_params: dict[str, object], - proxy_server_request: dict[str, object], + proxy_server_request: Mapping[str, object], ) -> list[str]: metadata: Final = cast( # cast-ok: request metadata is caller-provided Mapping[str, object], litellm_params.get("metadata") or {} diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index e4fe88690f5..6e972047712 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -674,6 +674,15 @@ def is_request_body_safe( return True +def carries_client_side_overrides(request_body: object) -> bool: + if not isinstance(request_body, dict): + return False + try: + return not is_request_body_safe(request_body=request_body, general_settings={}, llm_router=None, model="") + except ValueError: + return True + + def _coerce_metadata_to_dict(value: object) -> dict[str, object] | None: """Return ``value`` as a dict, parsing it from JSON if delivered as a string. diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 104a6452f79..8bf39b7a17a 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -12,7 +12,6 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response -from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -68,7 +67,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # update_batch_in_database, validate_managed_id_requirement, ) -from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_metadata +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_request_data from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy, is_known_model from litellm.repositories.managed_batch_repository import ManagedBatchRepository @@ -81,7 +80,6 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedObjectTable router: Final = APIRouter() -_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) def _hidden_param_string(hidden_params: Mapping[str, object], key: str) -> str: @@ -89,13 +87,6 @@ def _hidden_param_string(hidden_params: Mapping[str, object], key: str) -> str: return value if isinstance(value, str) else "" -def _request_tags(data: Mapping[str, object]) -> tuple[str, ...] | None: - metadata: Final = data.get("litellm_metadata") - if metadata is None: - return None - return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata)) - - def _require_batch_response(response: LLMResponseTypes) -> LiteLLMBatch: if not isinstance(response, LiteLLMBatch): raise TypeError("Batch endpoint received a non-batch response") @@ -444,7 +435,7 @@ async def create_batch( model=model, provider=executed_provider, user_api_key_dict=user_api_key_dict, - request_tags=_request_tags(_create_batch_data), + request_tags=request_tags_from_request_data(_create_batch_data), ) if executed_provider is not None else await _create_provider_batch_for_managed_file( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e0d6538a545..55cc13ff97d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1605,6 +1605,31 @@ def _get_cost_breakdown_from_logging_obj( ) +def _header_value(value: float | None) -> str | None: + return str(value) if value is not None else None + + +def cost_headers( + response_cost: float | str | None, litellm_logging_obj: LiteLLMLoggingObj | None +) -> Mapping[str, str | None]: + cost_breakdown: Final = _get_cost_breakdown_from_logging_obj( + litellm_logging_obj=litellm_logging_obj, response_cost=response_cost + ) + return { + "x-litellm-response-cost": str(response_cost), + "x-litellm-response-cost-original": _header_value(cost_breakdown.original_cost), + "x-litellm-response-cost-discount-amount": _header_value(cost_breakdown.discount_amount), + "x-litellm-response-cost-margin-amount": _header_value(cost_breakdown.margin_total_amount), + "x-litellm-response-cost-margin-percent": _header_value(cost_breakdown.margin_percent), + "x-litellm-response-cost-input": _header_value(cost_breakdown.input_cost), + "x-litellm-response-cost-output": _header_value(cost_breakdown.output_cost), + "x-litellm-response-cost-cache-read": _header_value(cost_breakdown.cache_read_cost), + "x-litellm-response-cost-cache-creation": _header_value(cost_breakdown.cache_creation_cost), + "x-litellm-response-cost-reasoning": _header_value(cost_breakdown.reasoning_cost), + "x-litellm-response-cost-tool-usage": _header_value(cost_breakdown.tool_usage_cost), + } + + def _classifier_cost_from_request_data(request_data: Mapping[str, object] | None) -> float | None: """Cost of the auto-router's LLM classifier call, read from the request's routing_decision. @@ -1878,10 +1903,6 @@ class ProxyBaseLLMRequestProcessing: use_logging_obj=read_timing_from_logging_obj, ) - cost_breakdown: Final = _get_cost_breakdown_from_logging_obj( - litellm_logging_obj=litellm_logging_obj, response_cost=response_cost - ) - # Calculate updated spend for header (include current response_cost) current_spend: Final = user_api_key_dict.spend or 0.0 updated_spend = current_spend @@ -1908,37 +1929,7 @@ class ProxyBaseLLMRequestProcessing: ), # don't include query params, risk of leaking sensitive info "x-litellm-version": version, "x-litellm-model-region": model_region, - "x-litellm-response-cost": str(response_cost), - "x-litellm-response-cost-original": ( - str(cost_breakdown.original_cost) if cost_breakdown.original_cost is not None else None - ), - "x-litellm-response-cost-discount-amount": ( - str(cost_breakdown.discount_amount) if cost_breakdown.discount_amount is not None else None - ), - "x-litellm-response-cost-margin-amount": ( - str(cost_breakdown.margin_total_amount) if cost_breakdown.margin_total_amount is not None else None - ), - "x-litellm-response-cost-margin-percent": ( - str(cost_breakdown.margin_percent) if cost_breakdown.margin_percent is not None else None - ), - "x-litellm-response-cost-input": ( - str(cost_breakdown.input_cost) if cost_breakdown.input_cost is not None else None - ), - "x-litellm-response-cost-output": ( - str(cost_breakdown.output_cost) if cost_breakdown.output_cost is not None else None - ), - "x-litellm-response-cost-cache-read": ( - str(cost_breakdown.cache_read_cost) if cost_breakdown.cache_read_cost is not None else None - ), - "x-litellm-response-cost-cache-creation": ( - str(cost_breakdown.cache_creation_cost) if cost_breakdown.cache_creation_cost is not None else None - ), - "x-litellm-response-cost-reasoning": ( - str(cost_breakdown.reasoning_cost) if cost_breakdown.reasoning_cost is not None else None - ), - "x-litellm-response-cost-tool-usage": ( - str(cost_breakdown.tool_usage_cost) if cost_breakdown.tool_usage_cost is not None else None - ), + **cost_headers(response_cost, litellm_logging_obj), "x-litellm-classifier-cost": (str(classifier_cost) if classifier_cost is not None else None), "x-litellm-key-tpm-limit": str(user_api_key_dict.tpm_limit), "x-litellm-key-rpm-limit": str(user_api_key_dict.rpm_limit), @@ -3164,7 +3155,7 @@ class ProxyBaseLLMRequestProcessing: ) @staticmethod - def _extract_completed_responses_response(stream_response: object) -> object: + def extract_completed_responses_response(stream_response: object) -> object: """Pull the assembled ``ResponsesAPIResponse`` off a streaming iterator. ``ResponsesAPIStreamingIterator`` stores the terminal stream event @@ -3200,9 +3191,7 @@ class ProxyBaseLLMRequestProcessing: appears, ahead of the chunk carrying ``response.completed``. """ async for chunk in wrapped_generator: - completed_obj = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( - original_stream_response - ) + completed_obj = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(original_stream_response) if completed_obj is None: yield chunk continue @@ -3214,7 +3203,7 @@ class ProxyBaseLLMRequestProcessing: async for remaining_chunk in wrapped_generator: yield remaining_chunk return - late_completed_obj: Final = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( + late_completed_obj: Final = ProxyBaseLLMRequestProcessing.extract_completed_responses_response( original_stream_response ) if late_completed_obj is not None: diff --git a/litellm/proxy/db/model_usage_rollup.py b/litellm/proxy/db/model_usage_rollup.py index d83500d66e8..4566b10b260 100644 --- a/litellm/proxy/db/model_usage_rollup.py +++ b/litellm/proxy/db/model_usage_rollup.py @@ -19,6 +19,7 @@ from litellm.constants import ( from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES, SpendLogsPayload from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL, rollup_lock_timeout_setting +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -59,6 +60,7 @@ class ModelUsageTransaction: prompt_tokens: int completion_tokens: int successful: bool + counts_as_request: bool = True def model_usage_task_type(request_tags: str) -> str: @@ -78,12 +80,12 @@ def model_usage_task_type(request_tags: str) -> str: ) -def _is_internal_call(metadata: str) -> bool: +def _internal_call_origin(metadata: str) -> object: try: decoded: Final = _METADATA.validate_json(metadata) except ValidationError: - return False - return bool(decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) + return None + return decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) def _date_from_start_time(start_time: datetime | str) -> str | None: @@ -94,7 +96,8 @@ def _date_from_start_time(start_time: datetime | str) -> str | None: def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransaction | None: date: Final = _date_from_start_time(payload["startTime"]) - if date is None or _is_internal_call(payload["metadata"]): + origin: Final = _internal_call_origin(payload["metadata"]) + if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN): return None model: Final = payload["model"] or "unknown" return ModelUsageTransaction( @@ -109,6 +112,7 @@ def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransa prompt_tokens=payload["prompt_tokens"], completion_tokens=payload["completion_tokens"], successful=payload["status"] == "success", + counts_as_request=not origin, ) @@ -142,8 +146,9 @@ async def flush_model_usage_transactions( spend = sum(entry.spend for entry in entries) prompt_tokens = sum(entry.prompt_tokens for entry in entries) completion_tokens = sum(entry.completion_tokens for entry in entries) - successful = sum(1 for entry in entries if entry.successful) - failed = len(entries) - successful + requests = tuple(entry for entry in entries if entry.counts_as_request) + successful = sum(1 for entry in requests if entry.successful) + failed = len(requests) - successful key_fields = { "date": key.date, "model_group": key.model_group, @@ -159,7 +164,7 @@ async def flush_model_usage_transactions( "spend": spend, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, - "request_count": len(entries), + "request_count": len(requests), "successful_requests": successful, "failed_requests": failed, }, @@ -167,7 +172,7 @@ async def flush_model_usage_transactions( "spend": {"increment": spend}, "prompt_tokens": {"increment": prompt_tokens}, "completion_tokens": {"increment": completion_tokens}, - "request_count": {"increment": len(entries)}, + "request_count": {"increment": len(requests)}, "successful_requests": {"increment": successful}, "failed_requests": {"increment": failed}, }, diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index c1384565a80..27b5924c540 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -140,12 +140,12 @@ class ResponsesIDSecurity(CustomLogger): ) if not isinstance(addressed_id, str) or not addressed_id: return data - authorized_id: Final = self._authorize_response_id(addressed_id, user_api_key_dict) + authorized_id: Final = self.authorize_response_id(addressed_id, user_api_key_dict) data[addressed_id_field] = authorized_id data[ADDRESSED_RESPONSE_ID_FIELD] = addressed_id return data - def _authorize_response_id( + def authorize_response_id( self, response_id: str, user_api_key_dict: "UserAPIKeyAuth", @@ -344,3 +344,9 @@ class ResponsesIDSecurity(CustomLogger): ): chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) yield chunk + + +def decrypted_response_id(response_id: str) -> str: + """The id a client-facing response id stands for, the same for every encryption of one response.""" + original_response_id, _, _ = ResponsesIDSecurity()._decrypt_response_id(response_id) + return original_response_id diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index bfda00aa39f..ae381775391 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -31,6 +31,7 @@ from litellm.constants import ( CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, + MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, @@ -378,6 +379,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( CONSUMED_REQUEST_TAGS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, + MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, "standard_logging_object", "litellm_roi_estimator", "proxy_server_request", diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py index 880fdad92bf..cc3b0c429ff 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py @@ -9,6 +9,8 @@ import asyncio from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import strip_null_bytes @@ -45,6 +47,17 @@ def request_tags_from_metadata(request_metadata: Mapping[str, object]) -> tuple[ return None +_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def request_tags_from_request_data(data: Mapping[str, object]) -> tuple[str, ...] | None: + """Tags of a request whose metadata travels as litellm_metadata (batches, responses)""" + metadata: Final = data.get("litellm_metadata") + if metadata is None: + return None + return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata)) + + def log_batch_registration_result( finished: asyncio.Task[None], provider: str, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 7f9fcca369b..7c52499bd6f 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,11 +1,23 @@ import asyncio import contextlib import json -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from enum import Enum from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeAlias, cast, get_args +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NamedTuple, + Protocol, + TypeAlias, + cast, + get_args, + runtime_checkable, +) import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -17,25 +29,32 @@ from starlette.websockets import WebSocket, WebSocketDisconnect from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.constants import EMPTY_MAPPING +from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, build_blocked_response, ) from litellm.proxy._types import * +from litellm.proxy.auth.auth_utils import carries_client_side_overrides from litellm.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, user_api_key_auth, user_api_key_auth_websocket, ) -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, create_response +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, cost_headers, create_response from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_set_request_parsed_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, safe_set_request_parsed_body, ) +from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( + request_tags_from_request_data, +) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -47,10 +66,13 @@ from litellm.types.responses.main import DeleteResponseResult from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.utils import ProxyLogging from litellm.router import Router router: Final = APIRouter() _RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None) +_PROXY_SERVER_REQUEST_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) _ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg @@ -344,7 +366,7 @@ async def responses_api( asyncio.create_task( background_streaming_task( polling_id=polling_id, - data=data.copy(), + data=dict(_without_managed_background_marker(data)), polling_handler=polling_handler, request=request, fastapi_response=fastapi_response, @@ -379,7 +401,11 @@ async def responses_api( llm_router=llm_router, general_settings=general_settings, proxy_config=proxy_config, - select_data_generator=native_data_generator, + select_data_generator=( + partial(_streamed_background_data_generator, native_data_generator, proxy_logging_obj) + if requests_background(_PROXY_SERVER_REQUEST_ADAPTER.validate_python(data)) + else native_data_generator + ), model=None, user_model=user_model, user_temperature=user_temperature, @@ -389,49 +415,38 @@ async def responses_api( version=version, ) - # Store in managed objects table if background mode is enabled - if data.get("background") and isinstance(response, ResponsesAPIResponse): - if response.status in ["queued", "in_progress"]: - from litellm_enterprise.proxy.hooks.managed_files import ( - PROXY_LiteLLMManagedFiles, - ) + if isinstance(response, ResponsesAPIResponse) and (data.get("background") or is_background_response(response)): + managed_files_obj: Final = _managed_files_hook(proxy_logging_obj) - managed_files_obj: Final = cast( - PROXY_LiteLLMManagedFiles | None, - proxy_logging_obj.get_proxy_hook("managed_files"), - ) + if managed_files_obj and llm_router: + try: + hidden_params: Final = getattr(response, "_hidden_params", {}) or {} + model_id: Final = hidden_params.get("model_id", None) - if managed_files_obj and llm_router: - try: - # Get the actual deployment model_id from hidden params - hidden_params: Final = getattr(response, "_hidden_params", {}) or {} - model_id: Final = hidden_params.get("model_id", None) - - if not model_id: - verbose_proxy_logger.warning( - "No model_id found in response hidden params for response %s, skipping managed object storage", - response.id, - ) - raise Exception("No model_id found in response hidden params") - # Store in managed objects table - await managed_files_obj.store_unified_object_id( - unified_object_id=response.id, - file_object=response, - litellm_parent_otel_span=None, - model_object_id=response.id, - file_purpose="response", - user_api_key_dict=user_api_key_dict, - ) - - verbose_proxy_logger.info( - "Stored background response %s in managed objects table with unified_id=%s", - response.id, + if not model_id: + verbose_proxy_logger.warning( + "No model_id found in response hidden params for response %s, skipping managed object storage", response.id, ) - except Exception as e: - verbose_proxy_logger.error( - "Failed to store background response in managed objects table: %s", e - ) + raise Exception("No model_id found in response hidden params") + await managed_files_obj.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=None, + model_object_id=decrypted_response_id(response.id), + file_purpose="response", + user_api_key_dict=user_api_key_dict, + request_tags=_background_response_spend_tags(processor.data), + persist_attribution=True, + ) + + verbose_proxy_logger.info( + "Stored background response %s in managed objects table with unified_id=%s", + response.id, + response.id, + ) + except Exception as e: + verbose_proxy_logger.error("Failed to store background response in managed objects table: %s", e) return response except ModifyResponseException as e: @@ -764,10 +779,11 @@ async def get_response( # Normal provider response flow data: Final = await read_request_body(request=request) + read_settles_billing: Final = not carries_client_side_overrides(data) data["response_id"] = response_id processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response: Final[object] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -792,6 +808,266 @@ async def get_response( proxy_logging_obj=proxy_logging_obj, version=version, ) + await _bill_finished_background_response( + response=response, + response_id=response_id, + user_api_key_dict=user_api_key_dict, + fastapi_response=fastapi_response, + read_settles_billing=read_settles_billing, + ) + return response + + +def _background_response_spend_tags(data: Mapping[str, object]) -> tuple[str, ...] | None: + proxy_server_request: Final = data.get("proxy_server_request") + client_tags: Final = ( + StandardLoggingPayloadSetup.get_request_tags( + {}, _PROXY_SERVER_REQUEST_ADAPTER.validate_python(proxy_server_request) + ) + if isinstance(proxy_server_request, dict) + else [] + ) + tags: Final = (*(request_tags_from_request_data(data) or ()), *client_tags) + return tags or None + + +@runtime_checkable +class _BackgroundResponseStore(Protocol): + async def store_unified_object_id( + self, + unified_object_id: str, + file_object: ResponsesAPIResponse, + litellm_parent_otel_span: None, + model_object_id: str, + file_purpose: Literal["response"], + user_api_key_dict: UserAPIKeyAuth, + request_tags: Sequence[str] | None = None, + persist_attribution: bool = False, + ) -> None: ... + + async def track_untracked_background_response( + self, + response_id: str, + response: ResponsesAPIResponse, + user_api_key_dict: UserAPIKeyAuth, + request_tags: Sequence[str] | None = None, + ) -> bool: ... + + +def _managed_files_hook(proxy_logging_obj: "ProxyLogging") -> _BackgroundResponseStore | None: + hook: Final = proxy_logging_obj.get_proxy_hook("managed_files") + return hook if isinstance(hook, _BackgroundResponseStore) else None + + +async def _bill_finished_background_response( + response: object, + response_id: str, + user_api_key_dict: UserAPIKeyAuth, + fastapi_response: Response, + read_settles_billing: bool, +) -> None: + from litellm.proxy.proxy_server import proxy_logging_obj + + if not isinstance(response, ResponsesAPIResponse) or not is_background_response(response): + return + if response.status not in TERMINAL_RESPONSE_STATUSES: + return + managed_files_hook: Final = _managed_files_hook(proxy_logging_obj) + track_untracked: Final = ( + partial( + managed_files_hook.track_untracked_background_response, + response_id, + user_api_key_dict=user_api_key_dict, + ) + if managed_files_hook is not None + else None + ) + charge: Final = await _settle_background_billing( + (lambda cost_check: cost_check.bill_if_finished(response_id, response, track_untracked)) + if read_settles_billing + else (lambda cost_check: cost_check.bill_from_deployment(response_id, track_untracked)) + ) + _report_charge(charge, fastapi_response) + + +class _BilledCharge(Protocol): + @property + def cost(self) -> float | None: ... + + @property + def logging_obj(self) -> "LiteLLMLoggingObj": ... + + +_UntrackedResponseTracker: TypeAlias = Callable[[ResponsesAPIResponse], Awaitable[bool]] + + +class _BackgroundCostCheck(Protocol): + async def bill_from_deployment( + self, response_id: str, track_untracked: _UntrackedResponseTracker | None = None + ) -> _BilledCharge | None: ... + + async def bill_if_finished( + self, + response_id: str, + response: ResponsesAPIResponse, + track_untracked: _UntrackedResponseTracker | None = None, + ) -> _BilledCharge | None: ... + + +def _background_cost_check() -> _BackgroundCostCheck | None: + from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_logging_obj + + if prisma_client is None or llm_router is None: + return None + try: + from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost + except ImportError: + return None + + return CheckResponsesCost( # pyright: ignore[reportUnknownVariableType] # litellm_enterprise is absent from the type-check venv + proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, llm_router=llm_router + ) + + +_BillingSettlement: TypeAlias = Callable[[_BackgroundCostCheck], Awaitable[_BilledCharge | None]] +_BILLING_RETRY_DELAYS_SECONDS: Final = (5.0, 30.0, 120.0) +_billing_retries: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: the event loop keeps only weak task refs + + +@dataclass(frozen=True, slots=True) +class _BillingFailed: + error: Exception + + +async def _attempt_billing( + settle: _BillingSettlement, cost_check: _BackgroundCostCheck +) -> _BilledCharge | None | _BillingFailed: + try: + return await settle(cost_check) + except Exception as e: + return _BillingFailed(e) + + +async def _retry_background_billing( + settle: _BillingSettlement, cost_check: _BackgroundCostCheck, retry_delays: Sequence[float] +) -> None: + for delay in retry_delays: + await asyncio.sleep(delay) + outcome = await _attempt_billing(settle, cost_check) + if not isinstance(outcome, _BillingFailed): + return + verbose_proxy_logger.warning("Retrying the bill for a finished background response failed: %s", outcome.error) + verbose_proxy_logger.error("Gave up billing a finished background response; the cost poll bills it when enabled") + + +def _schedule_billing_retry( + settle: _BillingSettlement, cost_check: _BackgroundCostCheck, retry_delays: Sequence[float] +) -> None: + retry: Final = asyncio.create_task(_retry_background_billing(settle, cost_check, retry_delays)) + _billing_retries.add(retry) + retry.add_done_callback(_billing_retries.discard) + + +async def _settle_background_billing( + settle: _BillingSettlement, retry_delays: Sequence[float] = _BILLING_RETRY_DELAYS_SECONDS +) -> _BilledCharge | None: + cost_check: Final = _background_cost_check() + if cost_check is None: + return None + outcome: Final = await _attempt_billing(settle, cost_check) + if not isinstance(outcome, _BillingFailed): + return outcome + verbose_proxy_logger.warning( + "Billing a finished background response failed, retrying in background: %s", outcome.error + ) + _schedule_billing_retry(settle, cost_check, retry_delays) + return None + + +def _report_charge(charge: _BilledCharge | None, fastapi_response: Response) -> None: + if charge is None or charge.cost is None: + return + billed_headers: Final = cost_headers(charge.cost, charge.logging_obj) + fastapi_response.headers.update({name: value for name, value in billed_headers.items() if value is not None}) + + +def _caller_may_address_response( + response_id: str, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: "ProxyLogging" +) -> bool: + security_hook: Final = proxy_logging_obj.get_proxy_hook("responses_id_security") + if not isinstance(security_hook, ResponsesIDSecurity): + return True + try: + security_hook.authorize_response_id(response_id, user_api_key_dict) + except HTTPException: + return False + return True + + +def _without_managed_background_marker(request_data: Mapping[str, object]) -> Mapping[str, object]: + metadata: Final = request_data.get("litellm_metadata") + if not isinstance(metadata, Mapping): + return dict(request_data) + return { + **request_data, + "litellm_metadata": { + key: value + for key, value in _PROXY_SERVER_REQUEST_ADAPTER.validate_python(metadata).items() + if key != MANAGED_BACKGROUND_RESPONSE_METADATA_KEY + }, + } + + +def _streamed_background_data_generator( + select_data_generator: Callable[..., AsyncIterator[str | bytes]], + proxy_logging_obj: "ProxyLogging", + response: object, + user_api_key_dict: UserAPIKeyAuth, + request_data: dict[str, object], + request: Request | None = None, +) -> AsyncIterator[str | bytes]: + chunks: Final = select_data_generator( + response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, request=request + ) + managed_files_hook: Final = _managed_files_hook(proxy_logging_obj) + if managed_files_hook is None: + return chunks + track_untracked: Final = partial( + managed_files_hook.track_untracked_background_response, + user_api_key_dict=user_api_key_dict, + request_tags=_background_response_spend_tags(request_data), + ) + return _bill_streamed_background_response(response, chunks, track_untracked) + + +async def _bill_streamed_background_response( + stream: object, + chunks: AsyncIterator[str | bytes], + track_untracked: Callable[[str, ResponsesAPIResponse], Awaitable[bool]], +) -> AsyncIterator[str | bytes]: + async for chunk in chunks: + completed: object = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream) + if not isinstance(completed, ResponsesAPIResponse): + yield chunk + continue + await _bill_finished_stream(completed, track_untracked) + yield chunk + async for remaining_chunk in chunks: + yield remaining_chunk + return + late_completed: Final = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream) + if isinstance(late_completed, ResponsesAPIResponse): + await _bill_finished_stream(late_completed, track_untracked) + + +async def _bill_finished_stream( + response: ResponsesAPIResponse, track_untracked: Callable[[str, ResponsesAPIResponse], Awaitable[bool]] +) -> None: + if not is_background_response(response): + return + await _settle_background_billing( + lambda cost_check: cost_check.bill_if_finished(response.id, response, partial(track_untracked, response.id)) + ) @router.delete( @@ -874,9 +1150,14 @@ async def delete_response( # Normal provider response flow data: Final = await read_request_body(request=request) data["response_id"] = response_id + charge: Final = ( + await _settle_background_billing(lambda cost_check: cost_check.bill_from_deployment(response_id)) + if _caller_may_address_response(response_id, user_api_key_dict, proxy_logging_obj) + else None + ) processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + deleted: Final = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -901,6 +1182,8 @@ async def delete_response( proxy_logging_obj=proxy_logging_obj, version=version, ) + _report_charge(charge, fastapi_response) + return deleted @router.get( @@ -1287,10 +1570,11 @@ async def cancel_response( # Normal provider response flow data: Final = await read_request_body(request=request) + read_settles_billing: Final = not carries_client_side_overrides(data) data["response_id"] = response_id processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response: Final[object] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -1315,6 +1599,14 @@ async def cancel_response( proxy_logging_obj=proxy_logging_obj, version=version, ) + await _bill_finished_background_response( + response=response, + response_id=response_id, + user_api_key_dict=user_api_key_dict, + fastapi_response=fastapi_response, + read_settles_billing=read_settles_billing, + ) + return response def _resolve_responses_ws_session_limit_seconds() -> float: @@ -1578,7 +1870,7 @@ async def _responses_websocket_session( } ) ) - except Exception: + except (WebSocketDisconnect, RuntimeError, OSError): pass await websocket.close(code=1008, reason="Pre-call error") return diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 12438614755..d40bac380aa 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") + end_user_id String? updated_at DateTime @updatedAt updated_by String? diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 36b414a9a87..5b02d6a6afd 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -17,6 +17,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( CLI_SESSION_KEY_PREFIX, EMPTY_MAPPING, + INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, @@ -24,6 +25,7 @@ from litellm.constants import ( LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, MCP_SPEND_LOG_MODEL_PREFIX, + NON_INFERENCE_CALL_TYPES, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL, @@ -63,6 +65,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token from litellm.types.router import DeploymentTypedDict, LiteLLM_Params from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, PROMPT_CARRYING_GUARDRAIL_FIELDS, AzureSpillover, CallTypes, @@ -292,9 +295,23 @@ def _get_spend_logs_metadata( BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" +BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" -def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | None: +def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object]) -> str: + if call_type == CallTypes.aretrieve_batch.value: + return BATCH_COST_REQUEST_ID_SUFFIX + if ( + call_type in NON_INFERENCE_CALL_TYPES + and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + ): + return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX + return "" + + +def get_spend_logs_id( + call_type: str, response_obj: dict, kwargs: dict, metadata: Mapping[str, object] = EMPTY_MAPPING +) -> str | None: standard_logging_payload = kwargs.get("standard_logging_object") candidate_ids: Final = ( response_obj.get("id"), @@ -304,9 +321,9 @@ def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | resolved_id: Final = next( (candidate for candidate in candidate_ids if isinstance(candidate, str) and candidate), None ) - if resolved_id is not None and call_type == CallTypes.aretrieve_batch.value: - return f"{resolved_id}{BATCH_COST_REQUEST_ID_SUFFIX}" - return resolved_id + if resolved_id is None: + return None + return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata)}" _MISSING_ATTRIBUTE: Final = object() @@ -560,7 +577,7 @@ def get_logging_payload( if not usage and isinstance(_combined_usage, litellm.Usage): usage = _combined_usage.model_dump() - id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs) + id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs, metadata) standard_logging_payload: Final = cast(StandardLoggingPayload | None, kwargs.get("standard_logging_object", None)) end_user_id = get_end_user_id_for_cost_tracking(litellm_params) diff --git a/litellm/router.py b/litellm/router.py index f9388781aee..18a26641b04 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -214,7 +214,7 @@ from litellm.router_utils.cooldown_handlers import ( async_get_cooldown_deployments_with_debug_info, get_cooldown_deployments, is_advisor_orchestration_failure, - is_background_response_cost_poll_not_found, + is_background_response_poll_not_found, is_caller_scoped_auth_failure, is_caller_timeout_408, set_cooldown_deployments, @@ -3617,7 +3617,7 @@ class Router: # fallback_iterator, or any future wrapper). Without this # the proxy's container-ownership hook (which reads # ``getattr(stream_response, "completed_response", None)`` - # via _extract_completed_responses_response) silently + # via extract_completed_responses_response) silently # records nothing on streaming /v1/responses calls — every # follow-up /v1/containers//files call then 403s for # the very key that created the container (#30210). @@ -8743,10 +8743,10 @@ class Router: litellm_params: Final = kwargs.get("litellm_params", {}) _model_info: Final = litellm_params.get("model_info", {}) - if is_background_response_cost_poll_not_found(exception, litellm_params): + if is_background_response_poll_not_found(exception, litellm_params): verbose_router_logger.debug( "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Provider 404 came from the background response cost poll, not the deployment's health." + "Provider 404 came from the background response poll, not the deployment's health." ) return False diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index a5f6277d2d7..70b32f59f81 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -30,7 +30,11 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCacheValue from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback from litellm.types.router import Deployment -from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN +from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, + InternalCallOrigin, +) from .router_callbacks.track_deployment_metrics import ( get_deployment_failures_for_current_minute, @@ -71,11 +75,16 @@ def is_advisor_orchestration_failure(exception: BaseException | None) -> bool: return bool(getattr(exception, _ADVISOR_ORCHESTRATION_FAILURE_ATTR, False)) -def is_background_response_cost_poll_not_found(exception: Exception, litellm_params: Mapping[str, object]) -> bool: - """Whether a background response cost poll failed with a provider 404.""" +_BACKGROUND_RESPONSE_POLL_CALL_ORIGINS: Final[frozenset[InternalCallOrigin]] = frozenset( + {BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN} +) + + +def is_background_response_poll_not_found(exception: Exception, litellm_params: Mapping[str, object]) -> bool: + """Whether a background response poll, a status check or the billed read, failed with a provider 404.""" return getattr(exception, "status_code", None) == 404 and any( isinstance(candidate, Mapping) - and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) in _BACKGROUND_RESPONSE_POLL_CALL_ORIGINS for candidate in (litellm_params.get("metadata"), litellm_params.get("litellm_metadata")) ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index dafd9f5b771..e7e4d8489d9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3164,6 +3164,7 @@ InternalCallOrigin = Literal[ "shadow_eval_judge", "llm_as_a_judge_guardrail", "background_response_cost_poll", + "background_response_status_poll", ] """Which internal litellm feature originated a billed sub-call, so a spend log row records that it is not traffic the caller sent.""" @@ -3173,6 +3174,7 @@ SHADOW_EVAL_ROUTER_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_router" SHADOW_EVAL_JUDGE_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_judge" LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN: Final[InternalCallOrigin] = "llm_as_a_judge_guardrail" BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_cost_poll" +BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_status_poll" class StandardLoggingHeuristicV2Forecast(TypedDict): diff --git a/schema.prisma b/schema.prisma index 12438614755..d40bac380aa 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") + end_user_id String? updated_at DateTime @updatedAt updated_by String? diff --git a/tests/integration/spend/_background_responses.py b/tests/integration/spend/_background_responses.py new file mode 100644 index 00000000000..ac044dedc9c --- /dev/null +++ b/tests/integration/spend/_background_responses.py @@ -0,0 +1,327 @@ +import itertools +import json +import math +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from typing import Final + +import httpx +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.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, SseResponse +from pydantic import JsonValue + +INPUT_COST_PER_TOKEN: Final = 0.001 +OUTPUT_COST_PER_TOKEN: Final = 0.002 +INPUT_TOKENS: Final = 19 +OUTPUT_TOKENS: Final = 7 +PRICE: Final = INPUT_TOKENS * INPUT_COST_PER_TOKEN + OUTPUT_TOKENS * OUTPUT_COST_PER_TOKEN +SCHEDULER_PERIOD_CEILING_SECONDS: Final = 31 +ONE_POLL_SECONDS: Final = 2 * SCHEDULER_PERIOD_CEILING_SECONDS + 8 +POLL_WINDOW_SECONDS: Final = SCHEDULER_PERIOD_CEILING_SECONDS + 9 +SPEND_LANDING_SECONDS: Final = 60 +PROVIDER_ID: Final = "resp_$REQUEST_ID" +PROMPT: Final = "bill me once" +READS: Final = 3 +CONCURRENT_READS: Final = 12 +FINISHED_STATUSES: Final = frozenset({"completed", "incomplete", "cancelled"}) +MANAGED_READ_MARKER: Final = "_managed_background_response" +ATTRIBUTION_COLUMNS: Final = ("api_key", "team_id", "organization_id", "user") + + +def response_body(status: str, *, background: bool = True, **fields: JsonValue) -> dict[str, JsonValue]: + finished: Final = status in FINISHED_STATUSES + return { + "id": PROVIDER_ID, + "object": "response", + "created_at": 1, + "status": status, + "background": background, + "store": False, + "error": None, + "incomplete_details": {"reason": "max_output_tokens"} if status == "incomplete" else None, + "model": "gpt-4o-mini", + "output": ( + [ + { + "id": "msg_$REQUEST_ID", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + } + ] + if finished + else [] + ), + "parallel_tool_calls": True, + "temperature": 1.0, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + "truncation": "disabled", + "usage": ( + { + "input_tokens": INPUT_TOKENS, + "output_tokens": OUTPUT_TOKENS, + "total_tokens": INPUT_TOKENS + OUTPUT_TOKENS, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + } + if finished + else None + ), + "metadata": {}, + **fields, + } + + +def json_reply(body: dict[str, JsonValue], status: int = 200) -> JsonResponse: + return JsonResponse(content_type="application/json", body=body, status=status) + + +def scripted_routes( + create: JsonResponse, read: JsonResponse, *, provider_id: str = PROVIDER_ID, **extra: JsonResponse +) -> RoutedResponse: + return RoutedResponse( + content_type="application/x-routed", + routes={ + "POST /responses": create, + f"GET /responses/{provider_id}": read, + **{f"POST /responses/{provider_id}/{action}": response for action, response in extra.items()}, + }, + ) + + +def background_routes(read_status: str, **fields: JsonValue) -> RoutedResponse: + return scripted_routes( + json_reply(response_body("queued", **fields)), json_reply(response_body(read_status, **fields)) + ) + + +def sse_frame(event: str, sequence_number: int, response: dict[str, JsonValue]) -> str: + return f"event: {event}\ndata: " + json.dumps( + {"type": event, "sequence_number": sequence_number, "response": response}, separators=(",", ":") + ) + + +def streamed_reply(*, background: bool = True) -> SseResponse: + return SseResponse( + content_type="text/event-stream", + frames=( + sse_frame("response.created", 0, response_body("queued", background=background)), + sse_frame("response.completed", 1, response_body("completed", background=background)), + ), + ) + + +@dataclass(frozen=True, slots=True) +class Upstream: + model: str + model_id: str + handle: ScenarioHandle + + def script(self, response: RoutedResponse | SseResponse) -> None: + register_scenario(self.handle.scenario_id, response, control_url=self.handle.control_url) + + +def scripted_upstream( + scenario: Scenario, + response: RoutedResponse | SseResponse, + *, + model_info: Mapping[str, JsonValue] | None = None, + **litellm_params: JsonValue, +) -> Upstream: + scenario_id: Final = f"bgonce-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario(scenario_id, response) + scenario.cleanups.callback(delete_scenario, handle) + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": scenario_id, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-scripted-provider", + "api_base": handle.api_base(), + "input_cost_per_token": INPUT_COST_PER_TOKEN, + "output_cost_per_token": OUTPUT_COST_PER_TOKEN, + **litellm_params, + }, + **({"model_info": dict(model_info)} if model_info else {}), + }, + ) + model_id: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, model_id) + return Upstream(scenario_id, model_id, handle) + + +def forget_row(response_id: str) -> None: + write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', (response_id,)) + + +def create_background_response( + scenario: Scenario, + key: str, + model: str, + *, + body: Mapping[str, JsonValue] | None = None, + headers: Mapping[str, str] | None = None, +) -> str: + response: Final = scenario.gateway.request( + "POST", + "/v1/responses", + body or {"model": model, "input": PROMPT, "background": True, "store": False}, + key=key, + headers=headers, + ) + assert response.status_code == 200, response.text + response_id: Final = string_value(object_value(response.json())["id"]) + scenario.cleanups.callback(forget_row, response_id) + return response_id + + +@dataclass(frozen=True, slots=True) +class Read: + status: str + text: str + cost: float | None + + +def reported_cost(headers: httpx.Headers) -> float | None: + reported: Final = headers.get("x-litellm-response-cost") + return None if reported is None else float(reported) + + +def read_of(response: httpx.Response) -> Read: + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + output: Final = list_value(body["output"]) + text: Final = ( + string_value(object_value(list_value(object_value(output[0])["content"])[0])["text"]) if output else "" + ) + return Read(string_value(body["status"]), text, reported_cost(response.headers)) + + +def read_response(gateway: Gateway, key: str, response_id: str, body: Mapping[str, JsonValue] | None = None) -> Read: + return read_of(gateway.request("GET", f"/v1/responses/{response_id}", body, key=key)) + + +def free_reads(status: str, count: int) -> tuple[Read, ...]: + return (Read(status, "pong", 0.0),) * count + + +def settling_reads(status: str, count: int) -> tuple[Read, ...]: + return (Read(status, "pong", pytest.approx(PRICE)), *free_reads(status, count - 1)) + + +def hashed(key: str) -> str: + return sha256(key.encode()).hexdigest() + + +def number(value: JsonValue) -> float: + assert isinstance(value, int | float), value + return float(value) + + +def key_spend(key: str) -> float: + rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed(key),)) + return number(rows[0]["spend"]) + + +def team_spend(team_id: str) -> float: + rows: Final = read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)) + return number(rows[0]["spend"]) + + +def user_spend(user_id: str) -> float: + rows: Final = read_rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = %s', (user_id,)) + return number(rows[0]["spend"]) + + +def org_spend(organization_id: str) -> float: + rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (organization_id,) + ) + return number(rows[0]["spend"]) + + +def tag_spend(tag: str) -> float: + rows: Final = read_rows( + 'SELECT COALESCE(SUM(spend), 0) AS spend FROM "LiteLLM_DailyTagSpend" WHERE tag = %s', (tag,) + ) + return number(rows[0]["spend"]) + + +def costs(expected: float) -> Callable[[float], bool]: + return lambda spend: math.isclose(spend, expected, abs_tol=1e-9) + + +def assert_charged_exactly(read_spend: Callable[[], float], expected: float) -> None: + held(read_spend, costs(expected), holding=POLL_WINDOW_SECONDS, seconds=POLL_WINDOW_SECONDS + SPEND_LANDING_SECONDS) + + +def charged_rows(model_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, call_type, api_key, team_id, organization_id, "user", request_tags, spend, ' + 'prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_id = %s AND spend > 0', + (model_id,), + ) + + +def create_rows(model_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT api_key, team_id, organization_id, "user", request_tags FROM "LiteLLM_SpendLogs" ' + "WHERE model_id = %s AND call_type = 'aresponses'", + (model_id,), + ) + + +def one_charged_row(model_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: charged_rows(model_id), lambda found: len(found) >= 1, seconds=30) + assert len(rows) == 1, rows + return rows[0] + + +def managed_rows(response_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT status, api_key, created_by, team_id FROM "LiteLLM_ManagedObjectTable" ' + "WHERE unified_object_id = %s OR model_object_id = %s", + (response_id, response_id), + ) + + +def assert_billed_like_its_creator( + row: dict[str, JsonValue], key: str, team_id: str | None, user_id: str | None +) -> None: + assert string_value(row["request_id"]).endswith("_background_cost"), row + assert row["call_type"] == "aget_responses" + assert row["api_key"] == hashed(key) + assert row["team_id"] == team_id + assert row["user"] == user_id + assert row["prompt_tokens"] == INPUT_TOKENS + assert row["completion_tokens"] == OUTPUT_TOKENS + assert number(row["spend"]) == pytest.approx(PRICE) + + +def observed_requests(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + return tuple(map(object_value, list_value(upstream.get("/__observations").json()["requests"]))) + + +def calls_to(requests: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> tuple[dict[str, JsonValue], ...]: + return tuple(request for request in requests if string_value(request["path"]).startswith(f"/{handle.scenario_id}/")) + + +def upstream_log(gateway: Gateway, handle: ScenarioHandle) -> Iterator[tuple[dict[str, JsonValue], ...]]: + fresh: Final = (calls_to(observed_requests(gateway), handle) for _ in itertools.count()) + return itertools.accumulate(fresh, lambda seen, calls: (*seen, *calls)) + + +def submissions(log: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> tuple[JsonValue, ...]: + return tuple( + call["body"] for call in log if call["method"] == "POST" and call["path"] == f"/{handle.scenario_id}/responses" + ) diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py new file mode 100644 index 00000000000..8a9818a7532 --- /dev/null +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -0,0 +1,771 @@ +import asyncio +import json +import uuid +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +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.upstream import delete_scenario, register_scenario +from integration.caching.response_cache_case import entry_key, redis_store +from integration.spend._background_responses import ( + ATTRIBUTION_COLUMNS, + CONCURRENT_READS, + INPUT_COST_PER_TOKEN, + INPUT_TOKENS, + MANAGED_READ_MARKER, + ONE_POLL_SECONDS, + OUTPUT_COST_PER_TOKEN, + OUTPUT_TOKENS, + POLL_WINDOW_SECONDS, + PRICE, + PROMPT, + READS, + SPEND_LANDING_SECONDS, + Read, + Upstream, + assert_billed_like_its_creator, + assert_charged_exactly, + background_routes, + charged_rows, + costs, + create_background_response, + create_rows, + forget_row, + free_reads, + hashed, + json_reply, + key_spend, + managed_rows, + number, + one_charged_row, + org_spend, + read_of, + read_response, + reported_cost, + response_body, + scripted_routes, + scripted_upstream, + settling_reads, + streamed_reply, + submissions, + tag_spend, + team_spend, + upstream_log, + user_spend, +) +from openai._legacy_response import LegacyAPIResponse +from openai.types.responses import Response +from pydantic import JsonValue + +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + + +def _sdk_read(raw: LegacyAPIResponse[Response]) -> Read: + parsed: Final = raw.parse() + return Read(parsed.status or "", parsed.output_text, reported_cost(raw.headers)) + + +def _httpx_flow(scenario: Scenario, key: str, model: str) -> tuple[str, tuple[Read, ...]]: + response_id: Final = create_background_response(scenario, key, model) + return response_id, tuple(read_response(scenario.gateway, key, response_id) for _ in range(READS)) + + +def _forged_metadata_flow(scenario: Scenario, key: str, model: str) -> tuple[str, tuple[Read, ...]]: + forged: Final[dict[str, JsonValue]] = { + "internal_call_origin": BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + MANAGED_READ_MARKER: False, + } + response_id: Final = create_background_response(scenario, key, model) + return response_id, tuple( + read_response(scenario.gateway, key, response_id, {"metadata": forged, "litellm_metadata": forged}) + for _ in range(READS) + ) + + +def _sdk_poll(client: openai.OpenAI, response_id: str, remaining: int) -> tuple[Read, ...]: + if remaining == 0: + return () + raw: Final = client.responses.with_raw_response.retrieve(response_id) + return (_sdk_read(raw), *_sdk_poll(client, raw.parse().id, remaining - 1)) + + +def _sdk_sync_flow(scenario: Scenario, key: str, model: str) -> tuple[str, tuple[Read, ...]]: + with openai.OpenAI(base_url=f"{scenario.gateway.client.base_url}/v1", api_key=key, max_retries=0) as client: + created: Final = client.responses.create(model=model, input=PROMPT, background=True, store=False) + scenario.cleanups.callback(forget_row, created.id) + return created.id, _sdk_poll(client, created.id, READS) + + +async def _sdk_async_poll(client: openai.AsyncOpenAI, response_id: str, remaining: int) -> tuple[Read, ...]: + if remaining == 0: + return () + raw: Final = await client.responses.with_raw_response.retrieve(response_id) + return (_sdk_read(raw), *await _sdk_async_poll(client, raw.parse().id, remaining - 1)) + + +async def _sdk_async_reads(base_url: str, key: str, model: str) -> tuple[str, tuple[Read, ...]]: + async with openai.AsyncOpenAI(base_url=base_url, api_key=key, max_retries=0) as client: + created: Final = await client.responses.create(model=model, input=PROMPT, background=True, store=False) + return created.id, await _sdk_async_poll(client, created.id, READS) + + +def _sdk_async_flow(scenario: Scenario, key: str, model: str) -> tuple[str, tuple[Read, ...]]: + response_id, reads = asyncio.run(_sdk_async_reads(f"{scenario.gateway.client.base_url}/v1", key, model)) + scenario.cleanups.callback(forget_row, response_id) + return response_id, reads + + +Flow = Callable[[Scenario, str, str], tuple[str, tuple[Read, ...]]] +FLOWS: Final[Mapping[str, Flow]] = { + "httpx": _httpx_flow, + "httpx_forged_metadata": _forged_metadata_flow, + "openai_sync": _sdk_sync_flow, + "openai_async": _sdk_async_flow, +} + + +def _team_member_key(scenario: Scenario, **fields: JsonValue) -> tuple[str, str, str]: + team_id: Final = scenario.team(**fields) + user_id: Final = scenario.member(team_id) + return scenario.key(team_id=team_id, user_id=user_id), team_id, user_id + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize( + ("client", "terminal_status"), + [ + ("httpx", "completed"), + ("httpx", "incomplete"), + ("httpx_forged_metadata", "completed"), + ("openai_sync", "completed"), + ("openai_async", "completed"), + ], +) +def test_a_finished_background_response_bills_its_creator_once_however_often_it_is_read( + gateway: Gateway, client: str, terminal_status: str +) -> None: + with gateway.scenario() as scenario: + key, team_id, user_id = _team_member_key(scenario) + upstream: Final = scripted_upstream(scenario, background_routes(terminal_status)) + log: Final = upstream_log(gateway, upstream.handle) + response_id, reads = FLOWS[client](scenario, key, upstream.model) + assert reads in (settling_reads(terminal_status, READS), free_reads(terminal_status, READS)) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert_billed_like_its_creator(one_charged_row(upstream.model_id), key, team_id, user_id) + assert team_spend(team_id) == pytest.approx(PRICE) + assert submissions(next(log), upstream.handle) == ( + {"model": "gpt-4o-mini", "input": PROMPT, "background": True, "store": False}, + ) + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] + + +def _org_team_member_key(scenario: Scenario) -> tuple[str, str, str, str]: + organization_id: Final = scenario.organization() + team_id: Final = scenario.team(organization_id=organization_id) + user_id: Final = scenario.org_member(organization_id, "internal_user") + scenario.gateway.post("/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}}) + return scenario.key(team_id=team_id, user_id=user_id), team_id, user_id, organization_id + + +def _billed_by_the_poll_alone( + scenario: Scenario, key: str, tags: tuple[str, str] +) -> tuple[Upstream, dict[str, JsonValue]]: + body_tag, header_tag = tags + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response( + scenario, + key, + upstream.model, + body={ + "model": upstream.model, + "input": PROMPT, + "background": True, + "store": False, + "litellm_metadata": {"tags": [body_tag]}, + }, + headers={"x-litellm-tags": header_tag}, + ) + eventually( + lambda: [row["status"] for row in managed_rows(response_id)], + lambda statuses: statuses == ["completed"], + seconds=ONE_POLL_SECONDS, + ) + held(lambda: key_spend(key), costs(PRICE), holding=5, seconds=SPEND_LANDING_SECONDS) + charged: Final = one_charged_row(upstream.model_id) + create_row: Final = eventually(lambda: create_rows(upstream.model_id), lambda rows: len(rows) == 1)[0] + assert {column: charged[column] for column in ATTRIBUTION_COLUMNS} == { + column: create_row[column] for column in ATTRIBUTION_COLUMNS + } + assert set(list_value(charged["request_tags"])) == set(list_value(create_row["request_tags"])) >= set(tags) + return upstream, charged + + +def _unique_tags() -> tuple[str, str]: + return f"audit-body-{uuid.uuid4().hex[:8]}", f"audit-header-{uuid.uuid4().hex[:8]}" + + +@pytest.mark.timeout(240) +def test_a_background_response_nobody_reads_is_billed_once_to_its_creating_key_team_org_and_tags( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + key, team_id, user_id, organization_id = _org_team_member_key(scenario) + tags: Final = _unique_tags() + _, charged = _billed_by_the_poll_alone(scenario, key, tags) + assert_billed_like_its_creator(charged, key, team_id, user_id) + assert charged["organization_id"] == organization_id + eventually(lambda: tuple(map(tag_spend, tags)), lambda spent: all(map(costs(PRICE), spent))) + eventually(lambda: org_spend(organization_id), costs(PRICE)) + listed: Final = gateway.request( + "GET", "/spend/logs", params={"request_id": string_value(charged["request_id"])} + ) + assert listed.status_code == 200, listed.text + assert [number(object_value(row)["spend"]) for row in list_value(listed.json())] == [pytest.approx(PRICE)] + assert number(object_value(gateway.get("/key/info", {"key": key})["info"])["spend"]) == pytest.approx(PRICE) + team_info: Final = object_value(gateway.get("/team/info", {"team_id": team_id})["team_info"]) + assert number(team_info["spend"]) == pytest.approx(PRICE) + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("in_a_team", [False, True], ids=["no_team", "team"]) +def test_a_background_response_nobody_reads_on_a_key_with_no_user_is_billed_once_to_the_key( + gateway: Gateway, in_a_team: bool +) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key(team_id=scenario.team()) if in_a_team else scenario.key() + _, charged = _billed_by_the_poll_alone(scenario, key, _unique_tags()) + assert string_value(charged["request_id"]).endswith("_background_cost"), charged + assert charged["api_key"] == hashed(key) + assert number(charged["spend"]) == pytest.approx(PRICE) + + +@pytest.mark.timeout(240) +def test_a_background_response_is_free_until_it_finishes_and_a_failed_read_bills_nothing(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream: Final = scripted_upstream(scenario, background_routes("in_progress")) + response_id: Final = create_background_response(scenario, key, upstream.model) + assert read_response(gateway, key, response_id) == Read("in_progress", "", 0.0) + outage: Final = json_reply({"error": {"message": "The server had an error.", "type": "server_error"}}, 500) + upstream.script(scripted_routes(json_reply(response_body("queued")), outage)) + failed: Final = gateway.request("GET", f"/v1/responses/{response_id}", key=key) + assert failed.status_code == 500, failed.text + assert "The server had an error." in failed.text + assert [row["status"] for row in managed_rows(response_id)] == ["queued"] + upstream.script(background_routes("completed")) + assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == settling_reads("completed", 2) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] + + +@pytest.mark.timeout(240) +def test_cancelling_a_background_response_bills_it_once(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + cancelled: Final = json_reply(response_body("cancelled")) + upstream: Final = scripted_upstream( + scenario, scripted_routes(json_reply(response_body("queued")), cancelled, cancel=cancelled) + ) + response_id: Final = create_background_response(scenario, key, upstream.model) + cancel: Final = gateway.request("POST", f"/v1/responses/{response_id}/cancel", key=key) + reads: Final = (read_of(cancel), read_response(gateway, key, response_id)) + assert reads in (settling_reads("cancelled", 2), free_reads("cancelled", 2)) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert number(one_charged_row(upstream.model_id)["spend"]) == pytest.approx(PRICE) + + +@pytest.mark.timeout(120) +def test_a_foreground_response_bills_its_create_once_and_its_read_is_free(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + foreground: Final = json_reply(response_body("completed", background=False)) + upstream: Final = scripted_upstream(scenario, scripted_routes(foreground, foreground)) + created: Final = gateway.request( + "POST", "/v1/responses", {"model": upstream.model, "input": PROMPT, "store": False}, key=key + ) + assert read_of(created) == Read("completed", "pong", pytest.approx(PRICE)) + response_id: Final = string_value(object_value(created.json())["id"]) + assert read_response(gateway, key, response_id) == Read("completed", "pong", 0.0) + held(lambda: key_spend(key), costs(PRICE), holding=10, seconds=SPEND_LANDING_SECONDS) + assert managed_rows(response_id) == [] + + +STREAMED_BACKGROUND_FIELDS: Final[Mapping[str, dict[str, JsonValue]]] = { + "top_level": {"background": True}, + "extra_body": {"extra_body": {"background": True}}, +} + + +def _streamed_response_id(gateway: Gateway, key: str, model: str, background: Mapping[str, JsonValue]) -> str: + with gateway.client.stream( + "POST", + "/v1/responses", + json={"model": model, "input": PROMPT, "store": False, "stream": True, **background}, + headers={"Authorization": f"Bearer {key}"}, + ) as stream: + assert stream.status_code == 200, stream.read() + events: Final = tuple( + object_value(json.loads(line.removeprefix("data: "))) + for line in stream.iter_lines() + if line.startswith("data: {") + ) + completed: Final = [event for event in events if event["type"] == "response.completed"] + assert len(completed) == 1, events + return string_value(object_value(completed[0]["response"])["id"]) + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("background_in", STREAMED_BACKGROUND_FIELDS.keys()) +def test_a_streamed_background_response_read_later_bills_its_stream_alone(gateway: Gateway, background_in: str) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream: Final = scripted_upstream(scenario, streamed_reply()) + response_id: Final = _streamed_response_id( + gateway, key, upstream.model, STREAMED_BACKGROUND_FIELDS[background_in] + ) + scenario.cleanups.callback(forget_row, response_id) + upstream.script(background_routes("completed")) + assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == free_reads("completed", 2) + assert_charged_exactly(lambda: key_spend(key), PRICE) + + +def _finished(scenario: Scenario, key: str) -> tuple[Upstream, str]: + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + return upstream, create_background_response(scenario, key, upstream.model) + + +def _lose_the_row_then_finish(scenario: Scenario, key: str) -> tuple[Upstream, str]: + upstream: Final = scripted_upstream(scenario, background_routes("queued")) + response_id: Final = create_background_response(scenario, key, upstream.model) + forget_row(response_id) + upstream.script(background_routes("completed")) + return upstream, response_id + + +@pytest.mark.timeout(240) +def test_a_finished_response_the_proxy_lost_track_of_bills_its_reader_once(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key, team_id, user_id = _team_member_key(scenario) + upstream, response_id = _lose_the_row_then_finish(scenario, key) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads == settling_reads("completed", READS) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert_billed_like_its_creator(one_charged_row(upstream.model_id), key, team_id, user_id) + assert [(row["status"], row["api_key"]) for row in managed_rows(response_id)] == [("completed", hashed(key))] + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("tracked", [True, False], ids=["tracked", "untracked"]) +def test_concurrent_reads_of_a_finished_response_bill_it_once(gateway: Gateway, tracked: bool) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream, response_id = _finished(scenario, key) if tracked else _lose_the_row_then_finish(scenario, key) + with ThreadPoolExecutor(max_workers=CONCURRENT_READS) as pool: + reads: Final = tuple(pool.map(lambda _: read_response(gateway, key, response_id), range(CONCURRENT_READS))) + assert [(read.status, read.text) for read in reads] == [("completed", "pong")] * CONCURRENT_READS + assert [read.cost for read in reads if read.cost != 0.0] in ([], [pytest.approx(PRICE)]) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert len(charged_rows(upstream.model_id)) == 1 + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] + + +@pytest.mark.timeout(240) +def test_another_key_of_the_same_owner_reads_for_free_and_the_creating_key_pays(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + creator, team_id, user_id = _team_member_key(scenario) + reader: Final = scenario.key(team_id=team_id, user_id=user_id) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, creator, upstream.model) + reads: Final = tuple(read_response(gateway, reader, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + assert_charged_exactly(lambda: key_spend(creator), PRICE) + assert key_spend(reader) == 0 + assert_billed_like_its_creator(one_charged_row(upstream.model_id), creator, team_id, user_id) + + +@pytest.mark.timeout(240) +def test_a_create_that_asks_for_background_inside_extra_body_is_billed_once(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + log: Final = upstream_log(gateway, upstream.handle) + response_id: Final = create_background_response( + scenario, + key, + upstream.model, + body={"model": upstream.model, "input": PROMPT, "store": False, "extra_body": {"background": True}}, + ) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert [object_value(body).get("background") for body in submissions(next(log), upstream.handle)] == [True] + + +@pytest.mark.timeout(240) +def test_a_response_whose_instructions_carry_a_nul_is_tracked_and_billed_once(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream: Final = scripted_upstream(scenario, background_routes("completed", instructions="be\u0000brief")) + response_id: Final = create_background_response(scenario, key, upstream.model) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + assert_charged_exactly(lambda: key_spend(key), PRICE) + stored: Final = read_rows( + "SELECT file_object->>'instructions' AS instructions FROM \"LiteLLM_ManagedObjectTable\" " + "WHERE unified_object_id = %s", + (response_id,), + ) + assert stored == [{"instructions": "bebrief"}] + + +@pytest.mark.timeout(240) +def test_a_request_tag_carrying_a_nul_is_stored_without_it_and_lands_on_the_billed_row(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + suffix: Final = uuid.uuid4().hex[:8] + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response( + scenario, + key, + upstream.model, + body={ + "model": upstream.model, + "input": PROMPT, + "background": True, + "store": False, + "litellm_metadata": {"tags": [f"audit-nul\u0000{suffix}"]}, + }, + ) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(2)) + assert reads in (settling_reads("completed", 2), free_reads("completed", 2)) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert f"audit-nul{suffix}" in list_value(one_charged_row(upstream.model_id)["request_tags"]) + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("managed_row", ["tracked", "lost"]) +def test_a_read_retargeted_with_a_client_side_api_base_bills_the_real_usage_once( + gateway: Gateway, managed_row: str +) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream: Final = scripted_upstream( + scenario, background_routes("queued"), configurable_clientside_auth_params=["api_base"] + ) + provider_id: Final = f"resp_{upstream.model}" + cheap_body: Final = { + **response_body("completed"), + "id": provider_id, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + decoy: Final = register_scenario( + f"bgonce-decoy-{uuid.uuid4().hex[:12]}", + scripted_routes(json_reply(cheap_body), json_reply(cheap_body), provider_id=provider_id), + ) + scenario.cleanups.callback(delete_scenario, decoy) + decoy_log: Final = upstream_log(gateway, decoy) + response_id: Final = create_background_response(scenario, key, upstream.model) + if managed_row == "lost": + forget_row(response_id) + upstream.script(background_routes("completed")) + retargeted: Final = read_response( + gateway, key, response_id, {"model": upstream.model, "api_base": decoy.api_base()} + ) + assert (retargeted,) in (settling_reads("completed", 1), free_reads("completed", 1)) + assert [call["method"] for call in next(decoy_log)] == ["GET"] + eventually( + lambda: [row["status"] for row in managed_rows(response_id)], + lambda statuses: statuses == ["completed"], + seconds=ONE_POLL_SECONDS, + ) + assert_charged_exactly(lambda: key_spend(key), PRICE) + charged: Final = one_charged_row(upstream.model_id) + assert (charged["prompt_tokens"], charged["completion_tokens"]) == (INPUT_TOKENS, OUTPUT_TOKENS) + + +def _id_a_read_returned(gateway: Gateway, key: str, response_id: str) -> str: + read: Final = gateway.request("GET", f"/v1/responses/{response_id}", None, key=key) + assert read.status_code == 200, read.text + return string_value(object_value(read.json())["id"]) + + +@pytest.mark.timeout(240 + ONE_POLL_SECONDS) +@pytest.mark.parametrize("read_by", ["create_id", "id_a_read_returned"]) +def test_a_row_stored_before_the_upgrade_bills_the_creating_team_and_user_once(gateway: Gateway, read_by: str) -> None: + with gateway.scenario() as scenario: + key, team_id, user_id = _team_member_key(scenario) + upstream: Final = scripted_upstream(scenario, background_routes("queued")) + created_id: Final = create_background_response(scenario, key, upstream.model) + response_id: Final = created_id if read_by == "create_id" else _id_a_read_returned(gateway, key, created_id) + write_rows( + 'UPDATE "LiteLLM_ManagedObjectTable" SET model_object_id = unified_object_id, api_key = NULL, ' + "request_tags = NULL WHERE unified_object_id = %s", + (created_id,), + ) + upstream.script(background_routes("completed")) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + held(lambda: team_spend(team_id), costs(PRICE), holding=ONE_POLL_SECONDS, seconds=SPEND_LANDING_SECONDS) + assert user_spend(user_id) == pytest.approx(PRICE) + charged: Final = one_charged_row(upstream.model_id) + assert (charged["team_id"], charged["user"]) == (team_id, user_id) + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("settled_status", ["completed", "stale_expired"]) +def test_reads_of_a_response_whose_row_already_settled_are_free(gateway: Gateway, settled_status: str) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream: Final = scripted_upstream(scenario, background_routes("queued")) + response_id: Final = create_background_response(scenario, key, upstream.model) + write_rows( + 'UPDATE "LiteLLM_ManagedObjectTable" SET status = %s WHERE unified_object_id = %s', + (settled_status, response_id), + ) + upstream.script(background_routes("completed")) + assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == free_reads("completed", 2) + assert_charged_exactly(lambda: key_spend(key), 0.0) + assert charged_rows(upstream.model_id) == [] + + +FOREIGN_IDS: Final[Mapping[str, Callable[[Upstream], str]]] = { + "unknown": lambda _: f"resp_{uuid.uuid4().hex}", + "five_kilobytes": lambda _: "resp_" + "a" * 5120, + "raw_vendor_id": lambda upstream: f"resp_{upstream.model}", +} + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("foreign", FOREIGN_IDS.keys()) +def test_a_read_of_an_id_the_proxy_never_issued_fails_without_tracking_or_billing_it( + gateway: Gateway, foreign: str +) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, key, upstream.model) + foreign_id: Final = FOREIGN_IDS[foreign](upstream) + refused: Final = gateway.request("GET", f"/v1/responses/{foreign_id}", key=key) + assert refused.status_code >= 400, refused.text + assert managed_rows(foreign_id) == [] + assert (read_response(gateway, key, response_id),) in ( + settling_reads("completed", 1), + free_reads("completed", 1), + ) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert len(charged_rows(upstream.model_id)) == 1 + + +@pytest.mark.timeout(240) +def test_a_background_create_answered_at_once_and_its_cached_twin_bill_the_create_alone(gateway: Gateway) -> None: + store: Final = redis_store() + with gateway.scenario() as scenario: + key: Final = scenario.key() + answered: Final = json_reply(response_body("completed")) + upstream: Final = scripted_upstream(scenario, scripted_routes(answered, answered)) + log: Final = upstream_log(gateway, upstream.handle) + first: Final = create_background_response(scenario, key, upstream.model) + eventually( + lambda: entry_key(store, f"msg_{upstream.model}"), + lambda found: found is not None, + ) + second: Final = create_background_response(scenario, key, upstream.model) + assert len(submissions(next(log), upstream.handle)) == 1 + reads: Final = [read_response(gateway, key, response_id) for response_id in (first, second, first, second)] + assert tuple(reads) == free_reads("completed", 4) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert len(charged_rows(upstream.model_id)) == 1 + + +def _end_user_spend(end_user: str) -> float: + rows: Final = read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id = %s', (end_user,)) + return number(rows[0]["spend"]) if rows else 0.0 + + +def _charged_end_users(model_id: str) -> list[JsonValue]: + rows: Final = read_rows('SELECT end_user FROM "LiteLLM_SpendLogs" WHERE model_id = %s AND spend > 0', (model_id,)) + return [row["end_user"] for row in rows] + + +@pytest.mark.timeout(240) +def test_a_background_response_bills_the_end_user_who_created_it_once(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + end_user: Final = f"audit-end-user-{uuid.uuid4().hex}" + scenario.cleanups.callback(gateway.request, "POST", "/end_user/delete", {"user_ids": [end_user]}) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response( + scenario, + key, + upstream.model, + body={"model": upstream.model, "input": PROMPT, "background": True, "store": False, "user": end_user}, + ) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + assert_charged_exactly(lambda: _end_user_spend(end_user), PRICE) + assert key_spend(key) == pytest.approx(PRICE) + assert _charged_end_users(upstream.model_id) == [end_user] + + +@pytest.mark.timeout(240) +def test_the_read_that_settles_a_background_response_splits_its_charge_into_input_and_output( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + _, response_id = _lose_the_row_then_finish(scenario, key) + settling: Final = gateway.request("GET", f"/v1/responses/{response_id}", key=key) + later: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS - 1)) + assert read_of(settling) == Read("completed", "pong", pytest.approx(PRICE)) + split: Final = tuple(float(settling.headers[f"x-litellm-response-cost-{part}"]) for part in ("input", "output")) + assert split == pytest.approx((INPUT_TOKENS * INPUT_COST_PER_TOKEN, OUTPUT_TOKENS * OUTPUT_COST_PER_TOKEN)) + assert later == free_reads("completed", READS - 1) + assert_charged_exactly(lambda: key_spend(key), PRICE) + + +def _answered_at_once(scenario: Scenario, key: str) -> tuple[Upstream, str]: + answered: Final = json_reply(response_body("completed")) + upstream: Final = scripted_upstream(scenario, scripted_routes(answered, answered)) + return upstream, create_background_response(scenario, key, upstream.model) + + +def _stored_output(response_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + "SELECT file_object->'output' AS output, position('pong' in file_object::text) > 0 AS holds_text " + 'FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', + (response_id,), + ) + + +STORED_BY: Final[Mapping[str, Callable[[Scenario, str], tuple[Upstream, str]]]] = { + "the_create": _answered_at_once, + "a_read": _lose_the_row_then_finish, +} + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("stored_by", STORED_BY.keys()) +def test_a_stored_background_response_keeps_none_of_the_model_output(gateway: Gateway, stored_by: str) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + _, response_id = STORED_BY[stored_by](scenario, key) + assert read_response(gateway, key, response_id).text == "pong" + assert _stored_output(response_id) == [{"output": [], "holds_text": False}] + + +def _model_usage(upstream: Upstream) -> tuple[float, int]: + rows: Final = read_rows( + "SELECT COALESCE(SUM(spend), 0) AS spend, COALESCE(SUM(request_count), 0)::bigint AS requests " + 'FROM "LiteLLM_DailyModelUsage" WHERE model_group IN (%s, %s)', + (upstream.model, upstream.model_id), + ) + requests: Final = rows[0]["requests"] + assert isinstance(requests, int), rows + return number(rows[0]["spend"]), requests + + +@pytest.mark.timeout(240) +def test_model_usage_counts_a_background_response_once_in_spend_and_only_client_calls_as_requests( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + upstream, response_id = _finished(scenario, key) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + assert_charged_exactly(lambda: key_spend(key), PRICE) + client_calls: Final = 1 + READS + held( + lambda: _model_usage(upstream), + lambda usage: costs(PRICE)(usage[0]) and usage[1] == client_calls, + holding=POLL_WINDOW_SECONDS, + seconds=POLL_WINDOW_SECONDS + SPEND_LANDING_SECONDS, + ) + + +BASE_MODEL: Final = "gpt-4o-mini" +BASE_MODEL_ONLY: Final[Mapping[str, JsonValue]] = {"base_model": BASE_MODEL} +BASE_MODEL_IN: Final[Mapping[str, tuple[Mapping[str, JsonValue] | None, Mapping[str, JsonValue]]]] = { + "model_info": (BASE_MODEL_ONLY, {}), + "litellm_params": (None, BASE_MODEL_ONLY), +} + + +def _unpriced_upstream( + scenario: Scenario, + unpriced: str, + model_info: Mapping[str, JsonValue] | None, + litellm_params: Mapping[str, JsonValue], +) -> Upstream: + foreground: Final = json_reply(response_body("completed", background=False, model=unpriced)) + return scripted_upstream( + scenario, + scripted_routes(foreground, foreground), + model_info=model_info, + model=f"openai/{unpriced}", + input_cost_per_token=None, + output_cost_per_token=None, + **litellm_params, + ) + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("base_model_in", BASE_MODEL_IN.keys()) +def test_a_background_response_on_a_base_model_deployment_bills_what_a_foreground_response_costs( + gateway: Gateway, base_model_in: str +) -> None: + with gateway.scenario() as scenario: + foreground_key: Final = scenario.key() + background_key: Final = scenario.key() + unpriced: Final = f"audit-unpriced-{uuid.uuid4().hex[:12]}" + reference: Final = _unpriced_upstream(scenario, unpriced, BASE_MODEL_ONLY, {}) + created: Final = gateway.request( + "POST", "/v1/responses", {"model": reference.model, "input": PROMPT, "store": False}, key=foreground_key + ) + foreground_cost: Final = read_of(created).cost + assert foreground_cost is not None and foreground_cost > 0, created.headers + upstream: Final = _unpriced_upstream(scenario, unpriced, *BASE_MODEL_IN[base_model_in]) + upstream.script(background_routes("completed", model=unpriced)) + response_id: Final = create_background_response(scenario, background_key, upstream.model) + reads: Final = tuple(read_response(gateway, background_key, response_id) for _ in range(READS)) + settled: Final = ( + Read("completed", "pong", pytest.approx(foreground_cost)), + *free_reads("completed", READS - 1), + ) + assert reads in (settled, free_reads("completed", READS)) + assert_charged_exactly(lambda: key_spend(background_key), foreground_cost) + background_rows: Final = [ + row for row in charged_rows(upstream.model_id) if row["api_key"] == hashed(background_key) + ] + assert [number(row["spend"]) for row in background_rows] == [pytest.approx(foreground_cost)] + + +@pytest.mark.timeout(240) +def test_a_client_cannot_make_a_background_create_free_by_forging_the_managed_marker(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key() + answered: Final = json_reply(response_body("completed")) + upstream: Final = scripted_upstream(scenario, scripted_routes(answered, answered)) + forged: Final[dict[str, JsonValue]] = {MANAGED_READ_MARKER: True} + created: Final = gateway.request( + "POST", + "/v1/responses", + { + "model": upstream.model, + "input": PROMPT, + "background": True, + "store": False, + "metadata": forged, + "litellm_metadata": forged, + }, + key=key, + ) + assert read_of(created) == Read("completed", "pong", pytest.approx(PRICE)) + response_id: Final = string_value(object_value(created.json())["id"]) + 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) diff --git a/tests/integration/spend/test_background_response_billing_proxies.py b/tests/integration/spend/test_background_response_billing_proxies.py new file mode 100644 index 00000000000..e83f2629f33 --- /dev/null +++ b/tests/integration/spend/test_background_response_billing_proxies.py @@ -0,0 +1,569 @@ +import itertools +import uuid +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, held, object_value, string_value +from integration._support.database import read_rows, scratch_database, write_rows +from integration._support.otlp_sink import Span, configure_sink, owned_sinks, recorded_spans +from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process +from integration._support.prometheus_series import Sample, scrape +from integration._support.timing import assert_not_delayed_by, timed +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import RoutedResponse, SseResponse +from integration.spend._background_responses import ( + INPUT_COST_PER_TOKEN, + MANAGED_READ_MARKER, + ONE_POLL_SECONDS, + OUTPUT_COST_PER_TOKEN, + PRICE, + PROMPT, + PROVIDER_ID, + READS, + Read, + assert_charged_exactly, + background_routes, + costs, + create_background_response, + free_reads, + hashed, + json_reply, + key_spend, + managed_rows, + number, + read_response, + reported_cost, + response_body, + scripted_routes, + settling_reads, + streamed_reply, +) +from pydantic import JsonValue, TypeAdapter + +CONFIGURATION: Final = TypeAdapter(dict[str, JsonValue]) +OWNED_PROXY_SECONDS: Final = 2 * graceful_stop_seconds() + 240 +NO_DATABASE: Final = ("DATABASE_URL", "DATABASE_URL_READ_REPLICA") +POLL_OFF: Final = {"PROXY_BATCH_POLLING_ENABLED": "false"} +SINK_DELAY_SECONDS: Final = 4.0 +SPEND_RETRY_SECONDS: Final = 60 +BURST_RESPONSES: Final = 6 +RESTART_RESPONSES: Final = 3 +COOLDOWN_SECONDS: Final = 300 +CHAT_COMPLETION: Final[dict[str, JsonValue]] = { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, +} + + +@contextmanager +def _scripted(response: RoutedResponse | SseResponse, count: int = 1) -> Iterator[tuple[ScenarioHandle, ...]]: + with ExitStack() as stack: + handles: Final = tuple(register_scenario(f"bgonce-{uuid.uuid4().hex[:12]}", response) for _ in range(count)) + for handle in handles: + stack.callback(delete_scenario, handle) + yield handles + + +def _rescript(handles: Sequence[ScenarioHandle], response: RoutedResponse) -> None: + for handle in handles: + register_scenario(handle.scenario_id, response, control_url=handle.control_url) + + +def _with_chat(response: RoutedResponse) -> RoutedResponse: + return RoutedResponse( + content_type="application/x-routed", + routes={**response.routes, "POST /chat/completions": json_reply(CHAT_COMPLETION)}, + ) + + +def _deployment(handle: ScenarioHandle, model_info: Mapping[str, JsonValue] | None = None) -> dict[str, JsonValue]: + return { + "model_name": handle.scenario_id, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-scripted-provider", + "api_base": handle.api_base(), + "input_cost_per_token": INPUT_COST_PER_TOKEN, + "output_cost_per_token": OUTPUT_COST_PER_TOKEN, + }, + **({"model_info": dict(model_info)} if model_info else {}), + } + + +def _write_config(directory: Path, config: Mapping[str, JsonValue]) -> Path: + path: Final = directory / f"background-billing-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(dict(config))) + return path + + +def _config( + directory: Path, + handles: Sequence[ScenarioHandle], + *, + model_info: Mapping[str, JsonValue] | None = None, + **sections: Mapping[str, JsonValue], +) -> Path: + shared: Final = CONFIGURATION.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + settings: Final = { + name: {**object_value(shared[name]), **sections.get(name, {})} + for name in ("general_settings", "litellm_settings", "router_settings") + } + extra: Final = {name: dict(section) for name, section in sections.items() if name not in settings} + return _write_config( + directory, + {"model_list": [_deployment(handle, model_info) for handle in handles], **settings, **extra}, + ) + + +@contextmanager +def _scratch_proxy( + gateway: Gateway, + directory: Path, + monkeypatch: pytest.MonkeyPatch, + config: Path, + overrides: Mapping[str, str] | None = None, +) -> Iterator[Gateway]: + with scratch_database() as database_url: + monkeypatch.setenv("DATABASE_URL", database_url) + with owned_proxy( + gateway, directory, {"DATABASE_URL": database_url, **(overrides or {})}, config=config + ) as proxy: + yield proxy + + +def _team_member_key(scenario: Scenario) -> str: + team_id: Final = scenario.team() + return scenario.key(team_id=team_id, user_id=scenario.member(team_id)) + + +def _billed_rows_of(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND spend > 0', (hashed(key),) + ) + + +def _assert_each_billed_once_as_background_cost(key: str, count: int) -> None: + rows: Final = _billed_rows_of(key) + assert len(rows) == count, rows + assert all(string_value(row["request_id"]).endswith("_background_cost") for row in rows), rows + assert len({row["request_id"] for row in rows}) == count, rows + assert [number(row["spend"]) for row in rows] == pytest.approx([PRICE] * count) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS) +def test_a_proxy_without_a_database_bills_every_read_of_a_finished_background_response( + gateway: Gateway, tmp_path: Path +) -> None: + forged: Final[dict[str, JsonValue]] = { + "internal_call_origin": "background_response_status_poll", + MANAGED_READ_MARKER: True, + } + with _scripted(background_routes("completed")) as (handle,): + config: Final = _write_config( + tmp_path, + { + "model_list": [_deployment(handle)], + "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}, + }, + ) + overrides: Final = {"STORE_MODEL_IN_DB": "False", "DISABLE_SCHEMA_UPDATE": "true"} + with owned_proxy( + gateway, tmp_path, overrides, config=config, database_setup=(), remove_environment=NO_DATABASE + ) as proxy: + created: Final = proxy.post( + "/v1/responses", {"model": handle.scenario_id, "input": PROMPT, "background": True, "store": False} + ) + response_id: Final = string_value(created["id"]) + reads: Final = ( + read_response(proxy, proxy.key, response_id), + read_response(proxy, proxy.key, response_id), + read_response(proxy, proxy.key, response_id, {"metadata": forged, "litellm_metadata": forged}), + ) + assert [(read.status, read.text) for read in reads] == [("completed", "pong")] * 3 + assert [read.cost for read in reads] == pytest.approx([PRICE] * 3) + + +def _span_cost(span: Span) -> float: + return number(span["attributes"].get("gen_ai.cost.total_cost", 0)) + + +def _spans_of(sink: str, key: str) -> tuple[Span, ...]: + _, spans = recorded_spans(sink) + return tuple(span for span in spans if span["attributes"].get("metadata.user_api_key_hash") == hashed(key)) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS + 120) +def test_with_the_cost_poll_off_the_first_finished_read_bills_the_creator_once_and_tells_the_destination_once( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with _scripted(background_routes("queued")) as (handle,), owned_sinks(tmp_path / "sinks") as sinks: + config: Final = _config( + tmp_path, + (handle,), + litellm_settings={"callbacks": ["otel"]}, + callback_settings={ + "otel": {"exporter": "http/json", "endpoint": sinks.operator, "use_simple_processor": True} + }, + ) + overrides: Final = {**POLL_OFF, "OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": sinks.operator} + with ( + _scratch_proxy(gateway, tmp_path, monkeypatch, config, overrides) as proxy, + proxy.scenario() as scenario, + ): + key: Final = _team_member_key(scenario) + response_id: Final = create_background_response(scenario, key, handle.scenario_id) + unfinished: Final = read_response(proxy, key, response_id) + _rescript((handle,), background_routes("completed")) + finished: Final = tuple(read_response(proxy, key, response_id) for _ in range(READS)) + assert (unfinished.status, unfinished.cost) == ("queued", 0.0) + assert finished == settling_reads("completed", READS) + assert_charged_exactly(lambda: key_spend(key), PRICE) + _assert_each_billed_once_as_background_cost(key, 1) + spans: Final = eventually( + lambda: _spans_of(sinks.operator, key), lambda found: len(found) >= READS + 3, seconds=30 + ) + assert [_span_cost(span) for span in spans if _span_cost(span) > 0] == pytest.approx([PRICE]) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS + 60) +def test_the_read_that_finishes_a_background_response_does_not_wait_on_a_slow_spend_destination( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with _scripted(background_routes("queued")) as (handle,), owned_sinks(tmp_path / "sinks") as sinks: + config: Final = _config(tmp_path, (handle,), litellm_settings={"callbacks": ["openmeter"]}) + overrides: Final = { + **POLL_OFF, + "OPENMETER_API_KEY": "om-scripted-key", + "OPENMETER_API_ENDPOINT": sinks.tenant, + } + with ( + _scratch_proxy(gateway, tmp_path, monkeypatch, config, overrides) as proxy, + proxy.scenario() as scenario, + ): + key: Final = scenario.key() + response_id: Final = create_background_response(scenario, key, handle.scenario_id) + configure_sink(sinks.tenant, delay_seconds=SINK_DELAY_SECONDS) + unfinished: Final = timed(lambda: read_response(proxy, key, response_id)) + _rescript((handle,), background_routes("completed")) + first_finished: Final = timed(lambda: read_response(proxy, key, response_id)) + later: Final = tuple(timed(lambda: read_response(proxy, key, response_id)) for _ in range(2)) + assert first_finished.value.status == "completed" + assert_not_delayed_by( + (first_finished.seconds,), + control=(unfinished.seconds, *(read.seconds for read in later)), + injected_delay=SINK_DELAY_SECONDS, + ) + + +def _provider_404() -> RoutedResponse: + gone: Final = json_reply( + { + "error": { + "message": f"Response with id '{PROVIDER_ID}' not found.", + "type": "invalid_request_error", + "param": None, + "code": None, + } + }, + 404, + ) + return scripted_routes(json_reply(response_body("queued")), gone) + + +def _create_status(proxy: Gateway, key: str, model: str) -> int: + return proxy.request( + "POST", "/v1/responses", {"model": model, "input": PROMPT, "background": True, "store": False}, key=key + ).status_code + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS + ONE_POLL_SECONDS) +def test_neither_a_poll_404_nor_a_client_read_404_takes_the_deployment_out_of_service( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with ExitStack() as stack: + (handle,) = stack.enter_context(_scripted(_provider_404())) + config: Final = _config( + tmp_path, + (handle,), + model_info={"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0, "NotFoundErrorAllowedFails": 0}}, + router_settings={"disable_cooldowns": False, "cooldown_time": COOLDOWN_SECONDS}, + ) + proxy: Final = stack.enter_context(_scratch_proxy(gateway, tmp_path, monkeypatch, config)) + scenario: Final = stack.enter_context(proxy.scenario()) + key: Final = scenario.key() + polled: Final = create_background_response(scenario, key, handle.scenario_id) + eventually( + lambda: [row["status"] for row in managed_rows(polled)], + lambda statuses: statuses == ["stale_expired"], + seconds=ONE_POLL_SECONDS, + ) + after_poll: Final = _create_status(proxy, key, handle.scenario_id) + client_read: Final = proxy.request("GET", f"/v1/responses/{polled}", key=key).status_code + after_client_read: Final = _create_status(proxy, key, handle.scenario_id) + assert (after_poll, client_read, after_client_read) == (200, 404, 200) + + +def _burst_call(proxy: Gateway, call: tuple[str, str, str]) -> int | None: + kind, key, target = call + try: + if kind == "read": + return proxy.request("GET", f"/v1/responses/{target}", key=key).status_code + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": target, "messages": [{"role": "user", "content": PROMPT}]}, + key=key, + ).status_code + except httpx.TransportError: + return None + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS + 2 * graceful_stop_seconds() + 60 + ONE_POLL_SECONDS + 100) +def test_background_responses_read_across_a_killed_worker_are_each_billed_once_to_their_creator( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with ( + _scripted(_with_chat(background_routes("queued")), BURST_RESPONSES) as handles, + scratch_database() as database_url, + ): + monkeypatch.setenv("DATABASE_URL", database_url) + config: Final = _config(tmp_path, handles) + with ( + owned_proxy_process(gateway, tmp_path, {"DATABASE_URL": database_url}, config=config, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = _team_member_key(scenario) + chat_key: Final = scenario.key() + response_ids: Final = tuple( + create_background_response(scenario, key, handle.scenario_id) for handle in handles + ) + reads: Final = tuple(("read", key, response_id) for response_id in response_ids * 3) + chats: Final = tuple(("chat", chat_key, handle.scenario_id) for handle in handles * 2) + calls: Final = tuple( + call for call in itertools.chain.from_iterable(itertools.zip_longest(reads, chats)) if call is not None + ) + with owned.replacing_worker(min(owned.worker_pids())) as kill, ThreadPoolExecutor(max_workers=10) as pool: + kill.freeze() + futures: Final = tuple(pool.submit(_burst_call, owned.gateway, call) for call in calls) + eventually( + lambda: sum(future.done() for future in futures), + lambda done: done >= len(futures) // 2, + seconds=30, + ) + kill() + outcomes: Final = tuple(future.result() for future in futures) + assert set(outcomes) <= {200, None}, outcomes + assert outcomes.count(200) >= len(outcomes) // 2, outcomes + _rescript(handles, _with_chat(background_routes("completed"))) + finished: Final = tuple(read_response(owned.gateway, key, response_id) for response_id in response_ids * 2) + assert [read.status for read in finished] == ["completed"] * 2 * BURST_RESPONSES + assert [read.cost in (0.0, pytest.approx(PRICE)) for read in finished] == [True] * 2 * BURST_RESPONSES + assert sum(read.cost or 0.0 for read in finished) <= BURST_RESPONSES * PRICE + 1e-9 + assert_charged_exactly(lambda: key_spend(key), BURST_RESPONSES * PRICE) + _assert_each_billed_once_as_background_cost(key, BURST_RESPONSES) + + +def _statuses(response_ids: Sequence[str]) -> list[JsonValue]: + rows: Final = itertools.chain.from_iterable(managed_rows(response_id) for response_id in response_ids) + return [row["status"] for row in rows] + + +def _submit(proxy: Gateway, key: str, model: str) -> str: + response: Final = proxy.request( + "POST", "/v1/responses", {"model": model, "input": PROMPT, "background": True, "store": False}, key=key + ) + assert response.status_code == 200, response.text + return string_value(object_value(response.json())["id"]) + + +@pytest.mark.timeout(2 * OWNED_PROXY_SECONDS + ONE_POLL_SECONDS + 100) +def test_background_responses_left_pending_across_a_proxy_restart_are_billed_once_to_their_creator( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with _scripted(background_routes("queued"), RESTART_RESPONSES) as handles, scratch_database() as database_url: + monkeypatch.setenv("DATABASE_URL", database_url) + config: Final = _config(tmp_path, handles) + with owned_proxy( + gateway, tmp_path / "before", {"DATABASE_URL": database_url, **POLL_OFF}, config=config + ) as first: + key: Final = string_value(first.post("/key/generate", {})["key"]) + response_ids: Final = tuple(_submit(first, key, handle.scenario_id) for handle in handles) + _rescript(handles, background_routes("completed")) + with owned_proxy(gateway, tmp_path / "after", {"DATABASE_URL": database_url}, config=config): + eventually( + lambda: _statuses(response_ids), + lambda statuses: statuses == ["completed"] * RESTART_RESPONSES, + seconds=ONE_POLL_SECONDS, + ) + assert_charged_exactly(lambda: key_spend(key), RESTART_RESPONSES * PRICE) + _assert_each_billed_once_as_background_cost(key, RESTART_RESPONSES) + + +DELETED: Final[dict[str, JsonValue]] = {"id": PROVIDER_ID, "object": "response.deleted", "deleted": True} +REFUSAL_HOLD_SECONDS: Final = 10 + + +def _deletable(read_status: str) -> RoutedResponse: + return RoutedResponse( + content_type="application/x-routed", + routes={**background_routes(read_status).routes, f"DELETE /responses/{PROVIDER_ID}": json_reply(DELETED)}, + ) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS + 120) +def test_a_delete_bills_a_finished_background_response_once_and_only_for_a_caller_that_may_address_it( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with _scripted(_deletable("queued")) as (handle,): + config: Final = _config(tmp_path, (handle,)) + with ( + _scratch_proxy(gateway, tmp_path, monkeypatch, config, POLL_OFF) as proxy, + proxy.scenario() as scenario, + ): + creator: Final = _team_member_key(scenario) + stranger: Final = _team_member_key(scenario) + response_id: Final = create_background_response(scenario, creator, handle.scenario_id) + _rescript((handle,), _deletable("completed")) + refused: Final = proxy.request("DELETE", f"/v1/responses/{response_id}", key=stranger) + assert refused.status_code == 403, refused.text + held(lambda: key_spend(creator), costs(0.0), holding=REFUSAL_HOLD_SECONDS) + assert [row["status"] for row in managed_rows(response_id)] == ["queued"] + deleted: Final = proxy.request("DELETE", f"/v1/responses/{response_id}", key=creator) + assert deleted.status_code == 200, deleted.text + assert object_value(deleted.json())["deleted"] is True + assert reported_cost(deleted.headers) == pytest.approx(PRICE) + assert_charged_exactly(lambda: key_spend(creator), PRICE) + _assert_each_billed_once_as_background_cost(creator, 1) + assert key_spend(stranger) == 0 + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] + + +@contextmanager +def _billing_claims_refused() -> Iterator[None]: + write_rows( + "CREATE FUNCTION audit_refuse_billing_claim() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN " + "IF NEW.status = 'completed' AND OLD.status IS DISTINCT FROM 'completed' THEN " + "RAISE EXCEPTION 'audit refuses the billing claim'; END IF; RETURN NEW; END $$", + (), + ) + write_rows( + 'CREATE TRIGGER audit_refuse_billing_claim BEFORE UPDATE ON "LiteLLM_ManagedObjectTable" ' + "FOR EACH ROW EXECUTE FUNCTION audit_refuse_billing_claim()", + (), + ) + try: + yield + finally: + write_rows('DROP TRIGGER audit_refuse_billing_claim ON "LiteLLM_ManagedObjectTable"', ()) + write_rows("DROP FUNCTION audit_refuse_billing_claim()", ()) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS + 180) +def test_a_finished_read_whose_bill_fails_to_settle_bills_the_creator_once_after_the_database_recovers( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with _scripted(background_routes("queued")) as (handle,): + config: Final = _config(tmp_path, (handle,)) + with ( + _scratch_proxy(gateway, tmp_path, monkeypatch, config, POLL_OFF) as proxy, + proxy.scenario() as scenario, + ): + key: Final = _team_member_key(scenario) + response_id: Final = create_background_response(scenario, key, handle.scenario_id) + _rescript((handle,), background_routes("completed")) + with _billing_claims_refused(): + unsettled: Final = read_response(proxy, key, response_id) + statuses: Final = [row["status"] for row in managed_rows(response_id)] + assert unsettled == Read("completed", "pong", 0.0) + assert statuses == ["queued"] + eventually(lambda: key_spend(key), costs(PRICE), seconds=SPEND_RETRY_SECONDS) + assert tuple(read_response(proxy, key, response_id) for _ in range(READS)) == free_reads("completed", READS) + assert_charged_exactly(lambda: key_spend(key), PRICE) + _assert_each_billed_once_as_background_cost(key, 1) + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] + + +REMAINING_KEY_BUDGET: Final = "litellm_remaining_api_key_budget_metric" +KEY_BUDGET: Final = 10.0 + + +def _remaining_key_budget(samples: Sequence[Sample], key: str) -> list[float]: + return [ + sample.value + for sample in samples + if sample.name == REMAINING_KEY_BUDGET and sample.labels.get("hashed_api_key") == hashed(key) + ] + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS + ONE_POLL_SECONDS + 60) +def test_a_background_response_the_poll_bills_lowers_its_creating_keys_remaining_budget_by_its_price( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with _scripted(background_routes("completed")) as (handle,): + config: Final = _config(tmp_path, (handle,), litellm_settings={"callbacks": ["prometheus"]}) + with _scratch_proxy(gateway, tmp_path, monkeypatch, config) as proxy, proxy.scenario() as scenario: + key: Final = scenario.key(max_budget=KEY_BUDGET) + response_id: Final = create_background_response(scenario, key, handle.scenario_id) + eventually( + lambda: [row["status"] for row in managed_rows(response_id)], + lambda statuses: statuses == ["completed"], + seconds=ONE_POLL_SECONDS, + ) + assert_charged_exactly(lambda: key_spend(key), PRICE) + eventually( + lambda: _remaining_key_budget(scrape(proxy), key), + lambda values: values == [pytest.approx(KEY_BUDGET - PRICE)], + seconds=30, + ) + + +CACHE_POLLING: Final[dict[str, JsonValue]] = {"responses": {"background_mode": {"polling_via_cache": "all"}}} +CACHE_POLLED_CREATES: Final[Mapping[str, tuple[dict[str, JsonValue], bool]]] = { + "streamed_with_background_in_extra_body": ({"stream": True, "extra_body": {"background": True}}, True), + "unstreamed": ({}, False), +} + + +def _cache_polled_create(proxy: Gateway, key: str, model: str, create: Mapping[str, JsonValue]) -> str: + created: Final = proxy.request( + "POST", + "/v1/responses", + {"model": model, "input": PROMPT, "background": True, "store": False, **create}, + key=key, + ) + assert created.status_code == 200, created.text + return string_value(object_value(created.json())["id"]) + + +@pytest.mark.timeout(OWNED_PROXY_SECONDS + 120) +@pytest.mark.parametrize("create", CACHE_POLLED_CREATES.keys()) +def test_a_background_response_the_proxy_polls_through_its_cache_bills_its_creator_once( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, create: str +) -> None: + body, provider_sees_background = CACHE_POLLED_CREATES[create] + with _scripted(streamed_reply(background=provider_sees_background)) as (handle,): + config: Final = _config(tmp_path, (handle,), litellm_settings=CACHE_POLLING) + with _scratch_proxy(gateway, tmp_path, monkeypatch, config, POLL_OFF) as proxy, proxy.scenario() as scenario: + key: Final = _team_member_key(scenario) + polling_id: Final = _cache_polled_create(proxy, key, handle.scenario_id, body) + finished: Final = eventually( + lambda: read_response(proxy, key, polling_id), + lambda read: read.status == "completed", + seconds=SPEND_RETRY_SECONDS, + ) + assert finished.text == "pong" + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert [read_response(proxy, key, polling_id).text for _ in range(READS)] == ["pong"] * READS + assert_charged_exactly(lambda: key_spend(key), PRICE) diff --git a/tests/integration/spend/test_background_response_poll_retirement.py b/tests/integration/spend/test_background_response_poll_retirement.py index f1b1e96901c..3b5864681fd 100644 --- a/tests/integration/spend/test_background_response_poll_retirement.py +++ b/tests/integration/spend/test_background_response_poll_retirement.py @@ -16,7 +16,7 @@ from integration.cost_calculation.cost_tracking_case import JsonResponse, Routed from pydantic import JsonValue from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE -from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.proxy.spend_tracking.spend_tracking_utils import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN INPUT_COST_PER_TOKEN: Final = 0.001 @@ -240,30 +240,21 @@ def test_a_provider_error_other_than_404_keeps_the_row_queued_for_retry(gateway: assert _row_status(unified_id) == [{"status": "queued"}] -def _polled_provider_id(row: dict[str, JsonValue]) -> str: - return ResponsesAPIRequestUtils.decode_responses_api_response_id(string_value(row["request_id"]))["response_id"] - - -def _poll_spend_rows(handle: ScenarioHandle) -> list[dict[str, JsonValue]]: - billed: Final = read_rows( - 'SELECT request_id, call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" ' - "WHERE metadata->>'internal_call_origin' = %s", - (BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,), +def _poll_spend_rows(unified_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" ' + "WHERE request_id = %s AND metadata->>'internal_call_origin' = %s", + (f"{unified_id}{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}", BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN), ) - return [ - {column: value for column, value in row.items() if column != "request_id"} - for row in billed - if _polled_provider_id(row) == f"resp_{handle.scenario_id}" - ] @pytest.mark.timeout(180) def test_a_completed_response_is_marked_completed_and_its_poll_is_billed(gateway: Gateway) -> None: with gateway.scenario() as scenario: - model, handle = _scripted_deployment(scenario, "done", _routes("queued", _json(_response_body("completed")))) + model, _ = _scripted_deployment(scenario, "done", _routes("queued", _json(_response_body("completed")))) unified_id: Final = _submit_background_response(scenario, scenario.key(), model) _await_status(unified_id, "completed", ONE_POLL_SECONDS) - spend_rows: Final = eventually(lambda: _poll_spend_rows(handle), lambda rows: len(rows) >= 1, seconds=70) + spend_rows: Final = eventually(lambda: _poll_spend_rows(unified_id), lambda rows: len(rows) >= 1, seconds=70) assert spend_rows[0] == { "call_type": "aget_responses", "status": "success", @@ -305,7 +296,7 @@ def test_responses_gone_at_the_provider_do_not_starve_a_newer_response_out_of_co ) for index in range(MAX_OBJECTS_PER_POLL_CYCLE) ) - completable_model, completable = _scripted_deployment( + completable_model, _ = _scripted_deployment( scenario, "done", _routes("queued", _json(_response_body("completed"))) ) completable_id: Final = _submit_background_response(scenario, key, completable_model) @@ -313,7 +304,7 @@ def test_responses_gone_at_the_provider_do_not_starve_a_newer_response_out_of_co assert [_row_status(gone_id) for gone_id in gone_ids] == [ [{"status": "stale_expired"}] ] * MAX_OBJECTS_PER_POLL_CYCLE - spend_rows: Final = eventually(lambda: _poll_spend_rows(completable), lambda rows: len(rows) >= 1, seconds=70) + spend_rows: Final = eventually(lambda: _poll_spend_rows(completable_id), lambda rows: len(rows) >= 1, seconds=70) assert spend_rows[0]["spend"] == pytest.approx( INPUT_TOKENS * INPUT_COST_PER_TOKEN + OUTPUT_TOKENS * OUTPUT_COST_PER_TOKEN ) diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 81efc36385e..91f6593575e 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -6048,3 +6048,103 @@ async def test_create_from_raw_provider_file_id_skips_managed_ownership_lookup(c assert isinstance(result, dict) and result[file_field] == "file-raw-provider-id", result assert not any("unified_file_id" in where for where in table.find_first_calls), table.find_first_calls + + +_MANAGED_READ_SALT_KEY: Final = "managed-read-unit-test-salt-key" + + +def _proxy_issued_response_id(encrypted: bool) -> str: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.types.llms.openai import ResponsesAPIResponse + + routed_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + "openai", "deployment-1", "resp_provider_job" + ) + if not encrypted: + return routed_id + issuer: Final = ResponsesIDSecurity( + general_settings_reader=lambda: {}, signing_key_reader=lambda: _MANAGED_READ_SALT_KEY + ) + return issuer._encrypt_response_id( + ResponsesAPIResponse(id=routed_id, created_at=1, output=[], status="queued"), + UserAPIKeyAuth(user_id="alice"), + ).id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) +@pytest.mark.parametrize("encrypted", [False, True], ids=["after_id_security", "before_id_security"]) +async def test_a_read_of_a_proxy_issued_response_id_is_marked_managed( + call_type: str, encrypted: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """The cost poll bills a background response the proxy issued, so a client read of that id is + marked managed and priced at zero, whichever order the ID security hook runs in.""" + from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setenv("LITELLM_SALT_KEY", _MANAGED_READ_SALT_KEY) + addressed_id: Final = _proxy_issued_response_id(encrypted) + + result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="alice"), + cache=MagicMock(), + data={"response_id": addressed_id, "litellm_metadata": {"user_api_key_user_id": "alice"}}, + call_type=call_type, + ) + + assert result == { + "response_id": addressed_id, + "litellm_metadata": {"user_api_key_user_id": "alice", MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True}, + } + + +@pytest.mark.asyncio +async def test_a_read_of_a_raw_provider_response_id_is_not_marked_managed() -> None: + """A raw provider id names no deployment, so the proxy never stored that response and its read + keeps pricing normally.""" + from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY + from litellm.proxy._types import UserAPIKeyAuth + + result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="alice"), + cache=MagicMock(), + data={"response_id": "resp_raw_provider_job", "metadata": {}}, + call_type="aget_responses", + ) + + assert isinstance(result, dict) + assert MANAGED_BACKGROUND_RESPONSE_METADATA_KEY not in result["metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "create_fields, marked", + [ + ({"stream": True, "background": True}, True), + ({"stream": True, "extra_body": {"background": True}}, True), + ({"stream": False, "background": True}, False), + ({"stream": True}, False), + ], + ids=["streamed_background", "streamed_background_in_extra_body", "not_streamed", "not_background"], +) +async def test_only_a_streamed_background_create_is_marked_for_the_managed_bill( + create_fields: Mapping[str, object], marked: bool +) -> None: + """A streamed background create is billed once through its managed row when its stream ends, + so only that create is marked free, and only in litellm_metadata, which never reaches the + provider.""" + from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY + from litellm.proxy._types import UserAPIKeyAuth + + result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="alice"), + cache=MagicMock(), + data={"model": "gpt-5.6", "input": "hello", "metadata": {"run": "nightly"}, "litellm_metadata": {}, **create_fields}, + call_type="aresponses", + ) + + assert isinstance(result, dict) + assert (MANAGED_BACKGROUND_RESPONSE_METADATA_KEY in result["litellm_metadata"]) is marked + assert result["metadata"] == {"run": "nightly"} diff --git a/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index cf305c660cd..18d2c8b9ace 100644 --- a/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -553,3 +553,59 @@ async def test_store_unified_object_id_omits_unset_columns(): assert "api_key" not in create_data assert "request_tags" not in create_data assert "unified-b" in store + + +@pytest.mark.asyncio +async def test_store_unified_object_id_drops_nul_characters_postgres_cannot_store(): + """Regression (spend loss): Postgres rejects a NUL inside jsonb, and a client can send one + in a background response's instructions or tags. The row the cost poll bills from must + still be written, with only the NUL characters removed.""" + from litellm.types.llms.openai import ResponsesAPIResponse + + instance, store = _in_memory_managed_files() + creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + response = ResponsesAPIResponse( + id="resp_nul", + created_at=1, + status="queued", + model="gpt-5.6", + instructions="be brief\x00", + metadata={"note\x00": "keep\x00 this"}, + output=[], + ) + + await instance.store_unified_object_id( + unified_object_id="unified-r", + file_object=response, + litellm_parent_otel_span=None, + model_object_id="resp_nul", + file_purpose="response", + user_api_key_dict=creator, + request_tags=["env\x00:prod"], + persist_attribution=True, + ) + + row = store["unified-r"] + assert "\\u0000" not in row["file_object"] + stored = json.loads(row["file_object"]) + assert stored == json.loads(response.model_dump_json().replace("\\u0000", "")) + assert stored["instructions"] == "be brief" + assert stored["metadata"] == {"note": "keep this"} + assert row["request_tags"].data == ["env:prod"] + + +@pytest.mark.asyncio +async def test_store_unified_object_id_keeps_the_provider_json_when_it_has_no_nul(): + instance, store = _in_memory_managed_files() + batch = _build_batch_response(batch_id="b", status="validating") + + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=batch, + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=UserAPIKeyAuth(user_id="alice", api_key="hash-alice"), + ) + + assert store["unified-b"]["file_object"] == batch.model_dump_json() diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index 7bf53b11cdb..01f20689e89 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -1499,6 +1499,51 @@ async def test_create_still_upserts_and_claims_attribution(): assert created["team_id"] == "team-alpha" +@pytest.mark.asyncio +async def test_a_stored_background_response_names_its_creators_end_user_and_holds_none_of_its_output(): + """Billing reads only a stored response's status and model, so neither the row nor its cache + entry keeps what the model generated, and the row names the end user the create billed to.""" + from litellm.types.llms.openai import ResponsesAPIResponse + + managed_files, mock_prisma = _make_object_store_instance() + generated_text = "the generated answer" + response = ResponsesAPIResponse( + id="resp_unified", + object="response", + status="completed", + created_at=1, + model="gpt-5.6", + output=[ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": generated_text, "annotations": []}], + } + ], + background=True, + ) + + await managed_files.store_unified_object_id( + unified_object_id="resp_unified", + file_object=response, + litellm_parent_otel_span=None, + model_object_id="resp_provider", + file_purpose="response", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-creator", end_user_id="customer-42", parent_otel_span=None), + persist_attribution=True, + ) + + written = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"] + cached = managed_files.internal_usage_cache.async_set_cache.await_args.kwargs["value"] + assert generated_text not in json.dumps(written, default=str) + assert generated_text not in json.dumps(cached, default=str) + assert json.loads(written["create"]["file_object"])["status"] == "completed" + assert written["create"]["end_user_id"] == "customer-42" + assert generated_text in response.model_dump_json() + + @pytest.mark.asyncio async def test_default_callers_still_create_their_rows(): """create_if_missing defaults to True, so the fine-tune, Responses and managed diff --git a/tests/unit/integrations/test_responses_background_cost.py b/tests/unit/integrations/test_responses_background_cost.py index 0d4218f2137..1bf95e35c4d 100644 --- a/tests/unit/integrations/test_responses_background_cost.py +++ b/tests/unit/integrations/test_responses_background_cost.py @@ -360,14 +360,14 @@ class TestCheckResponsesCost: mock_job.id = "job-123" mock_job.unified_object_id = "resp_test_id" mock_job.created_by = "test-user" + mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_id"} # Mock find_many to return the job mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( return_value=[mock_job] ) - # Mock update_many - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock() + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) # Create a completed response completed_response = ResponsesAPIResponse( @@ -406,7 +406,10 @@ class TestCheckResponsesCost: if c.kwargs.get("where", {}).get("id") is not None ] assert len(completion_calls) == 1 - assert completion_calls[0].kwargs["where"]["id"]["in"] == ["job-123"] + assert completion_calls[0].kwargs["where"] == { + "id": "job-123", + "status": {"in": ["queued", "in_progress"]}, + } assert completion_calls[0].kwargs["data"]["status"] == "completed" @pytest.mark.asyncio @@ -423,6 +426,7 @@ class TestCheckResponsesCost: mock_job.id = "job-456" mock_job.unified_object_id = "resp_failed" mock_job.created_by = "test-user" + mock_job.file_object = {"model": "gpt-4o", "id": "resp_failed"} mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( return_value=[mock_job] diff --git a/tests/unit/litellm_core_utils/test_internal_call_metadata.py b/tests/unit/litellm_core_utils/test_internal_call_metadata.py index 73923dc75a5..dfb232e96ce 100644 --- a/tests/unit/litellm_core_utils/test_internal_call_metadata.py +++ b/tests/unit/litellm_core_utils/test_internal_call_metadata.py @@ -1,11 +1,17 @@ """Unit tests for internal-call metadata forwarding: budget-reservation stripping and origin stamping.""" -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY from litellm.litellm_core_utils.internal_call_metadata import ( forwarded_internal_call_metadata, + is_unbilled_non_inference_call, + is_unbilled_non_inference_call_from_params, sanitized_forwardable_call_metadata, ) -from litellm.types.utils import SHADOW_EVAL_ROUTER_CALL_ORIGIN +from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, + SHADOW_EVAL_ROUTER_CALL_ORIGIN, +) PARENT = { "user_api_key": "sk-hash", @@ -119,3 +125,69 @@ class TestSubCallMetadataSanitization: assert sanitized_auth.team_id == "team-1" assert sanitized_auth.api_key == auth.api_key assert auth.budget_reservation == {"reserved_cost": 1.0} + + +class TestNonInferenceCallBilling: + """A read of a stored object replays the usage of the call that created it, so it is free. A + background create returns no usage, so its reads are where the job's spend shows up: the cost + poll's stamped read bills it, and any unmarked read (the SDK, a proxy with no database) bills + it too, because nothing else ever will.""" + + BACKGROUND = {"id": "resp_1", "background": True, "status": "completed"} + FOREGROUND = {"id": "resp_2", "background": False, "status": "completed"} + COST_POLL = {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + STATUS_POLL = {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN} + MANAGED_READ = {"user_api_key": "sk-hash", MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True} + + def test_unmarked_read_of_a_background_response_is_billed(self): + assert is_unbilled_non_inference_call("aget_responses", None, self.BACKGROUND) is False + assert is_unbilled_non_inference_call("aget_responses", {"user_api_key": "sk-hash"}, self.BACKGROUND) is False + assert is_unbilled_non_inference_call_from_params("aget_responses", {}, self.BACKGROUND) is False + + def test_managed_read_and_status_poll_of_a_background_response_are_free(self): + assert is_unbilled_non_inference_call("aget_responses", self.MANAGED_READ, self.BACKGROUND) is True + assert is_unbilled_non_inference_call("acancel_responses", self.MANAGED_READ, self.BACKGROUND) is True + assert is_unbilled_non_inference_call("aget_responses", self.STATUS_POLL, self.BACKGROUND) is True + assert is_unbilled_non_inference_call_from_params( + "aget_responses", {"metadata": self.MANAGED_READ}, self.BACKGROUND + ) is True + assert is_unbilled_non_inference_call_from_params( + "aget_responses", {"litellm_metadata": self.STATUS_POLL}, self.BACKGROUND + ) is True + + def test_a_marker_that_is_not_true_does_not_waive_the_charge(self): + not_true = {MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: "true"} + assert is_unbilled_non_inference_call("aget_responses", not_true, self.BACKGROUND) is False + + def test_cost_poll_read_is_billed_even_when_marked(self): + assert is_unbilled_non_inference_call("aget_responses", self.COST_POLL, self.BACKGROUND) is False + assert is_unbilled_non_inference_call( + "aget_responses", {**self.MANAGED_READ, **self.COST_POLL}, self.BACKGROUND + ) is False + assert is_unbilled_non_inference_call_from_params( + "aget_responses", {"litellm_metadata": self.COST_POLL}, self.FOREGROUND + ) is False + + def test_read_of_a_foreground_response_is_free_whatever_it_carries(self): + assert is_unbilled_non_inference_call("aget_responses", None, self.FOREGROUND) is True + assert is_unbilled_non_inference_call("aget_responses", {"user_api_key": "sk-hash"}, self.FOREGROUND) is True + assert is_unbilled_non_inference_call( + "aget_responses", {INTERNAL_CALL_ORIGIN_METADATA_KEY: SHADOW_EVAL_ROUTER_CALL_ORIGIN}, self.FOREGROUND + ) is True + + def test_inference_calls_are_billed_unless_the_proxy_tracks_their_background_response(self): + assert is_unbilled_non_inference_call("aresponses", None, self.BACKGROUND) is False + assert is_unbilled_non_inference_call("aresponses", self.MANAGED_READ, self.FOREGROUND) is False + assert is_unbilled_non_inference_call_from_params("aresponses", {}, self.BACKGROUND) is False + assert is_unbilled_non_inference_call_from_params("acompletion", {"metadata": self.MANAGED_READ}, None) is False + + def test_a_tracked_streamed_background_create_is_free_in_every_shape_its_stream_logs(self): + from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse + + finished = ResponsesAPIResponse(**{**self.BACKGROUND, "created_at": 1, "output": []}) + completed_event = ResponseCompletedEvent(type="response.completed", response=finished) + for logged in (finished, completed_event, completed_event.model_dump(), self.BACKGROUND): + assert is_unbilled_non_inference_call("aresponses", self.MANAGED_READ, logged) is True + assert is_unbilled_non_inference_call_from_params( + "aresponses", {"litellm_metadata": self.MANAGED_READ}, logged + ) is True diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 870378e93f6..fbe3766f654 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -7071,6 +7071,21 @@ class TestNonInferenceCallTypesAreNotBilled: assert payload is not None assert payload["total_tokens"] == 6000 + @pytest.mark.parametrize( + "litellm_metadata", + [ + {"_managed_background_response": True}, + {"internal_call_origin": "background_response_status_poll"}, + ], + ids=["managed_read", "status_poll"], + ) + def test_managed_read_of_a_background_response_is_free(self, litellm_metadata): + """The proxy's cost poll bills a background response it tracks once, so a read the managed files + hook marks, or the poll's own status check, must not bill the same tokens again.""" + logging_obj = self._logging_obj("aget_responses", litellm_metadata=litellm_metadata) + cost = logging_obj.response_cost_calculator(result=self._retrieved_response(background=True)) + assert cost == 0.0 + def test_reading_a_foreground_response_is_still_free(self): """Guards the test above against a blanket exemption: an explicit background=false read was already billed by its create and must stay at zero.""" diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 58d6a44e4a4..3b9e7a30302 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -2366,17 +2366,21 @@ class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" - def _instance(self, key_row=None, team_row=None, user_row=None): - from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost - + def _prisma(self, key_row=None, team_row=None, user_row=None): prisma = MagicMock() prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_row) prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) - return CheckBatchCost( - proxy_logging_obj=MagicMock(), - prisma_client=prisma, - llm_router=MagicMock(), + return prisma + + def _instance(self, key_row=None, team_row=None, user_row=None, prisma=None): + from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( + ManagedObjectCreatorAttribution, + ) + + return ManagedObjectCreatorAttribution( + prisma_client=prisma if prisma is not None else self._prisma(key_row, team_row, user_row), + poller_name="CheckBatchCost", ) def _job(self, **overrides): @@ -2402,7 +2406,7 @@ class TestBatchCostAttribution: user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None), ) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert metadata["user_api_key"] == "hash-alice" assert metadata["user_api_key_hash"] == "hash-alice" @@ -2419,7 +2423,7 @@ class TestBatchCostAttribution: instance = self._instance() job = self._job(api_key=None, request_tags=None) - metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + metadata = await instance.batch_metadata(job, "batch-1") assert metadata["user_api_key"] is None assert metadata["user_api_key_user_id"] == "alice" @@ -2432,15 +2436,32 @@ class TestBatchCostAttribution: a None user_id) and the key hash still drives key-level attribution.""" from types import SimpleNamespace - instance = self._instance(key_row=SimpleNamespace(key_alias="svc-key")) + prisma = self._prisma(key_row=SimpleNamespace(key_alias="svc-key")) + instance = self._instance(prisma=prisma) job = self._job(created_by=None) - metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + metadata = await instance.batch_metadata(job, "batch-1") assert metadata["user_api_key"] == "hash-alice" assert metadata["user_api_key_user_id"] is None assert metadata["user_api_key_alias"] == "svc-key" - instance.prisma_client.db.litellm_usertable.find_unique.assert_not_called() + prisma.db.litellm_usertable.find_unique.assert_not_called() + + @pytest.mark.asyncio + async def test_batch_metadata_keeps_the_creator_column_and_adds_no_budgets(self): + """The batch cost event keeps the user id and fields the batch poll has always written, + even though background responses read budgets and drop key-scoped creators off the same rows.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="svc-key", spend=4.0, max_budget=5.0), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None, spend=40.0, max_budget=50.0), + ) + + metadata = await instance.batch_metadata(self._job(created_by="key:hash-alice"), "batch-1") + + assert metadata["user_api_key_user_id"] == "key:hash-alice" + assert [field for field in metadata if field.endswith(("_spend", "_max_budget"))] == [] @pytest.mark.asyncio async def test_metadata_drops_non_string_tags(self): @@ -2449,7 +2470,7 @@ class TestBatchCostAttribution: instance = self._instance() job = self._job(request_tags=["env:prod", 7, None, "team:ml"]) - metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + metadata = await instance.batch_metadata(job, "batch-1") assert metadata["tags"] == ["env:prod", "team:ml"] @@ -2457,10 +2478,11 @@ class TestBatchCostAttribution: async def test_key_alias_lookup_failure_does_not_break_attribution(self): """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" - instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + prisma = self._prisma() + prisma.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + instance = self._instance(prisma=prisma) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert metadata["user_api_key"] == "hash-alice" assert metadata.get("user_api_key_alias") is None @@ -2469,7 +2491,7 @@ class TestBatchCostAttribution: async def test_cli_session_batch_keeps_its_alias_without_a_key_row(self): instance = self._instance(key_row=None) - metadata = await instance._build_creator_attribution_metadata(self._job(api_key="cli-session-alice"), "batch-1") + metadata = await instance.batch_metadata(self._job(api_key="cli-session-alice"), "batch-1") assert metadata["user_api_key"] == "cli-session-alice" assert metadata["user_api_key_alias"] == "cli-session-alice" @@ -2478,7 +2500,7 @@ class TestBatchCostAttribution: async def test_raw_cli_session_token_on_a_legacy_batch_row_is_not_treated_as_the_alias(self): instance = self._instance(key_row=None) - metadata = await instance._build_creator_attribution_metadata( + metadata = await instance.batch_metadata( self._job(api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa"), "batch-1" ) @@ -2496,7 +2518,7 @@ class TestBatchCostAttribution: user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"), ) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert metadata["user_api_key_alias"] == "Alice Chen" assert metadata["user_api_key"] == "hash-alice" @@ -2512,7 +2534,7 @@ class TestBatchCostAttribution: user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"), ) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert metadata["user_api_key_alias"] == "Alice Chen" @@ -2527,7 +2549,7 @@ class TestBatchCostAttribution: user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"), ) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert metadata["user_api_key_alias"] == "prod-key" @@ -2542,7 +2564,7 @@ class TestBatchCostAttribution: team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), ) - metadata = await instance._build_creator_attribution_metadata(self._job(org_id="org-at-creation"), "batch-1") + metadata = await instance.batch_metadata(self._job(org_id="org-at-creation"), "batch-1") assert metadata["user_api_key_org_id"] == "org-at-creation" @@ -2557,7 +2579,7 @@ class TestBatchCostAttribution: team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), ) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert metadata["user_api_key_org_id"] == "org-42" @@ -2572,7 +2594,7 @@ class TestBatchCostAttribution: team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), ) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert metadata["user_api_key_org_id"] == "org-team" @@ -2582,12 +2604,11 @@ class TestBatchCostAttribution: organization: the two lookups fail independently, so org spend still lands.""" from types import SimpleNamespace - instance = self._instance( - team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), - ) - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + prisma = self._prisma(team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team")) + prisma.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + instance = self._instance(prisma=prisma) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert metadata["user_api_key_org_id"] == "org-team" @@ -2602,7 +2623,7 @@ class TestBatchCostAttribution: team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None), ) - metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + metadata = await instance.batch_metadata(self._job(), "batch-1") assert "user_api_key_org_id" not in metadata @@ -2625,7 +2646,7 @@ class TestBatchCostAttribution: key_row=SimpleNamespace(key_alias="prod-key"), user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None), ) - metadata = await instance._build_creator_attribution_metadata(self._job(api_key=token_hash), "batch-1") + metadata = await instance.batch_metadata(self._job(api_key=token_hash), "batch-1") assert metadata["user_api_key"] == token_hash assert metadata["user_api_key_hash"] == token_hash 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 8ba3f0b000a..f606658fa09 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -3,15 +3,48 @@ Unit tests for CheckResponsesCost class """ import asyncio +import dataclasses from collections.abc import Mapping +from dataclasses import dataclass from datetime import datetime +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest +from prisma.errors import UniqueViolationError +from litellm_enterprise.proxy.common_utils.check_responses_cost import BilledBackgroundResponse, BilledCharge from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + +class _RecordingBilledLogger: + """Stands in for the spend logger so a test reads the billing events a poller wrote, and can + make the next writes raise.""" + + def __init__(self, failures: tuple[BaseException, ...] = (), cost: float = 0.25) -> None: + self.billed: tuple[BilledBackgroundResponse, ...] = () + self._failures = failures + self.charge: Final = BilledCharge(cost=cost, logging_obj=MagicMock()) + + async def __call__(self, billed: BilledBackgroundResponse) -> BilledCharge: + await asyncio.sleep(0) + if self._failures: + failure, *remaining = self._failures + self._failures = tuple(remaining) + raise failure + self.billed = (*self.billed, billed) + return self.charge + + +def _deployment(model_id: str) -> Deployment: + return Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id=model_id), + ) class _RecordingManagedObjectTable: @@ -27,6 +60,114 @@ class _RecordingManagedObjectTable: return len(self.rows) +@dataclass(frozen=True) +class _StoredResponse: + id: str + unified_object_id: str + status: str + model_object_id: str = "" + file_purpose: str = "response" + file_object: Mapping[str, object] = dataclasses.field(default_factory=lambda: {"model": "gpt-5"}) + created_by: str | None = None + api_key: str | None = None + team_id: str | None = None + org_id: str | None = None + request_tags: object = None + end_user_id: str | None = None + + +class _InMemoryManagedObjectTable: + """Applies where filters and status writes the way Postgres does, so two pollers sharing it + race over the same rows.""" + + def __init__(self, rows: tuple[_StoredResponse, ...]) -> None: + self._rows: Final = rows + self.statuses: Final = {row.id: row.status for row in rows} + self.model_object_ids: Final = {row.id: row.model_object_id for row in rows} + + def _matches(self, row: _StoredResponse, where: Mapping[str, object]) -> bool: + def value_of(field: str) -> object: + if field == "status": + return self.statuses[row.id] + if field == "model_object_id": + return self.model_object_ids[row.id] + return getattr(row, field) + + def holds(field: str, condition: object) -> bool: + if field == "OR" and isinstance(condition, list): + return any(self._matches(row, alternative) for alternative in condition) + if isinstance(condition, dict): + return value_of(field) in condition["in"] + return value_of(field) == condition + + return all(holds(field, condition) for field, condition in where.items()) + + def _current(self, where: Mapping[str, object]) -> list[_StoredResponse]: + return [ + dataclasses.replace(row, status=self.statuses[row.id], model_object_id=self.model_object_ids[row.id]) + for row in self._rows + if self._matches(row, where) + ] + + async def find_many(self, where: Mapping[str, object], **_: object) -> list[_StoredResponse]: + return self._current(where) + + async def find_first(self, where: Mapping[str, object]) -> _StoredResponse | None: + return next(iter(self._current(where)), None) + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, str]) -> int: + matched: Final = self._current(where) + if "model_object_id" in data and any( + other.id not in {row.id for row in matched} and value == data["model_object_id"] + for other, value in ((other, self.model_object_ids[other.id]) for other in self._rows) + ): + raise UniqueViolationError({"user_facing_error": {"error_code": "P2002"}}) + self.statuses.update({row.id: data["status"] for row in matched if "status" in data}) + self.model_object_ids.update({row.id: data["model_object_id"] for row in matched if "model_object_id" in data}) + return len(matched) + + +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.db.litellm_managedobjecttable = table + poller: Final = CheckResponsesCost( + proxy_logging_obj=MagicMock(), + prisma_client=prisma_client, + llm_router=MagicMock(), + log_billed_response=billed_logger, + ) + poller._expire_stale_rows = AsyncMock(return_value=0) + return poller + + +def _terminal_response(status: str = "completed") -> ResponsesAPIResponse: + return ResponsesAPIResponse( + id="resp_upstream", + object="response", + status=status, + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150), + ) + + +def _claim(job_id: str) -> tuple[Mapping[str, object], Mapping[str, object]]: + return ({"id": job_id, "status": {"in": ["queued", "in_progress"]}}, {"status": "completed"}) + + +_READ_OF_A_BACKGROUND_RESPONSE: Final = MappingProxyType({"background": True}) + + +def _is_billed(litellm_metadata: Mapping[str, object]) -> bool: + from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call + + return not is_unbilled_non_inference_call( + "aget_responses", dict(litellm_metadata), _READ_OF_A_BACKGROUND_RESPONSE + ) + + class TestCheckResponsesCost: """Test suite for CheckResponsesCost class""" @@ -50,12 +191,16 @@ class TestCheckResponsesCost: """Create a mock LLM Router""" router = MagicMock() router.aget_responses = AsyncMock() - router.get_deployment = MagicMock() + router.get_deployment = MagicMock(side_effect=lambda model_id: _deployment(model_id)) return router + @pytest.fixture + def billed_logger(self) -> _RecordingBilledLogger: + return _RecordingBilledLogger() + @pytest.fixture def check_responses_cost_instance( - self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router, billed_logger ): """Create a CheckResponsesCost instance with mocked dependencies""" from litellm_enterprise.proxy.common_utils.check_responses_cost import ( @@ -66,6 +211,7 @@ class TestCheckResponsesCost: proxy_logging_obj=mock_proxy_logging_obj, prisma_client=mock_prisma_client, llm_router=mock_llm_router, + log_billed_response=billed_logger, ) # Mock _expire_stale_rows (raw SQL) so _cleanup_stale_managed_objects # succeeds without a real DB. Individual tests can override this. @@ -121,7 +267,7 @@ class TestCheckResponsesCost: @pytest.mark.asyncio async def test_check_responses_cost_with_completed_response( - self, check_responses_cost_instance, mock_prisma_client, mock_llm_router + self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, billed_logger ): """Test check_responses_cost with a completed response""" # Mock job with response ID @@ -150,23 +296,22 @@ class TestCheckResponsesCost: ) mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) - # Run the check with mocked litellm.aget_responses with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: mock_aget.return_value = mock_response await check_responses_cost_instance.check_responses_cost() - # update_many should only contain the job completion call calls = ( mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list ) - assert len(calls) == 1 - completion_call = calls[0] - assert completion_call[1]["data"]["status"] == "completed" - assert completion_call[1]["where"]["id"]["in"] == ["job-123"] + assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-123")] + assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False] + assert [(billed.response.id, billed.response.usage) for billed in billed_logger.billed] == [ + ("resp_test_123", mock_response.usage) + ] @pytest.mark.asyncio async def test_check_responses_cost_with_failed_response( @@ -399,7 +544,8 @@ class TestCheckResponsesCost: mock_job.id = "job-404" mock_job.file_object = {"model": "gpt-4o", "id": encoded_response_id} - mock_llm_router.get_deployment.return_value = {"model_id": "deployment-404"} + mock_llm_router.get_deployment.side_effect = None + mock_llm_router.get_deployment.return_value = _deployment("deployment-404") mock_llm_router.aget_responses = AsyncMock( side_effect=litellm.NotFoundError( message="Response with id 'resp_upstream_404' not found.", model="gpt-5", llm_provider="openai" @@ -461,7 +607,8 @@ class TestCheckResponsesCost: mock_job.id = "job-404-mapped" mock_job.file_object = {"model": "gpt-5.5", "id": encoded_response_id} - mock_llm_router.get_deployment.return_value = {"model_id": "deployment-404-mapped"} + mock_llm_router.get_deployment.side_effect = None + mock_llm_router.get_deployment.return_value = _deployment("deployment-404-mapped") mock_llm_router.aget_responses = AsyncMock(side_effect=mapped.value) table = _RecordingManagedObjectTable(rows=(mock_job,)) @@ -493,6 +640,7 @@ class TestCheckResponsesCost: mock_job.id = "job-404-fallback" mock_job.file_object = {"model": "gpt-4o", "id": encoded_response_id} + mock_llm_router.get_deployment.side_effect = None mock_llm_router.get_deployment.return_value = None table = _RecordingManagedObjectTable(rows=(mock_job,)) @@ -531,7 +679,8 @@ class TestCheckResponsesCost: mock_job.id = "job-404-other" mock_job.file_object = {"model": "azure-gpt-5", "id": encoded_response_id} - mock_llm_router.get_deployment.return_value = {"model_id": "deployment-renamed"} + mock_llm_router.get_deployment.side_effect = None + mock_llm_router.get_deployment.return_value = _deployment("deployment-renamed") mock_llm_router.aget_responses = AsyncMock( side_effect=litellm.NotFoundError( message="Resource not found", model="gpt-5", llm_provider="azure" @@ -550,7 +699,7 @@ class TestCheckResponsesCost: @pytest.mark.asyncio async def test_check_responses_cost_deployment_lookup_error_skips_only_that_job( - self, check_responses_cost_instance, mock_prisma_client, mock_llm_router + self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, billed_logger ): """A deployment lookup that raises skips its own job and the cycle still records the next job.""" from litellm.responses.utils import ResponsesAPIRequestUtils @@ -576,7 +725,7 @@ class TestCheckResponsesCost: mock_llm_router.get_deployment.side_effect = [ Exception("Model invalid format - "), - {"model_id": "deployment-healthy"}, + _deployment("deployment-healthy"), ] mock_llm_router.aget_responses = AsyncMock( return_value=ResponsesAPIResponse( @@ -596,9 +745,11 @@ class TestCheckResponsesCost: await check_responses_cost_instance.check_responses_cost() mock_sdk_aget.assert_not_called() - mock_llm_router.aget_responses.assert_awaited_once() - assert mock_llm_router.aget_responses.call_args.kwargs["response_id"] == healthy_response_id - assert table.updates == (({"id": {"in": ["job-healthy"]}}, {"status": "completed"}),) + assert [call.kwargs["response_id"] for call in mock_llm_router.aget_responses.call_args_list] == [ + healthy_response_id, + ] + assert table.updates == (_claim("job-healthy"),) + assert [billed.response.id for billed in billed_logger.billed] == [healthy_response_id] @pytest.mark.asyncio async def test_check_responses_cost_non_404_error_keeps_row_for_retry( @@ -620,7 +771,8 @@ class TestCheckResponsesCost: mock_job.id = "job-500" mock_job.file_object = {"model": "gpt-4o", "id": encoded_response_id} - mock_llm_router.get_deployment.return_value = {"model_id": "deployment-500"} + mock_llm_router.get_deployment.side_effect = None + mock_llm_router.get_deployment.return_value = _deployment("deployment-500") mock_llm_router.aget_responses = AsyncMock( side_effect=litellm.InternalServerError( message="boom", model="gpt-5", llm_provider="openai" @@ -639,7 +791,7 @@ class TestCheckResponsesCost: @pytest.mark.asyncio async def test_check_responses_cost_multiple_jobs( - self, check_responses_cost_instance, mock_prisma_client + self, check_responses_cost_instance, mock_prisma_client, billed_logger ): """Test check_responses_cost with multiple jobs""" # Mock multiple jobs @@ -702,25 +854,23 @@ class TestCheckResponsesCost: ) mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) - # Run the check with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: mock_aget.side_effect = [mock_response1, mock_response2, mock_response3] await check_responses_cost_instance.check_responses_cost() - # update_many should only contain the job completion call calls = ( mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list ) - assert len(calls) == 1 - completion_call = calls[0] - assert len(completion_call[1]["where"]["id"]["in"]) == 2 - assert "job-1" in completion_call[1]["where"]["id"]["in"] - assert "job-3" in completion_call[1]["where"]["id"]["in"] - assert "job-2" not in completion_call[1]["where"]["id"]["in"] + assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-1"), _claim("job-3")] + assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False, False, False] + assert [(billed.response.id, billed.response.usage) for billed in billed_logger.billed] == [ + ("resp_test_1", mock_response1.usage), + ("resp_test_3", mock_response3.usage), + ] @pytest.mark.asyncio async def test_encoded_response_id_is_fetched_through_router( @@ -787,9 +937,7 @@ class TestCheckResponsesCost: calls = ( mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list ) - assert len(calls) == 1 - assert calls[0][1]["data"]["status"] == "completed" - assert calls[0][1]["where"]["id"]["in"] == ["job-router"] + assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-router")] @pytest.mark.asyncio async def test_encrypted_response_id_is_fetched_through_router( @@ -860,8 +1008,7 @@ class TestCheckResponsesCost: calls = ( mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list ) - assert len(calls) == 1 - assert calls[0][1]["where"]["id"]["in"] == ["job-encrypted"] + assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-encrypted")] @pytest.mark.asyncio async def test_response_id_without_model_id_uses_sdk( @@ -955,9 +1102,7 @@ class TestCheckResponsesCost: calls = ( mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list ) - assert len(calls) == 1 - assert calls[0][1]["data"]["status"] == "completed" - assert calls[0][1]["where"]["id"]["in"] == ["job-missing-deployment"] + assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-missing-deployment")] @pytest.mark.asyncio async def test_check_responses_cost_with_incomplete_response( @@ -993,9 +1138,7 @@ class TestCheckResponsesCost: calls = ( mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list ) - assert len(calls) == 1 - assert calls[0][1]["data"]["status"] == "completed" - assert calls[0][1]["where"]["id"]["in"] == ["job-incomplete"] + assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-incomplete")] @pytest.mark.asyncio async def test_check_responses_cost_no_model_in_file_object( @@ -1012,7 +1155,7 @@ class TestCheckResponsesCost: return_value=[mock_job] ) mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_response = MagicMock() @@ -1028,12 +1171,12 @@ class TestCheckResponsesCost: assert "model_group" not in call_kwargs.get("litellm_metadata", {}) @pytest.mark.asyncio - async def test_poll_stamps_internal_call_origin_so_the_read_is_billed( - self, check_responses_cost_instance, mock_prisma_client + async def test_poll_stamps_the_billing_event_so_it_is_priced_and_its_status_read_is_free( + self, check_responses_cost_instance, mock_prisma_client, billed_logger ): - """A background create returns queued with no usage, so this poll's retrieval is the only - place the job's spend is ever seen. Without the origin stamp it is priced at zero like a - user-facing read (LIT-5602) and the job is never billed.""" + """A background create returns queued with no usage, so the event the poll logs from the + finished response is the only place the job's spend is ever seen. Without the origin stamp + it is priced at zero like a user-facing read (LIT-5602) and the job is never billed.""" from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.litellm_core_utils.internal_call_metadata import ( is_unbilled_non_inference_call, @@ -1049,7 +1192,7 @@ class TestCheckResponsesCost: return_value=[mock_job] ) mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_response = MagicMock() @@ -1059,8 +1202,549 @@ class TestCheckResponsesCost: mock_aget.return_value = mock_response await check_responses_cost_instance.check_responses_cost() - metadata = mock_aget.call_args[1]["litellm_metadata"] - foreground_read = {"background": False} - assert metadata[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "background_response_cost_poll" - assert is_unbilled_non_inference_call("aget_responses", metadata, foreground_read) is False - assert is_unbilled_non_inference_call("aget_responses", None, foreground_read) is True + assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False] + (billed,) = billed_logger.billed + assert billed.metadata[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "background_response_cost_poll" + assert ( + is_unbilled_non_inference_call("aget_responses", dict(billed.metadata), _READ_OF_A_BACKGROUND_RESPONSE) + is False + ) + + @pytest.mark.asyncio + async def test_poll_bills_the_response_to_the_creating_key_team_org_and_tags( + self, check_responses_cost_instance, mock_prisma_client, billed_logger + ): + """The poll's billing event is the one billed event of a background job, so its spend row must name + the key, team, org, user, and tags of the create it settles, read off the managed object row, + never the poller itself.""" + from types import SimpleNamespace + + mock_job = MagicMock() + mock_job.unified_object_id = "resp_test_attributed" + mock_job.id = "job-attributed" + mock_job.created_by = "alice" + mock_job.api_key = "hash-alice" + mock_job.team_id = "team-alpha" + mock_job.org_id = "org-1" + mock_job.request_tags = ["env:prod"] + mock_job.file_object = {"model": "gpt-5", "id": "resp_test_attributed"} + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_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) + ) + mock_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) + ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=SimpleNamespace(user_email="alice@example.com", user_alias=None, spend=1.0, max_budget=2.0) + ) + + with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: + mock_aget.return_value = MagicMock(status="completed") + await check_responses_cost_instance.check_responses_cost() + + (billed,) = billed_logger.billed + metadata = billed.metadata + assert metadata["user_api_key"] == "hash-alice" + assert metadata["user_api_key_hash"] == "hash-alice" + assert metadata["user_api_key_user_id"] == "alice" + assert metadata["user_api_key_user_email"] == "alice@example.com" + assert metadata["user_api_key_team_id"] == "team-alpha" + assert metadata["user_api_key_org_id"] == "org-1" + assert metadata["user_api_key_alias"] == "prod-key" + assert metadata["user_api_key_team_alias"] == "Team Alpha" + assert (metadata["user_api_key_spend"], metadata["user_api_key_max_budget"]) == (4.0, 5.0) + assert (metadata["user_api_key_team_spend"], metadata["user_api_key_team_max_budget"]) == (40.0, 50.0) + assert (metadata["user_api_key_user_spend"], metadata["user_api_key_user_max_budget"]) == (1.0, 2.0) + assert metadata["tags"] == ["env:prod"] + assert metadata["model_group"] == "gpt-5" + + +@pytest.mark.asyncio +async def test_two_pollers_racing_over_one_completed_response_bill_it_once(): + """Every worker and pod runs its own poll, so two of them can pick up the same completed + response in one window; only the one that moves the row out of a pending status bills it.""" + table: Final = _InMemoryManagedObjectTable( + rows=(_StoredResponse(id="job-race", unified_object_id="resp_race", status="in_progress"),) + ) + reads: Final[list[Mapping[str, object]]] = [] + billed_logger: Final = _RecordingBilledLogger() + + async def provider_read(response_id: str, litellm_metadata: Mapping[str, object]) -> ResponsesAPIResponse: + reads.append(litellm_metadata) + await asyncio.sleep(0) + return _terminal_response() + + with patch("litellm.aget_responses", side_effect=provider_read): + await asyncio.gather( + _poller_over(table, billed_logger).check_responses_cost(), + _poller_over(table, billed_logger).check_responses_cost(), + ) + + assert [_is_billed(metadata) for metadata in reads] == [False, False] + assert len(billed_logger.billed) == 1 + assert table.statuses == {"job-race": "completed"} + + +@pytest.mark.asyncio +async def test_a_failed_billing_event_puts_the_response_back_for_the_next_poll(): + table: Final = _InMemoryManagedObjectTable( + rows=(_StoredResponse(id="job-retry", unified_object_id="resp_retry", status="queued"),) + ) + billed_logger: Final = _RecordingBilledLogger(failures=(RuntimeError("spend log unavailable"),)) + poller: Final = _poller_over(table, billed_logger) + with patch("litellm.aget_responses", new_callable=AsyncMock, return_value=_terminal_response()): + await poller.check_responses_cost() + assert table.statuses == {"job-retry": "queued"} + assert billed_logger.billed == () + await poller.check_responses_cost() + + assert len(billed_logger.billed) == 1 + assert table.statuses == {"job-retry": "completed"} + + +@pytest.mark.asyncio +async def test_only_the_first_read_that_sees_a_response_finished_bills_it(): + """A client polling its background response bills it on the first read that sees it finished, + under the creating key, so it is billed even if the provider forgets it before the poll runs; + reads before and after that one stay free.""" + table: Final = _InMemoryManagedObjectTable( + rows=( + _StoredResponse( + id="job-read", + unified_object_id="resp_read", + status="queued", + created_by="alice", + api_key="hash-alice", + team_id="team-alpha", + ), + ) + ) + billed_logger: Final = _RecordingBilledLogger() + poller: Final = _poller_over(table, billed_logger) + poller.prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + poller.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + poller.prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + finished: Final = _terminal_response("completed") + + with patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")): + costs: Final = ( + await poller.bill_if_finished(response_id="resp_read", response=_terminal_response("in_progress")), + await poller.bill_if_finished(response_id="resp_read", response=finished), + await poller.bill_if_finished(response_id="resp_read", response=finished), + ) + await poller.check_responses_cost() + + assert costs == (None, billed_logger.charge, None) + (billed,) = billed_logger.billed + assert billed.response.usage == finished.usage + assert billed.response.id == "resp_read" + assert billed.metadata["user_api_key_hash"] == "hash-alice" + assert billed.metadata["user_api_key_team_id"] == "team-alpha" + assert table.statuses == {"job-read": "completed"} + + +_ROTATION_SALT_KEY: Final = "check-responses-cost-unit-test-salt" + + +@pytest.mark.asyncio +async def test_a_read_through_a_reissued_response_id_bills_the_stored_response(monkeypatch: pytest.MonkeyPatch): + """The proxy re-encrypts a response id on every read, so a client polling with the id its last + read returned still bills the row stored under the id the create returned.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity + + monkeypatch.setenv("LITELLM_SALT_KEY", _ROTATION_SALT_KEY) + issuer: Final = ResponsesIDSecurity(general_settings_reader=lambda: {}, signing_key_reader=lambda: _ROTATION_SALT_KEY) + + def issued_id() -> str: + return issuer._encrypt_response_id( + ResponsesAPIResponse(id="resp_provider_job", created_at=1, output=[], status="queued"), + UserAPIKeyAuth(user_id="alice"), + ).id + + created_id: Final = issued_id() + reissued_id: Final = issued_id() + assert created_id != reissued_id + table: Final = _InMemoryManagedObjectTable( + rows=( + _StoredResponse( + id="job-rotated", + unified_object_id=created_id, + model_object_id="resp_provider_job", + status="queued", + ), + ) + ) + + billed_logger: Final = _RecordingBilledLogger() + with patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")): + await _poller_over(table, billed_logger).bill_if_finished(response_id=reissued_id, response=_terminal_response()) + + assert len(billed_logger.billed) == 1 + assert table.statuses == {"job-rotated": "completed"} + + +def _issued_ids(monkeypatch: pytest.MonkeyPatch, count: int) -> tuple[str, ...]: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity + + monkeypatch.setenv("LITELLM_SALT_KEY", _ROTATION_SALT_KEY) + issuer: Final = ResponsesIDSecurity(general_settings_reader=lambda: {}, signing_key_reader=lambda: _ROTATION_SALT_KEY) + return tuple( + issuer._encrypt_response_id( + ResponsesAPIResponse(id="resp_provider_job", created_at=1, output=[], status="queued"), + UserAPIKeyAuth(user_id="alice"), + ).id + for _ in range(count) + ) + + +def _stored_before_the_upgrade(created_id: str) -> _StoredResponse: + return _StoredResponse( + id="job-legacy", unified_object_id=created_id, model_object_id=created_id, status="queued" + ) + + +@pytest.mark.asyncio +async def test_a_response_stored_before_the_upgrade_bills_once_whichever_issued_id_reads_it( + monkeypatch: pytest.MonkeyPatch, +): + """A row the old proxy stored keys the response by the encrypted id the create returned, so the + read that sends that id rekeys it by the provider id, and a read through any later id finds it + already billed.""" + created_id, reissued_id = _issued_ids(monkeypatch, 2) + table: Final = _InMemoryManagedObjectTable(rows=(_stored_before_the_upgrade(created_id),)) + billed_logger: Final = _RecordingBilledLogger() + poller: Final = _poller_over(table, billed_logger) + track_untracked: Final = AsyncMock(return_value=False) + + with patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")): + await poller.bill_if_finished(response_id=created_id, response=_terminal_response()) + await poller.bill_if_finished( + response_id=reissued_id, response=_terminal_response(), track_untracked=track_untracked + ) + + assert len(billed_logger.billed) == 1 + assert table.model_object_ids == {"job-legacy": "resp_provider_job"} + track_untracked.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "read_rows, billed_by_the_poll", + [((), 1), ((_StoredResponse(id="job-read", unified_object_id="read-id", model_object_id="resp_provider_job", status="completed"),), 0)], + ids=["no_read_yet", "a_read_already_billed_it"], +) +async def test_the_poll_bills_a_response_stored_before_the_upgrade_only_when_no_read_billed_it_first( + monkeypatch: pytest.MonkeyPatch, read_rows: tuple[_StoredResponse, ...], billed_by_the_poll: int +): + """A read through a later id misses the old row and tracks the response under its provider id, + so the poll retires the old row unbilled when that row exists and bills it once otherwise.""" + (created_id,) = _issued_ids(monkeypatch, 1) + table: Final = _InMemoryManagedObjectTable(rows=(_stored_before_the_upgrade(created_id), *read_rows)) + billed_logger: Final = _RecordingBilledLogger() + + with patch("litellm.aget_responses", new_callable=AsyncMock, return_value=_terminal_response()): + await _poller_over(table, billed_logger).check_responses_cost() + await _poller_over(table, billed_logger).check_responses_cost() + + assert len(billed_logger.billed) == billed_by_the_poll + assert table.statuses["job-legacy"] == "completed" + + +@pytest.mark.asyncio +async def test_a_cancelled_billing_event_puts_the_response_back_for_the_next_poll(): + """A client that disconnects mid-read cancels the billing event, so the claim is released and + the next poll still bills the response.""" + table: Final = _InMemoryManagedObjectTable( + rows=(_StoredResponse(id="job-cancelled", unified_object_id="resp_cancelled", status="in_progress"),) + ) + billed_logger: Final = _RecordingBilledLogger(failures=(asyncio.CancelledError(),)) + + with pytest.raises(asyncio.CancelledError): + await _poller_over(table, billed_logger).bill_if_finished( + response_id="resp_cancelled", response=_terminal_response() + ) + + assert billed_logger.billed == () + assert table.statuses == {"job-cancelled": "in_progress"} + + +@pytest.mark.asyncio +async def test_the_billing_event_carries_the_creators_end_user(): + """The charge belongs to the end user the create named, whoever reads the response later.""" + table: Final = _InMemoryManagedObjectTable( + rows=( + _StoredResponse( + id="job-end-user", unified_object_id="resp_end_user", status="queued", end_user_id="customer-42" + ), + ) + ) + billed_logger: Final = _RecordingBilledLogger() + + await _poller_over(table, billed_logger).bill_if_finished( + response_id="resp_end_user", response=_terminal_response() + ) + + assert table.statuses == {"job-end-user": "completed"} + (billed,) = billed_logger.billed + assert billed.metadata["user_api_key_end_user_id"] == "customer-42" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("created_by", "billed_user_id"), + [("alice", "alice"), ("key:hashed-token", None)], + ids=["key_with_a_user", "key_with_no_user_or_team"], +) +async def test_the_billing_event_bills_the_creating_user_and_never_a_keys_owner_scope( + created_by: str, billed_user_id: str | None +): + table: Final = _InMemoryManagedObjectTable( + rows=( + _StoredResponse( + id="job-owner", + unified_object_id="resp_owner", + status="queued", + created_by=created_by, + api_key="hashed-token", + ), + ) + ) + billed_logger: Final = _RecordingBilledLogger() + + await _poller_over(table, billed_logger).bill_if_finished(response_id="resp_owner", response=_terminal_response()) + + (billed,) = billed_logger.billed + assert billed.metadata["user_api_key_user_id"] == billed_user_id + + +@pytest.mark.asyncio +async def test_settling_from_the_deployment_bills_a_finished_response_once_from_a_free_fetch(): + """A read the caller retargets, or a delete, bills from a fetch through the configured + deployment, since the client controls the body it would otherwise hand in.""" + table: Final = _InMemoryManagedObjectTable( + rows=(_StoredResponse(id="job-settle", unified_object_id="resp_settle", status="in_progress"),) + ) + billed_logger: Final = _RecordingBilledLogger() + poller: Final = _poller_over(table, billed_logger) + fetched: Final = _terminal_response() + + with patch("litellm.aget_responses", new_callable=AsyncMock, return_value=fetched) as fetch: + charge: Final = await poller.bill_from_deployment("resp_settle") + with patch("litellm.aget_responses", side_effect=AssertionError("a billed response is never fetched again")): + assert await poller.bill_from_deployment("resp_settle") is None + + assert charge == billed_logger.charge + assert [_is_billed(call.kwargs["litellm_metadata"]) for call in fetch.call_args_list] == [False] + (billed,) = billed_logger.billed + assert billed.response.usage == fetched.usage + assert table.statuses == {"job-settle": "completed"} + + +@pytest.mark.asyncio +async def test_settling_from_the_deployment_leaves_a_running_response_for_the_poll(): + table: Final = _InMemoryManagedObjectTable( + rows=(_StoredResponse(id="job-running", unified_object_id="resp_running", status="queued"),) + ) + billed_logger: Final = _RecordingBilledLogger() + + with patch("litellm.aget_responses", new_callable=AsyncMock, return_value=_terminal_response("in_progress")): + assert await _poller_over(table, billed_logger).bill_from_deployment("resp_running") is None + + assert billed_logger.billed == () + assert table.statuses == {"job-running": "queued"} + + +class _SpendCapture: + def __init__(self) -> None: + self.payloads: tuple[Mapping[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload = kwargs["standard_logging_object"] + assert isinstance(payload, Mapping) + self.payloads = (*self.payloads, payload) + + +def _spend_capture_logger(): + from litellm.integrations.custom_logger import CustomLogger + + class _CapturingLogger(_SpendCapture, CustomLogger): + pass + + return _CapturingLogger() + + +_PRICED_COST: Final = 100 * 0.001 + 50 * 0.002 + + +@dataclass(frozen=True) +class _PricedBilling: + table: _InMemoryManagedObjectTable + response_id: str + cost: float | None + + +_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} +) + + +async def _bill_through_a_priced_deployment( + read: ResponsesAPIResponse, + litellm_params: Mapping[str, object] = _PRICED_LITELLM_PARAMS, + model_info: Mapping[str, object] = MappingProxyType({}), +) -> _PricedBilling: + 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( + custom_llm_provider="openai", model_id="deployment-custom", response_id="resp_upstream" + ) + table: Final = _InMemoryManagedObjectTable( + rows=(_StoredResponse(id="job-priced", unified_object_id=response_id, status="queued", api_key="hash-alice"),) + ) + prisma_client: Final = MagicMock() + 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) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "priced-gpt", + "litellm_params": dict(litellm_params), + "model_info": {"id": "deployment-custom", **model_info}, + } + ] + ) + 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) + + +@pytest.mark.asyncio +async def test_the_billing_event_prices_the_fetched_usage_at_the_deployment_rates(monkeypatch: pytest.MonkeyPatch): + """Billing logs the body the read already fetched instead of reading it again, so the event must + carry the deployment's own pricing (100 input tokens at 0.001 plus 50 output tokens at 0.002 is + 0.2), and neither take the read's zero price nor change the response the reader gets back.""" + import litellm + + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + read: Final = _terminal_response() + read.hidden_params = {"response_cost": 0.0} + + billing: Final = await _bill_through_a_priced_deployment(read) + + (payload,) = capture.payloads + assert payload["response_cost"] == pytest.approx(_PRICED_COST) + 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 read.hidden_params == {"response_cost": 0.0} + assert billing.table.statuses == {"job-priced": "completed"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("base_model_in", ["litellm_params", "model_info"]) +async def test_the_billing_event_prices_the_usage_at_the_deployments_base_model( + monkeypatch: pytest.MonkeyPatch, base_model_in: str +): + """A deployment that names a base_model for cost tracking bills its background responses at that + model's price, whether the base_model sits in litellm_params or model_info.""" + import litellm + + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + monkeypatch.setitem( + litellm.model_cost, + "qa-priced-base-model", + { + "litellm_provider": "openai", + "mode": "responses", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + ) + base_model: Final = {"base_model": "qa-priced-base-model"} + + billing: Final = await _bill_through_a_priced_deployment( + _terminal_response(), + litellm_params={ + "model": "openai/gpt-5.6", + "api_key": "sk-unused", + **(base_model if base_model_in == "litellm_params" else {}), + }, + model_info=base_model if base_model_in == "model_info" else {}, + ) + + (payload,) = capture.payloads + assert payload["response_cost"] == pytest.approx(_PRICED_COST) + assert billing.cost == pytest.approx(_PRICED_COST) + + +@pytest.mark.asyncio +async def test_the_billing_event_reaches_sync_only_success_callbacks(monkeypatch: pytest.MonkeyPatch): + """Sync-only success callbacks (legacy string integrations, plain functions) get every async + call's event from a worker thread, so they must get the billed background response too.""" + import queue + + import litellm + + priced_costs: Final[queue.Queue[object]] = queue.Queue() + + def record_cost(kwargs: Mapping[str, object], *_: object) -> None: + priced_costs.put(kwargs.get("response_cost")) + + monkeypatch.setattr(litellm, "success_callback", [record_cost]) + monkeypatch.setattr(litellm, "_async_success_callback", []) + + await _bill_through_a_priced_deployment(_terminal_response()) + + assert priced_costs.get(timeout=10) == pytest.approx(_PRICED_COST) + + +@pytest.mark.asyncio +async def test_the_billing_event_keys_its_spend_row_apart_from_the_read_it_bills(monkeypatch: pytest.MonkeyPatch): + """The billing event logs the body the read returned, under the same response id, and spend logs + keep one row per request id. Its row must take its own request id, or the read's free row takes + the slot and the charge reaches the key with no spend row for an admin to see.""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + class _RequestIdCapture(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.request_ids: tuple[str, ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.request_ids = ( + *self.request_ids, + get_logging_payload(kwargs, response_obj, start_time, end_time)["request_id"], + ) + + capture: Final = _RequestIdCapture() + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + read: Final = _terminal_response() + + billing: Final = await _bill_through_a_priced_deployment(read) + + (billed_request_id,) = capture.request_ids + assert billed_request_id != billing.response_id + assert billed_request_id.startswith(billing.response_id) + diff --git a/tests/unit/proxy/db/test_model_usage_rollup.py b/tests/unit/proxy/db/test_model_usage_rollup.py index 0b7427d71f2..1e79d06370b 100644 --- a/tests/unit/proxy/db/test_model_usage_rollup.py +++ b/tests/unit/proxy/db/test_model_usage_rollup.py @@ -1,4 +1,5 @@ import asyncio +import json from datetime import datetime, timezone from typing import Any from unittest.mock import MagicMock @@ -6,6 +7,7 @@ from unittest.mock import MagicMock import httpx import pytest +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.proxy.db import rollup_lock_timeout from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.model_usage_rollup import ( @@ -16,6 +18,7 @@ from litellm.proxy.db.model_usage_rollup import ( model_usage_task_type, ) from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN class _FakeBatcher: @@ -155,6 +158,27 @@ async def test_flush_aggregates_each_rollup_row_into_one_upsert() -> None: assert upserts["claude"]["create"]["request_count"] == 1 +@pytest.mark.asyncio +async def test_a_background_response_billing_row_adds_spend_but_no_request() -> None: + """A background response is billed by a later event, not by the create the client sent, so + model insights counts its spend and tokens once and its request once, from the create.""" + batcher = _FakeBatcher() + prisma = _prisma(MagicMock(return_value=batcher)) + create = build_model_usage_transaction(_payload(spend=0.0)) + billing = build_model_usage_transaction( + _payload(spend=0.5, metadata=json.dumps({INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN})) + ) + assert create is not None and billing is not None + + await flush_model_usage_transactions(prisma_client=prisma, transactions=[create, billing]) + + (upsert,) = batcher.litellm_dailymodelusage.upsert.call_args_list + created = upsert.kwargs["data"]["create"] + assert created["spend"] == 0.5 + assert created["prompt_tokens"] == 20 + assert (created["request_count"], created["successful_requests"], created["failed_requests"]) == (1, 1, 0) + + @pytest.mark.asyncio async def test_flush_with_no_transactions_touches_nothing() -> None: prisma = _prisma(MagicMock()) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index c1ff8205a04..1af42947da7 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2,9 +2,13 @@ Test for response_api_endpoints/endpoints.py """ +import asyncio +import contextlib +import dataclasses +import json import unittest from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import Any, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -15,7 +19,7 @@ from httpx import Response import litellm from litellm.proxy.proxy_server import app -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from tests._master_key import MASTER_KEY @@ -300,6 +304,1060 @@ async def test_responses_api_background_polling_accepts_input_from_prompt_templa assert request_data["input"] == "hello from prompt" +@pytest.mark.asyncio +async def test_a_cache_polled_background_stream_is_priced_by_its_own_stream(): + """Cache polling runs the stream without the settlement wrapper, so a stream the managed files hook + marked as billed later must still be priced when it finishes, or the generation is never billed.""" + from fastapi import Response as FastAPIResponse + + from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY + from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import responses_api + + marked_by_the_hook: Final = { + "model": "gpt-5.6", + "input": "hello", + "stream": True, + "background": True, + "extra_body": {"background": True}, + "litellm_metadata": {MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True, "user_api_key_alias": "creator"}, + } + processor: Final = MagicMock() + processor.common_processing_pre_call_logic = AsyncMock(return_value=(marked_by_the_hook, MagicMock())) + body: Final = json.dumps( + {"model": "gpt-5.6", "input": "hello", "stream": True, "background": True, "extra_body": {"background": True}} + ).encode() + + with ( + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: polling decision is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", + return_value=True, + ), + patch( # test-quality-ok: background task is imported inside the endpoint + "litellm.proxy.response_polling.background_streaming.background_streaming_task", + new_callable=AsyncMock, + ) as polled_stream, + patch( # test-quality-ok: polling handler is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.ResponsePollingHandler.create_initial_state", + new_callable=AsyncMock, + ), + ): + await responses_api( + request=_request_for("POST", "/v1/responses", body), + fastapi_response=FastAPIResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-creator"), + ) + await asyncio.sleep(0) + + polled_metadata: Final = polled_stream.call_args.kwargs["data"]["litellm_metadata"] + finished: Final = _background_response("resp_cache_polled", "completed", 100) + assert is_unbilled_non_inference_call("aresponses", marked_by_the_hook["litellm_metadata"], finished) + assert not is_unbilled_non_inference_call("aresponses", polled_metadata, finished) + assert polled_metadata["user_api_key_alias"] == "creator" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["queued", "completed"]) +async def test_responses_api_background_create_persists_the_creating_keys_attribution(status: str): + """The cost poll bills a background job off its managed object row, so the create must persist + the key hash and tags the poll bills against, the way the batch create does, plus the client's + User-Agent tags that a spend-by-client view groups the charge under. A create the provider + already finished is stored too, so a later read finds it billed instead of billing it again.""" + from fastapi import Response as FastAPIResponse + from starlette.requests import Request + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import responses_api + + queued = ResponsesAPIResponse( + id="resp_background", object="response", status=status, created_at=1, output=[], usage=None + ) + queued._hidden_params = {"model_id": "deployment-1"} + processor = MagicMock() + processor.data = { + "model": "gpt-4o", + "input": "hello", + "background": True, + "litellm_metadata": {"tags": ["env:prod"]}, + "proxy_server_request": {"headers": {"user-agent": "codex_cli_rs/0.44.0"}}, + } + processor.base_process_llm_request = AsyncMock(return_value=queued) + managed_files = MagicMock() + managed_files.store_unified_object_id = AsyncMock() + managed_files.track_untracked_background_response = AsyncMock(return_value=False) + proxy_logging = MagicMock() + proxy_logging.get_proxy_hook = MagicMock(return_value=managed_files) + + async def receive(): + return { + "type": "http.request", + "body": b'{"model":"gpt-4o","input":"hello","background":true}', + "more_body": False, + } + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + with ( + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: polling decision is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", + return_value=False, + ), + patch( # test-quality-ok: the endpoint reads the proxy singletons at call time + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging + ), + patch( # test-quality-ok: the endpoint reads the proxy singletons at call time + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), + ): + result = await responses_api( + request=request, + fastapi_response=FastAPIResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert result is queued + store_kwargs = managed_files.store_unified_object_id.await_args.kwargs + assert store_kwargs["file_object"].status == status + assert store_kwargs["file_purpose"] == "response" + assert store_kwargs["persist_attribution"] is True + assert store_kwargs["request_tags"] == ( + "env:prod", + "User-Agent: codex_cli_rs", + "User-Agent: codex_cli_rs/0.44.0", + ) + + +class _SpendEvents: + def __init__(self) -> None: + self.payloads: tuple[Mapping[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload = kwargs["standard_logging_object"] + assert isinstance(payload, Mapping) + self.payloads = (*self.payloads, payload) + + +async def _drain_spend_events() -> None: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + +def _billed_costs(payloads: tuple[Mapping[str, object], ...]) -> list[str]: + return [str(payload["response_cost"]) for payload in payloads] + + +def _advertised_cost_split(headers: Mapping[str, str]) -> list[tuple[float, float]]: + if "x-litellm-response-cost-original" not in headers: + return [] + return [ + ( + float(headers["x-litellm-response-cost-original"]), + float(headers["x-litellm-response-cost-input"]) + float(headers["x-litellm-response-cost-output"]), + ) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "background, stored, status, expected_billed_events", + [(True, True, "completed", 1), (True, True, "in_progress", 0), (True, False, "completed", 0), (False, True, "completed", 0)], +) +async def test_reading_a_finished_background_response_bills_it( + monkeypatch: pytest.MonkeyPatch, background: bool, stored: bool, status: str, expected_billed_events: int +): + """The read that first sees a background response finished bills it from the body it fetched, + since the provider may drop a store=false response before the hourly cost poll reaches it. + Billing never reads the provider a second time, other reads stay free, the charge stays on the + end user the create named, and only the read that billed reports the cost it billed, split + into the components the billed event priced.""" + from types import SimpleNamespace + + from fastapi import Response as FastAPIResponse + + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import _bill_finished_background_response + + class _CapturingLogger(_SpendEvents, CustomLogger): + pass + + spend_events = _CapturingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [spend_events]) + read = ResponsesAPIResponse( + id="resp_read", + object="response", + status=status, + created_at=1, + model="gpt-5.6", + output=[], + usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150), + background=background, + ) + stored_row = SimpleNamespace( + id="job-read", + unified_object_id="resp_read", + status="queued", + file_object={"model": "gpt-5.6"}, + created_by=None, + api_key=None, + team_id=None, + org_id=None, + request_tags=None, + end_user_id="customer-42", + ) + fastapi_response = FastAPIResponse() + prisma_client = MagicMock() + prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row if stored else None) + prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")), + ): + await _bill_finished_background_response( + response=read, + response_id="resp_read", + user_api_key_dict=UserAPIKeyAuth(end_user_id="another-reader"), + fastapi_response=fastapi_response, + read_settles_billing=True, + ) + await _drain_spend_events() + + assert len(spend_events.payloads) == expected_billed_events + assert all(payload["metadata"]["user_api_key_end_user_id"] == "customer-42" for payload in spend_events.payloads) + assert fastapi_response.headers.getlist("x-litellm-response-cost") == _billed_costs(spend_events.payloads) + assert _advertised_cost_split(fastapi_response.headers) == [ + (pytest.approx(float(cost)), pytest.approx(float(cost))) for cost in _billed_costs(spend_events.payloads) + ] + + +@dataclasses.dataclass(frozen=True) +class _ManagedObjectRow: + id: str + unified_object_id: str + model_object_id: str + status: str + file_object: Mapping[str, object] + file_purpose: str = "response" + created_by: str | None = None + api_key: str | None = None + team_id: str | None = None + org_id: str | None = None + request_tags: object = None + end_user_id: str | None = None + + +class _UniqueManagedObjectTable: + """Rejects a second row for a response id and applies status claims atomically, the way the + Postgres table does, and yields on every query so concurrent reads interleave like two pods.""" + + def __init__(self, rows: tuple[_ManagedObjectRow, ...] = ()) -> None: + self.rows = rows + + def _matches(self, row: _ManagedObjectRow, where: Mapping[str, object]) -> bool: + def holds(field: str, condition: object) -> bool: + if isinstance(condition, list): + return any(self._matches(row, alternative) for alternative in condition) + if isinstance(condition, dict): + return getattr(row, field) in condition["in"] + return getattr(row, field) == condition + + return all(holds(field, condition) for field, condition in where.items()) + + async def find_first(self, where: Mapping[str, object]) -> _ManagedObjectRow | None: + await asyncio.sleep(0) + return next((row for row in self.rows if self._matches(row, where)), None) + + async def create(self, data: Mapping[str, object]) -> _ManagedObjectRow: + from prisma.errors import UniqueViolationError + + await asyncio.sleep(0) + if any( + row.unified_object_id == data["unified_object_id"] or row.model_object_id == data["model_object_id"] + for row in self.rows + ): + raise UniqueViolationError({"user_facing_error": {"error_code": "P2002"}}) + row = _ManagedObjectRow( + id=f"job-{len(self.rows)}", + unified_object_id=str(data["unified_object_id"]), + model_object_id=str(data["model_object_id"]), + status=str(data["status"]), + file_object=json.loads(str(data["file_object"])), + created_by=cast("str | None", data.get("created_by")), + api_key=cast("str | None", data.get("api_key")), + team_id=cast("str | None", data.get("team_id")), + org_id=cast("str | None", data.get("org_id")), + end_user_id=cast("str | None", data.get("end_user_id")), + request_tags=getattr(data.get("request_tags"), "data", None), + ) + self.rows = (*self.rows, row) + return row + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: + await asyncio.sleep(0) + matched = frozenset(row.id for row in self.rows if self._matches(row, where)) + self.rows = tuple( + dataclasses.replace(row, status=str(data["status"])) if row.id in matched else row for row in self.rows + ) + return len(matched) + + +async def _read_twice_on_two_pods( + monkeypatch: pytest.MonkeyPatch, table: _UniqueManagedObjectTable, response_id: str, reader: object +) -> tuple[Mapping[str, object], ...]: + from fastapi import Response as FastAPIResponse + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + from litellm.caching.caching import DualCache + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.response_api_endpoints.endpoints import _bill_finished_background_response + + class _CapturingLogger(_SpendEvents, CustomLogger): + pass + + spend_events: Final = _CapturingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [spend_events]) + prisma_client: Final = MagicMock() + 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) + proxy_logging: Final = MagicMock() + proxy_logging.get_proxy_hook = MagicMock( + return_value=_PROXY_LiteLLMManagedFiles(internal_usage_cache=DualCache(), prisma_client=prisma_client) + ) + finished: Final = ResponsesAPIResponse( + id=response_id, + object="response", + status="completed", + created_at=1, + model="gpt-5.6", + output=[], + usage=None, + background=True, + ) + + router: Final = _configured_deployment_router(finished) + router.aget_responses = AsyncMock(side_effect=AssertionError("billing must not read the provider again")) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")), + ): + await asyncio.gather( + *( + _bill_finished_background_response( + response=finished, + response_id=response_id, + user_api_key_dict=reader, + fastapi_response=FastAPIResponse(), + read_settles_billing=True, + ) + for _ in range(2) + ) + ) + await _drain_spend_events() + return spend_events.payloads + + +@pytest.mark.asyncio +async def test_a_finished_response_the_create_could_not_store_is_billed_once_under_the_reader( + monkeypatch: pytest.MonkeyPatch, +): + """A create that could not write its managed row (a database outage) leaves the reads marked free + and the poll blind to the job, so the first read that sees it finished stores the row under the + reading key and bills it, once even when two pods read it at the same moment.""" + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.responses.utils import ResponsesAPIRequestUtils + + response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-untracked", response_id="resp_untracked" + ) + table: Final = _UniqueManagedObjectTable() + + payloads: Final = await _read_twice_on_two_pods( + monkeypatch, + table, + response_id, + UserAPIKeyAuth(api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42"), + ) + + (billed,) = payloads + assert billed["metadata"]["user_api_key_hash"] == hash_token("sk-reader") + assert billed["metadata"]["user_api_key_team_id"] == "team-beta" + assert billed["metadata"]["user_api_key_org_id"] == "org-gamma" + assert billed["metadata"]["user_api_key_end_user_id"] == "customer-42" + (row,) = table.rows + assert (row.unified_object_id, row.status) == (response_id, "completed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_issued, stored_status", [(True, "completed"), (False, None)]) +async def test_a_finished_read_never_bills_a_response_already_billed_or_billed_by_the_read_itself( + monkeypatch: pytest.MonkeyPatch, proxy_issued: bool, stored_status: str | None +): + """A row the create stored as finished means the create billed it, and an id the proxy did not + issue is priced by the read itself, so neither read stores a row or bills the response again.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.responses.utils import ResponsesAPIRequestUtils + + response_id: Final = ( + ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-billed", response_id="resp_billed" + ) + if proxy_issued + else "resp_not_issued_by_the_proxy" + ) + stored: Final = ( + ( + _ManagedObjectRow( + id="job-billed", + unified_object_id=response_id, + model_object_id="resp_billed", + status=stored_status, + file_object={"model": "gpt-5.6"}, + ), + ) + if stored_status is not None + else () + ) + table: Final = _UniqueManagedObjectTable(stored) + + payloads: Final = await _read_twice_on_two_pods(monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader")) + + assert payloads == () + assert table.rows == stored + + +def _configured_deployment_router(fetched: ResponsesAPIResponse) -> MagicMock: + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router: Final = MagicMock() + router.get_deployment = MagicMock( + return_value=Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id="deployment-read"), + ) + ) + router.aget_responses = AsyncMock(return_value=fetched) + return router + + +def _background_response(response_id: str, status: str, input_tokens: int) -> ResponsesAPIResponse: + return ResponsesAPIResponse( + id=response_id, + object="response", + status=status, + created_at=1, + model="gpt-5.6", + output=[], + usage=ResponseAPIUsage(input_tokens=input_tokens, output_tokens=50, total_tokens=input_tokens + 50), + background=True, + ) + + +@dataclasses.dataclass(frozen=True) +class _StoredBackgroundResponse: + response_id: str + table: _UniqueManagedObjectTable + prisma_client: MagicMock + proxy_logging: MagicMock + + +def _one_pending_background_response() -> _StoredBackgroundResponse: + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + from litellm.caching.caching import DualCache + from litellm.responses.utils import ResponsesAPIRequestUtils + + response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-read", response_id="resp_read" + ) + table: Final = _UniqueManagedObjectTable( + ( + _ManagedObjectRow( + id="job-read", + unified_object_id=response_id, + model_object_id="resp_read", + status="queued", + file_object={"model": "gpt-5.6"}, + ), + ) + ) + prisma_client: Final = MagicMock() + 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) + proxy_logging: Final = MagicMock() + proxy_logging.get_proxy_hook = MagicMock( + return_value=_PROXY_LiteLLMManagedFiles(internal_usage_cache=DualCache(), prisma_client=prisma_client) + ) + return _StoredBackgroundResponse( + response_id=response_id, table=table, prisma_client=prisma_client, proxy_logging=proxy_logging + ) + + +def _request_for(method: str, path: str, body: bytes): + from starlette.requests import Request + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + {"type": "http", "method": method, "path": path, "headers": [(b"content-type", b"application/json")]}, + receive, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "endpoint_name, finished_status", + [("get_response", "completed"), ("cancel_response", "cancelled")], +) +@pytest.mark.parametrize( + "read_body, billed_input_tokens, deployment_fetches", + [(b"", 1000, 0), (b'{"api_base": "https://caller.example/v1"}', 100, 1)], +) +async def test_a_finished_read_or_cancel_bills_usage_only_the_configured_deployment_reported( + monkeypatch: pytest.MonkeyPatch, + endpoint_name: str, + finished_status: str, + read_body: bytes, + billed_input_tokens: int, + deployment_fetches: int, +): + """A read or cancel the caller sends to its own endpoint returns whatever usage that endpoint + claims, so it bills from a fetch through the configured deployment instead, and either way the + response reports the cost it billed.""" + from fastapi import Response as FastAPIResponse + + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints import endpoints + + class _CapturingLogger(_SpendEvents, CustomLogger): + pass + + spend_events: Final = _CapturingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [spend_events]) + stored: Final = _one_pending_background_response() + router: Final = _configured_deployment_router(_background_response(stored.response_id, finished_status, 100)) + processor: Final = MagicMock() + processor.base_process_llm_request = AsyncMock( + return_value=_background_response(stored.response_id, finished_status, 1000) + ) + fastapi_response: Final = FastAPIResponse() + + with ( + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.aget_responses", side_effect=AssertionError("the configured deployment answers billing reads")), + ): + await getattr(endpoints, endpoint_name)( + response_id=stored.response_id, + request=_request_for("GET", f"/v1/responses/{stored.response_id}", read_body), + fastapi_response=fastapi_response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-reader"), + ) + await _drain_spend_events() + + (billed,) = spend_events.payloads + assert billed["prompt_tokens"] == billed_input_tokens + assert router.aget_responses.await_count == deployment_fetches + assert fastapi_response.headers.getlist("x-litellm-response-cost") == _billed_costs(spend_events.payloads) + (row,) = stored.table.rows + assert row.status == "completed" + + +@pytest.mark.asyncio +async def test_deleting_a_finished_unbilled_background_response_bills_it_before_the_provider_forgets_it( + monkeypatch: pytest.MonkeyPatch, +): + """A delete leaves the cost poll nothing to fetch, so a finished response no read has billed yet + is billed from the configured deployment before the delete reaches the provider.""" + from fastapi import Response as FastAPIResponse + + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import delete_response + from litellm.types.responses.main import DeleteResponseResult + + class _CapturingLogger(_SpendEvents, CustomLogger): + pass + + spend_events: Final = _CapturingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [spend_events]) + stored: Final = _one_pending_background_response() + deleted: Final[list[str]] = [] + + async def provider_fetch(response_id: str, litellm_metadata: Mapping[str, object]) -> ResponsesAPIResponse: + if deleted: + raise litellm.NotFoundError(message="response not found", model="gpt-5.6", llm_provider="openai") + return _background_response(response_id, "completed", 100) + + async def provider_delete(fastapi_response: FastAPIResponse, **_: object) -> DeleteResponseResult: + deleted.append(stored.response_id) + fastapi_response.headers["x-litellm-response-cost"] = "0.0" + return DeleteResponseResult(id=stored.response_id, object="response", deleted=True) + + router: Final = _configured_deployment_router(_background_response(stored.response_id, "completed", 100)) + router.aget_responses = AsyncMock(side_effect=provider_fetch) + processor: Final = MagicMock() + processor.base_process_llm_request = AsyncMock(side_effect=provider_delete) + fastapi_response: Final = FastAPIResponse() + + with ( + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + ): + result: Final = await delete_response( + response_id=stored.response_id, + request=_request_for("DELETE", f"/v1/responses/{stored.response_id}", b""), + fastapi_response=fastapi_response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-reader"), + ) + await _drain_spend_events() + + assert result.deleted is True + (billed,) = spend_events.payloads + assert billed["prompt_tokens"] == 100 + assert fastapi_response.headers.getlist("x-litellm-response-cost") == _billed_costs(spend_events.payloads) + (row,) = stored.table.rows + assert row.status == "completed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("caller_role, billed_count", [("proxy_admin", 1), ("internal_user", 0)]) +async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(caller_role: str, billed_count: int): + """A delete settles the bill only for a caller the response id check lets through, so a key that + cannot address the response gets its 403 without the proxy fetching or billing it first.""" + from fastapi import HTTPException + from fastapi import Response as FastAPIResponse + + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity + from litellm.proxy.response_api_endpoints.endpoints import delete_response + from litellm.types.responses.main import DeleteResponseResult + + class _CapturingLogger(_SpendEvents, CustomLogger): + pass + + spend_events: Final = _CapturingLogger() + stored: Final = _one_pending_background_response() + managed_files: Final = stored.proxy_logging.get_proxy_hook.return_value + security: Final = ResponsesIDSecurity(general_settings_reader=lambda: {}, signing_key_reader=lambda: "salt") + stored.proxy_logging.get_proxy_hook = MagicMock( + side_effect=lambda name: security if name == "responses_id_security" else managed_files + ) + router: Final = _configured_deployment_router(_background_response(stored.response_id, "completed", 100)) + caller: Final = UserAPIKeyAuth(api_key="sk-caller", user_role=caller_role) + + async def pipeline(**_: object) -> DeleteResponseResult: + await security.async_pre_call_hook(caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses") + return DeleteResponseResult(id=stored.response_id, object="response", deleted=True) + + processor: Final = MagicMock() + processor.base_process_llm_request = AsyncMock(side_effect=pipeline) + processor.handle_llm_api_exception = AsyncMock(side_effect=lambda e, **_: e) + + with ( + patch.object(litellm, "_async_success_callback", [spend_events]), + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + ): + with contextlib.suppress(HTTPException): + await delete_response( + response_id=stored.response_id, + request=_request_for("DELETE", f"/v1/responses/{stored.response_id}", b""), + fastapi_response=FastAPIResponse(), + user_api_key_dict=caller, + ) + await _drain_spend_events() + + assert len(spend_events.payloads) == billed_count + assert router.aget_responses.await_count == billed_count + + +def _proxy_issued_id(provider_response_id: str) -> str: + from litellm.responses.utils import ResponsesAPIRequestUtils + + return ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-read", response_id=provider_response_id + ) + + +def _managed_files_over(table: _UniqueManagedObjectTable) -> tuple[MagicMock, MagicMock]: + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + from litellm.caching.caching import DualCache + + prisma_client: Final = MagicMock() + 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) + proxy_logging: Final = MagicMock() + proxy_logging.get_proxy_hook = MagicMock( + return_value=_PROXY_LiteLLMManagedFiles(internal_usage_cache=DualCache(), prisma_client=prisma_client) + ) + return prisma_client, proxy_logging + + +def _capture_spend_events(monkeypatch: pytest.MonkeyPatch) -> _SpendEvents: + from litellm.integrations.custom_logger import CustomLogger + + class _CapturingLogger(_SpendEvents, CustomLogger): + pass + + spend_events: Final = _CapturingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [spend_events]) + return spend_events + + +@contextlib.contextmanager +def _proxy_over(prisma_client: MagicMock, proxy_logging: MagicMock, router: MagicMock): + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + ): + yield + + +@pytest.mark.asyncio +@pytest.mark.parametrize("completed_after_last_event", [False, True], ids=["before_last_event", "after_last_event"]) +async def test_a_streamed_background_create_bills_its_finished_response_once_under_the_creator( + monkeypatch: pytest.MonkeyPatch, completed_after_last_event: bool +): + """A tracked streamed background create is free on its own and bills the finished response once + through its managed row, under the creating key and tags, before the client sees it finish + whenever the stream records the finished response in time, and no later read bills it again.""" + from fastapi import Response as FastAPIResponse + + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.response_api_endpoints.endpoints import ( + _bill_finished_background_response, + _streamed_background_data_generator, + ) + + spend_events: Final = _capture_spend_events(monkeypatch) + table: Final = _UniqueManagedObjectTable() + prisma_client, proxy_logging = _managed_files_over(table) + finished: Final = _background_response(_proxy_issued_id("resp_streamed"), "completed", 100) + stream: Final = MagicMock(completed_response=None) + + async def provider_events(**_: object): + yield "event: response.created\n\n" + if not completed_after_last_event: + stream.completed_response = MagicMock(response=finished) + yield "event: response.completed\n\n" + stream.completed_response = MagicMock(response=finished) + + billed_before_each_event: Final[list[int]] = [] + with _proxy_over(prisma_client, proxy_logging, _configured_deployment_router(finished)): + async for _ in _streamed_background_data_generator( + provider_events, + proxy_logging, + stream, + UserAPIKeyAuth(api_key="sk-creator"), + {"litellm_metadata": {"tags": ["nightly"]}}, + ): + billed_before_each_event.append(len(table.rows)) + await _bill_finished_background_response( + response=finished, + response_id=finished.id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-reader"), + fastapi_response=FastAPIResponse(), + read_settles_billing=True, + ) + await _drain_spend_events() + + assert billed_before_each_event == ([0, 0] if completed_after_last_event else [0, 1]) + (billed,) = spend_events.payloads + assert billed["metadata"]["user_api_key_hash"] == hash_token("sk-creator") + assert billed["request_tags"] == ["nightly"] + + +@pytest.mark.asyncio +async def test_a_read_that_bills_a_response_before_its_stream_ends_leaves_the_stream_nothing_to_bill( + monkeypatch: pytest.MonkeyPatch, +): + """A read can see a streamed background response finished before a slow client lets the stream + reach its last event, and then the stream bills nothing on top of the read.""" + from fastapi import Response as FastAPIResponse + + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.response_api_endpoints.endpoints import ( + _bill_finished_background_response, + _streamed_background_data_generator, + ) + + spend_events: Final = _capture_spend_events(monkeypatch) + table: Final = _UniqueManagedObjectTable() + prisma_client, proxy_logging = _managed_files_over(table) + finished: Final = _background_response(_proxy_issued_id("resp_raced"), "completed", 100) + stream: Final = MagicMock(completed_response=None) + + async def provider_events(**_: object): + yield "event: response.created\n\n" + await _bill_finished_background_response( + response=finished, + response_id=finished.id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-reader"), + fastapi_response=FastAPIResponse(), + read_settles_billing=True, + ) + stream.completed_response = MagicMock(response=finished) + yield "event: response.completed\n\n" + + with _proxy_over(prisma_client, proxy_logging, _configured_deployment_router(finished)): + async for _ in _streamed_background_data_generator( + provider_events, proxy_logging, stream, UserAPIKeyAuth(api_key="sk-creator"), {} + ): + pass + await _drain_spend_events() + + (billed,) = spend_events.payloads + assert billed["metadata"]["user_api_key_hash"] == hash_token("sk-reader") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "background_fields", + [{"background": True}, {"extra_body": {"background": True}}], + ids=["top_level", "extra_body"], +) +async def test_a_streamed_background_create_is_billed_wherever_it_sets_background( + monkeypatch: pytest.MonkeyPatch, background_fields: Mapping[str, object] +): + """A streamed create that sets background inside extra_body runs as a background job all the + same, so its stream bills the finished response once.""" + from fastapi import Response as FastAPIResponse + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.response_api_endpoints.endpoints import responses_api + + spend_events: Final = _capture_spend_events(monkeypatch) + table: Final = _UniqueManagedObjectTable() + prisma_client, proxy_logging = _managed_files_over(table) + finished: Final = _background_response(_proxy_issued_id("resp_streamed_create"), "completed", 100) + stream: Final = MagicMock(completed_response=None) + + async def provider_events(**_: object): + stream.completed_response = MagicMock(response=finished) + yield "event: response.completed\n\n" + + async def stream_to_client(select_data_generator, **_: object) -> MagicMock: + async for _ in select_data_generator(response=stream, user_api_key_dict=UserAPIKeyAuth(), request_data={}): + pass + return MagicMock() + + processor: Final = MagicMock() + processor.base_process_llm_request = AsyncMock(side_effect=stream_to_client) + processor_class: Final = MagicMock( + return_value=processor, + extract_completed_responses_response=ProxyBaseLLMRequestProcessing.extract_completed_responses_response, + ) + body: Final = json.dumps({"model": "gpt-5.6", "input": "hello", "stream": True, **background_fields}).encode() + + with ( + _proxy_over(prisma_client, proxy_logging, _configured_deployment_router(finished)), + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", processor_class + ), + patch( # test-quality-ok: polling decision is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", + return_value=False, + ), + patch("litellm.proxy.proxy_server.select_data_generator", provider_events), # test-quality-ok: the endpoint reads the proxy singletons at call time + ): + await responses_api( + request=_request_for("POST", "/v1/responses", body), + fastapi_response=FastAPIResponse(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-creator"), + ) + await _drain_spend_events() + + assert len(spend_events.payloads) == 1 + + +@pytest.mark.asyncio +async def test_a_retargeted_read_of_a_response_nobody_tracked_bills_it_once_from_the_deployment( + monkeypatch: pytest.MonkeyPatch, +): + """A read that carries its own api_base returns whatever usage that endpoint claims, so a finished + response with no managed row is tracked and billed from a fetch through the configured deployment, + once even when two pods read it at the same moment.""" + from fastapi import Response as FastAPIResponse + + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.response_api_endpoints.endpoints import _bill_finished_background_response + + spend_events: Final = _capture_spend_events(monkeypatch) + table: Final = _UniqueManagedObjectTable() + prisma_client, proxy_logging = _managed_files_over(table) + response_id: Final = _proxy_issued_id("resp_untracked_override") + router: Final = _configured_deployment_router(_background_response(response_id, "completed", 100)) + + with _proxy_over(prisma_client, proxy_logging, router): + await asyncio.gather( + *( + _bill_finished_background_response( + response=_background_response(response_id, "completed", 1000), + response_id=response_id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-reader"), + fastapi_response=FastAPIResponse(), + read_settles_billing=False, + ) + for _ in range(2) + ) + ) + await _drain_spend_events() + + (billed,) = spend_events.payloads + assert billed["prompt_tokens"] == 100 + assert billed["metadata"]["user_api_key_hash"] == hash_token("sk-reader") + + +@pytest.mark.asyncio +async def test_a_bill_that_fails_to_settle_is_retried_in_the_background(monkeypatch: pytest.MonkeyPatch): + """A read is free once marked, so a bill that fails while the read settles it (a database blip) + is retried in the background instead of waiting on a cost poll that may be turned off.""" + from litellm.proxy.response_api_endpoints import endpoints + + spend_events: Final = _capture_spend_events(monkeypatch) + stored: Final = _one_pending_background_response() + finished: Final = _background_response(stored.response_id, "completed", 100) + attempts: Final[list[int]] = [] + + async def bill_after_one_failure(cost_check: endpoints._BackgroundCostCheck): + attempts.append(len(attempts)) + if len(attempts) == 1: + raise RuntimeError("database unavailable") + return await cost_check.bill_if_finished(stored.response_id, finished) + + with _proxy_over(stored.prisma_client, stored.proxy_logging, _configured_deployment_router(finished)): + assert await endpoints._settle_background_billing(bill_after_one_failure, retry_delays=(0.0, 0.0)) is None + await asyncio.gather(*endpoints._billing_retries) + await _drain_spend_events() + + assert len(attempts) == 2 + assert len(spend_events.payloads) == 1 + + +@pytest.mark.asyncio +async def test_background_create_the_provider_confirms_is_stored_under_its_provider_id( + monkeypatch: pytest.MonkeyPatch, +): + """A create that sets background through extra_body reaches the provider as a background job + all the same, so the proxy stores it off the response, keyed by the id every re-encryption of + the client's id decrypts to.""" + from fastapi import Response as FastAPIResponse + from starlette.requests import Request + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity + from litellm.proxy.response_api_endpoints.endpoints import responses_api + + salt_key = "background-create-unit-test-salt" + monkeypatch.setenv("LITELLM_SALT_KEY", salt_key) + queued = ResponsesIDSecurity( + general_settings_reader=lambda: {}, signing_key_reader=lambda: salt_key + )._encrypt_response_id( + ResponsesAPIResponse( + id="resp_provider_job", object="response", status="queued", created_at=1, output=[], background=True + ), + UserAPIKeyAuth(user_id="alice"), + ) + queued._hidden_params = {"model_id": "deployment-1"} + processor = MagicMock() + processor.data = {"model": "gpt-4o", "input": "hello", "extra_body": {"background": True}} + processor.base_process_llm_request = AsyncMock(return_value=queued) + managed_files = MagicMock() + managed_files.store_unified_object_id = AsyncMock() + managed_files.track_untracked_background_response = AsyncMock(return_value=False) + proxy_logging = MagicMock() + proxy_logging.get_proxy_hook = MagicMock(return_value=managed_files) + + async def receive(): + return { + "type": "http.request", + "body": b'{"model":"gpt-4o","input":"hello","extra_body":{"background":true}}', + "more_body": False, + } + + request = Request( + {"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"content-type", b"application/json")]}, + receive, + ) + + with ( + patch( # test-quality-ok: endpoint constructs the processor directly + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: polling decision is imported inside the endpoint + "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", + return_value=False, + ), + patch( # test-quality-ok: the endpoint reads the proxy singletons at call time + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging + ), + patch( # test-quality-ok: the endpoint reads the proxy singletons at call time + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), + ): + await responses_api( + request=request, fastapi_response=FastAPIResponse(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test") + ) + + store_kwargs = managed_files.store_unified_object_id.await_args.kwargs + assert store_kwargs["unified_object_id"] == queued.id + assert store_kwargs["model_object_id"] == "resp_provider_job" + + class TestResponsesAPIEndpoints(unittest.TestCase): @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.llm_router") @@ -479,9 +1537,6 @@ class TestResponsesAPIEndpoints(unittest.TestCase): assert response_cost_value == pytest.approx(0.0005, abs=1e-10) -import json - - class TestManagedResponsesWSFirstMessage: @pytest.mark.asyncio async def test_first_message_processed_before_loop(self): diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 5eb738a2984..f5f690c4bfd 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -15,6 +15,7 @@ import litellm import litellm.constants as litellm_constants import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -48,6 +49,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( from litellm.proxy.utils import hash_token from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, StandardLoggingHiddenParams, StandardLoggingMetadata, StandardLoggingModelInformation, @@ -4400,7 +4402,10 @@ def test_user_traffic_carries_no_internal_call_origin(): def _spend_log_for_call_type( - call_type: str, internal_call_origin: str | None = None, background: bool | None = None + call_type: str, + internal_call_origin: str | None = None, + background: bool | None = None, + managed_read: bool = False, ) -> dict: from litellm.types.llms.openai import ResponsesAPIResponse @@ -4415,6 +4420,7 @@ def _spend_log_for_call_type( "metadata": { "user_api_key": "test-key", "internal_call_origin": internal_call_origin, + **({"_managed_background_response": True} if managed_read else {}), } }, }, @@ -4460,6 +4466,15 @@ def test_spend_log_for_background_response_retrieval_counts_tokens(): assert payload["total_tokens"] == 6000 +def test_spend_log_for_managed_background_response_read_counts_nothing(): + """A proxy that tracks a background response bills it once through the cost poll, so a client + read the managed files hook marked must not count the job's tokens a second time.""" + payload = _spend_log_for_call_type("aget_responses", background=True, managed_read=True) + + assert payload["total_tokens"] == 0 + assert payload["spend"] == 0.0 + + def test_spend_log_for_foreground_response_retrieval_still_counts_nothing(): """Guards the test above against a blanket exemption: an explicit background=false read was already billed by its create and must stay at zero.""" @@ -4573,6 +4588,34 @@ def test_get_spend_logs_id_separates_distinct_batches_whose_bodies_were_both_red assert ids == ["batch_first_batch_cost", "batch_second_batch_cost"] +def test_a_background_response_billing_event_writes_a_row_apart_from_the_read_that_fetched_it(): + """The one billing event of a background response logs the body a read just fetched, so both + carry the same response id. Spend logs keep one row per request id, so the billing event must + key its row apart from the read's, or the charge reaches the key and no row ever shows it.""" + response: Final = {"id": "resp_shared", "object": "response", "status": "completed", "background": True} + read_metadata: Final = {"user_api_key": "test-key", MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True} + billing_metadata: Final = { + "user_api_key": "test-key", + INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + } + request_ids: Final = tuple( + get_logging_payload( + kwargs={ + "call_type": "aget_responses", + "model": "gpt-5.6", + "litellm_call_id": call_id, + "litellm_params": {"metadata": metadata}, + }, + response_obj=dict(response), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + )["request_id"] + for call_id, metadata in (("call-read", read_metadata), ("call-bill", billing_metadata)) + ) + + assert request_ids == ("resp_shared", "resp_shared_background_cost") + + def test_get_spend_logs_id_prefers_the_response_id_over_the_standard_logging_id(): """An unredacted response keeps deciding its own row key, so cache-hit ids and every other call type behave exactly as they did before.""" diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 809ade05613..96c11be2650 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1189,6 +1189,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "litellm_gateway_injected_cache": "forged-deployment-id", "_session_deployment_affinity_ttl": 999999, "internal_call_origin": "autorouter_classifier", + "_managed_background_response": True, "_guardrail_pipelines": [{"name": "spoofed"}], "_pipeline_managed_guardrails": ["evaded"], "safe_user_metadata": "kept", @@ -1246,6 +1247,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "litellm_gateway_injected_cache", "_session_deployment_affinity_ttl", "internal_call_origin", + "_managed_background_response", "_guardrail_pipelines", "_pipeline_managed_guardrails", } diff --git a/tests/unit/test_responses_streaming_container_ownership.py b/tests/unit/test_responses_streaming_container_ownership.py index 07cc309798b..65889052bce 100644 --- a/tests/unit/test_responses_streaming_container_ownership.py +++ b/tests/unit/test_responses_streaming_container_ownership.py @@ -7,7 +7,7 @@ which returns ``FallbackResponsesStreamWrapper``. That wrapper set ``self.completed_response = None`` in __init__ and never updated it, so the proxy's container-ownership hook (which reads ``getattr(stream_response, "completed_response", None)`` via -``ProxyBaseLLMRequestProcessing._extract_completed_responses_response``) +``ProxyBaseLLMRequestProcessing.extract_completed_responses_response``) saw None on every streaming call and silently recorded nothing — follow-up ``GET /v1/containers//files`` then 403'd for the very key that created the container. @@ -194,7 +194,7 @@ class TestProxyOwnershipHookReadsCompletedResponse: wrapper = wrapper_cls(gen()) asyncio.run(_drain(wrapper)) - extracted = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( + extracted = ProxyBaseLLMRequestProcessing.extract_completed_responses_response( wrapper ) assert extracted is not None diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 3c6ba5ed7b4..b779658e54a 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -10410,14 +10410,16 @@ class TestBackgroundResponseCostPollCooldown: ) assert "dep-1" in self._cooled_down_ids(router) - def test_cost_poll_not_found_does_not_cool_down_deployment(self): + @pytest.mark.parametrize("poll_origin", ["background_response_cost_poll", "background_response_status_poll"]) + def test_background_response_poll_not_found_does_not_cool_down_deployment(self, poll_origin: str): + """A client can delete its background response before the poll reads it, so a 404 on either + poll read says nothing about the deployment's health.""" from datetime import datetime from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.router_utils.router_callbacks.track_deployment_metrics import ( get_deployment_failures_for_current_minute, ) - from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN router = self._router() now = datetime.now() @@ -10427,9 +10429,7 @@ class TestBackgroundResponseCostPollCooldown: "exception": self._not_found(), "litellm_params": { "model_info": {"id": "dep-1"}, - "litellm_metadata": { - INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - }, + "litellm_metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: poll_origin}, }, }, None,