diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index bd9440d4754..df1f263a947 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, ManagedObjectRow, + row_metadata, ) import litellm @@ -90,11 +91,6 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li return TeamRepository(prisma_client).table -def _row_metadata(row: object) -> dict[str, object]: - metadata: Final = getattr(row, "metadata", None) - return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column - - class CheckBatchCost: def __init__( self, @@ -1063,7 +1059,9 @@ class CheckBatchCost: if (has_key_row and key_row is None) or (team_id and team_row is None): return True settings: Final = get_dynamic_logging_metadata( - UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id), + UserAPIKeyAuth( + metadata=dict(row_metadata(key_row)), team_metadata=dict(row_metadata(team_row)), team_id=team_id + ), proxy_config, ) callback_vars: Final = (settings.callback_vars if settings is not None else None) or {} diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 1fa461d7ab7..07048cf2c2f 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -7,6 +7,12 @@ attribution, priced against the deployment that served it. The poll's status che with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the managed files hook, so both stay free. The spend event carries the id the client holds, the row's unified_object_id, so its spend log is found under that id with the _background_cost suffix. +The spend event reaches the creating key's or team's own logging callbacks the way that creator's +requests do. Its message content is redacted when the creator's key or team logging turns message +logging off, else as the proxy is set, the way the batch poll decides it. The row is claimed before +the creator is read, so reads that lose the race skip the lookups. A creator whose lookup fails, or +whose logging settings the logger rejects, is still billed, through the proxy's own callbacks only +and with its message content redacted. A read that finds no row, because the create could not store one, stores it under the reading key and bills it the same way, so a response the create left untracked is still billed once. A streamed background create is free on its own; its stream stores the row under the creating key and bills @@ -29,6 +35,7 @@ from typing import TYPE_CHECKING, Final, cast from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, ManagedObjectRow, + ResponseCreator, ) from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -54,11 +61,13 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.proxy.proxy_server import ProxyConfig from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.prisma_protocols import TableActions from litellm.router import Router _STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) +_CALLBACK_LIST_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) class _RoutedModelInfo(BaseModel): @@ -73,6 +82,8 @@ class _RoutedDeployment(BaseModel): _ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...]) _RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses" +_CALLBACK_LIST_FIELDS: Final = frozenset({"success_callback", "failure_callback", "litellm_disabled_callbacks"}) +_REJECTED_CREATOR_LOGGING: Final[Mapping[str, object]] = MappingProxyType({"turn_off_message_logging": True}) @dataclass(frozen=True, slots=True) @@ -89,6 +100,7 @@ class BilledBackgroundResponse: custom_llm_provider: str | None litellm_params: Mapping[str, object] metadata: Mapping[str, object] + creator_logging: Mapping[str, object] @dataclass(frozen=True, slots=True) @@ -107,11 +119,51 @@ def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse: return unpriced -async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +def _string_callbacks(callbacks: object) -> tuple[str, ...]: + if not isinstance(callbacks, list): + return () + return tuple( + callback for callback in _CALLBACK_LIST_ADAPTER.validate_python(callbacks) if isinstance(callback, str) + ) - logging_obj: Final = LiteLLMLogging( + +def _running_proxy_config() -> "ProxyConfig": + from litellm.proxy.proxy_server import proxy_config + + return proxy_config + + +def creator_logging_fields(creator: ResponseCreator, proxy_config: "ProxyConfig | None") -> Mapping[str, object]: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields + + if not creator.settings_read: + return _REJECTED_CREATOR_LOGGING + try: + return key_and_team_logging_fields( + user_api_key_dict=UserAPIKeyAuth( + metadata=dict(creator.key_metadata), + team_id=creator.team_id, + team_metadata=dict(creator.team_metadata), + ), + proxy_config=proxy_config if proxy_config is not None else _running_proxy_config(), + ) + except Exception as e: + verbose_proxy_logger.error( + f"CheckResponsesCost: could not read the creator's logging settings, so only the proxy's " + f"callbacks get the billed event: {e}" + ) + return _REJECTED_CREATOR_LOGGING + + +def _billed_logging_obj(billed: BilledBackgroundResponse, creator_logging: Mapping[str, object]) -> "LiteLLMLogging": + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.utils import split_dynamic_success_callbacks + + sync_callbacks, async_callbacks = split_dynamic_success_callbacks( + _string_callbacks(creator_logging.get("success_callback")) + ) + return LiteLLMLogging( model=billed.model, messages=[{"role": "user", "content": ""}], stream=False, @@ -119,7 +171,27 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi start_time=datetime.now(), litellm_call_id=str(uuid.uuid4()), function_id=str(uuid.uuid4()), + dynamic_success_callbacks=list(sync_callbacks), + dynamic_async_success_callbacks=list(async_callbacks) or None, + kwargs={key: value for key, value in creator_logging.items() if key not in _CALLBACK_LIST_FIELDS}, ) + + +def _creator_or_proxy_logging_obj(billed: BilledBackgroundResponse) -> "LiteLLMLogging": + try: + return _billed_logging_obj(billed, billed.creator_logging) + except Exception as e: + verbose_proxy_logger.error( + f"CheckResponsesCost: the creator's logging settings were rejected, so only the proxy's " + f"callbacks get the billed event: {e}" + ) + return _billed_logging_obj(billed, _REJECTED_CREATOR_LOGGING) + + +async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj: Final = _creator_or_proxy_logging_obj(billed) logging_obj.update_environment_variables( litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)}, optional_params={}, @@ -187,6 +259,7 @@ class CheckResponsesCost: prisma_client: "PrismaClient", llm_router: "Router", log_billed_response: BilledResponseLogger = log_billed_background_response, + proxy_config: "ProxyConfig | None" = None, ): from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -196,6 +269,7 @@ class CheckResponsesCost: self.llm_router: Router = llm_router self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost") self._log_billed_response: Final = log_billed_response + self._proxy_config: Final = proxy_config def _resolve_deployment(self, response_id: str) -> Deployment | None: model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) @@ -307,11 +381,12 @@ class CheckResponsesCost: data={"status": job.status or PENDING_RESPONSE_STATUSES[0]}, ) - async def _billed_response( + def _billed_response( self, job: ManagedObjectRow, target: _PollTarget, response: ResponsesAPIResponse, + creator: ResponseCreator, ) -> BilledBackgroundResponse: deployment: Final = target.deployment return BilledBackgroundResponse( @@ -323,13 +398,14 @@ class CheckResponsesCost: litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}), metadata=MappingProxyType( { - **(await self._creator_attribution.response_metadata(job, job.unified_object_id)), + **creator.metadata, INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, **target.model_metadata, **(_deployment_metadata(deployment) if deployment is not None else {}), **self._requested_model_group(target, job.team_id), } ), + creator_logging=creator_logging_fields(creator, self._proxy_config), ) def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]: @@ -354,7 +430,8 @@ class CheckResponsesCost: verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping") return None try: - charge: Final = await self._log_billed_response(await self._billed_response(job, target, response)) + creator: Final = await self._creator_attribution.response_creator(job, job.unified_object_id) + charge: Final = await self._log_billed_response(self._billed_response(job, target, response, creator)) except BaseException: await self._release_billing_claim(job) raise diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index 2e2c07fd544..a4a40a7a3d7 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -10,12 +10,16 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime -from typing import TYPE_CHECKING, Final, Protocol, cast +from typing import TYPE_CHECKING, Final, Protocol, TypeGuard, cast from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.constants import CLI_SESSION_KEY_PREFIX +from litellm.repositories.table_repositories import ( + DeletedVerificationTokenRepository, + DeprecatedVerificationTokenRepository, +) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository @@ -70,6 +74,7 @@ class ManagedObjectRow(Protocol): _TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object]) _KEY_OWNER_SCOPE_PREFIX: Final = "key:" +_MAX_ROTATION_HOPS: Final = 16 def _string_tags(request_tags: object) -> list[str] | None: @@ -96,6 +101,20 @@ def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteL return TeamRepository(prisma_client).table +def _rotated_token_table( + prisma_client: PrismaClient, +) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]: + return DeprecatedVerificationTokenRepository(prisma_client).table + + +def _deleted_token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedVerificationToken]: + return DeletedVerificationTokenRepository(prisma_client).table + + +def _deleted_team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedTeamTable]: + return TeamRepository(prisma_client).deleted_table + + @dataclass(frozen=True, slots=True) class _CreatorRows: key: prisma_models.LiteLLM_VerificationToken | None @@ -103,6 +122,15 @@ class _CreatorRows: user: prisma_models.LiteLLM_UserTable | None +@dataclass(frozen=True, slots=True) +class ResponseCreator: + metadata: Mapping[str, object] + key_metadata: Mapping[str, object] + team_id: str | None + team_metadata: Mapping[str, object] + settings_read: bool = True + + class ManagedObjectCreatorAttribution: """Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll job's spend row is attributed like the creating request: key hash, team, organization, @@ -117,7 +145,12 @@ class ManagedObjectCreatorAttribution: A background response also carries its end user, the creating key's project and its alias, and the key, team, and user budgets and per-model budgets, the way the read that billed it before did, and never names a key-scoped creator as its user. A key that sets no budget of - its own takes its budget tier's, the same rule auth applies. + its own takes its budget tier's, the same rule auth applies. Its creator's key and team + metadata come along too, so the billed event reaches the key's and team's logging callbacks + like that read did. A rotated key's metadata is read from the live key its rotations lead + to, and a deleted key's or team's from its deleted row. A lookup that fails leaves the + creator's settings unread and the attribution to the rows that still load, so the + response still bills and still counts toward their budgets. """ def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None: @@ -128,31 +161,70 @@ class ManagedObjectCreatorAttribution: rows: Final = await self._rows(job, object_id) return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)} - async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]: - rows: Final = await self._rows(job, object_id) - end_user_id: Final = getattr(job, "end_user_id", None) - return { - "user_api_key_user_id": _creator_user_id(job.created_by), - **_shared_metadata(job, rows), - **({"user_api_key_end_user_id": end_user_id} if end_user_id else {}), - "user_api_key_project_id": getattr(rows.key, "project_id", None), - "user_api_key_project_alias": getattr( - getattr(rows.key, "litellm_project_table", None), "project_alias", None + async def response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator: + try: + return await self._read_response_creator(job, object_id) + except Exception as e: + verbose_proxy_logger.error( + f"{self._poller_name}: could not read the creator of {object_id}, so it bills without their " + f"settings: {e}" + ) + return ResponseCreator( + metadata=_response_creator_metadata(job, await self._rows(job, object_id)), + key_metadata={}, + team_id=getattr(job, "team_id", None), + team_metadata={}, + settings_read=False, + ) + + async def _read_response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator: + api_key: Final = getattr(job, "api_key", None) + team_id: Final = getattr(job, "team_id", None) + rows: Final = _CreatorRows( + key=await self._live_key_row(api_key, job.created_by), + team=await self._live_team_row(team_id), + user=await self._user_row(object_id, _creator_user_id(job.created_by)), + ) + return ResponseCreator( + metadata=_response_creator_metadata(job, rows), + key_metadata=row_metadata(rows.key) if rows.key is not None else await self._retired_key_metadata(api_key), + team_id=team_id, + team_metadata=( + row_metadata(rows.team) if rows.team is not None else await self._deleted_team_metadata(team_id) ), - **_key_budget_fields(rows.key), - **_budget_fields( - rows.team, - "user_api_key_team_spend", - "user_api_key_team_max_budget", - "user_api_key_team_model_max_budget", - ), - **_budget_fields( - rows.user, - "user_api_key_user_spend", - "user_api_key_user_max_budget", - "user_api_key_user_model_max_budget", - ), - } + ) + + async def _retired_key_metadata(self, api_key: str | None) -> Mapping[str, object]: + if not _names_a_key_row(api_key): + return {} + current_token: Final = await self._token_rotated_to(api_key, _MAX_ROTATION_HOPS) + current_key: Final = ( + await _token_table(self._prisma_client).find_unique(where={"token": current_token}) + if current_token != api_key + else None + ) + if current_key is not None: + return row_metadata(current_key) + return row_metadata( + await _deleted_token_table(self._prisma_client).find_first( + where={"token": current_token}, order={"deleted_at": "desc"} + ) + ) + + async def _token_rotated_to(self, token: str, hops_left: int) -> str: + rotation: Final = await _rotated_token_table(self._prisma_client).find_unique(where={"token": token}) + if rotation is None or hops_left == 0: + return token + return await self._token_rotated_to(rotation.active_token_id, hops_left - 1) + + async def _deleted_team_metadata(self, team_id: str | None) -> Mapping[str, object]: + if not team_id: + return {} + return row_metadata( + await _deleted_team_table(self._prisma_client).find_first( + where={"team_id": team_id}, order={"deleted_at": "desc"} + ) + ) async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows: return _CreatorRows( @@ -175,25 +247,57 @@ class ManagedObjectCreatorAttribution: async def _key_row( self, object_id: str, api_key: str | None, created_by: str | None ) -> prisma_models.LiteLLM_VerificationToken | None: - if not api_key or _is_cli_session_key(api_key, created_by): - return None try: - return await _token_table(self._prisma_client).find_unique( - where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True} - ) + return await self._live_key_row(api_key, created_by) except Exception as e: verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}") return None - async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: - if not team_id: + async def _live_key_row( + self, api_key: str | None, created_by: str | None + ) -> prisma_models.LiteLLM_VerificationToken | None: + if not api_key or _is_cli_session_key(api_key, created_by): return None + return await _token_table(self._prisma_client).find_unique( + where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True} + ) + + async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: try: - return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + return await self._live_team_row(team_id) except Exception as e: verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}") return None + async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: + if not team_id: + return None + return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + + +def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: + end_user_id: Final = getattr(job, "end_user_id", None) + return { + "user_api_key_user_id": _creator_user_id(job.created_by), + **_shared_metadata(job, rows), + **({"user_api_key_end_user_id": end_user_id} if end_user_id else {}), + "user_api_key_project_id": getattr(rows.key, "project_id", None), + "user_api_key_project_alias": getattr(getattr(rows.key, "litellm_project_table", None), "project_alias", None), + **_key_budget_fields(rows.key), + **_budget_fields( + rows.team, + "user_api_key_team_spend", + "user_api_key_team_max_budget", + "user_api_key_team_model_max_budget", + ), + **_budget_fields( + rows.user, + "user_api_key_user_spend", + "user_api_key_user_max_budget", + "user_api_key_user_model_max_budget", + ), + } + def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: api_key: Final = getattr(job, "api_key", None) @@ -221,6 +325,19 @@ def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, } +def row_metadata(row: object) -> Mapping[str, object]: + metadata: Final = getattr(row, "metadata", None) + if not isinstance(metadata, dict): + return {} + return cast(Mapping[str, object], metadata) # cast-ok: prisma Json object column + + +def _names_a_key_row(api_key: object) -> TypeGuard[str]: + from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash + + return isinstance(api_key, str) and is_valid_sha256_hash(api_key) + + def _is_cli_session_key(api_key: str, created_by: str | None) -> bool: return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index f8e6498502b..3ac270fcf2c 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1218,6 +1218,32 @@ def get_dynamic_logging_metadata( return callback_settings_obj +def _callback_settings_fields(callback_settings_obj: TeamCallbackMetadata) -> Mapping[str, object]: + callback_vars: Final = callback_settings_obj.callback_vars + return { + "success_callback": callback_settings_obj.success_callback, + "failure_callback": callback_settings_obj.failure_callback, + **({**callback_vars, TRUSTED_CALLBACK_VARS_FIELD: callback_vars} if callback_vars is not None else {}), + } + + +def key_and_team_logging_fields(user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig) -> Mapping[str, object]: + callback_settings_obj: Final = get_dynamic_logging_metadata( + user_api_key_dict=user_api_key_dict, proxy_config=proxy_config + ) + disabled_callbacks: Final = (user_api_key_dict.metadata or {}).get("litellm_disabled_callbacks") + return MappingProxyType( + { + **(_callback_settings_fields(callback_settings_obj) if callback_settings_obj is not None else {}), + **( + {"litellm_disabled_callbacks": disabled_callbacks} + if disabled_callbacks and isinstance(disabled_callbacks, list) + else {} + ), + } + ) + + _get_dynamic_logging_metadata: Final = get_dynamic_logging_metadata @@ -2841,27 +2867,7 @@ async def add_litellm_data_to_request( request=request, request_data=data, metadata_key=_metadata_variable_name ) - # Team Callbacks controls - callback_settings_obj: Final = get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=proxy_config - ) - if callback_settings_obj is not None: - data["success_callback"] = callback_settings_obj.success_callback - data["failure_callback"] = callback_settings_obj.failure_callback - - if callback_settings_obj.callback_vars is not None: - # unpack callback_vars in data - for k, v in callback_settings_obj.callback_vars.items(): - data[k] = v - # Callbacks that must not honour request-supplied credentials read this - # proxy-owned field instead of the raw request kwargs. - data[TRUSTED_CALLBACK_VARS_FIELD] = callback_settings_obj.callback_vars - - # Add disabled callbacks from key metadata - if user_api_key_dict.metadata and "litellm_disabled_callbacks" in user_api_key_dict.metadata: - disabled_callbacks: Final = user_api_key_dict.metadata["litellm_disabled_callbacks"] - if disabled_callbacks and isinstance(disabled_callbacks, list): - data["litellm_disabled_callbacks"] = disabled_callbacks + data.update(key_and_team_logging_fields(user_api_key_dict=user_api_key_dict, proxy_config=proxy_config)) # Guardrails from key/team metadata and policy engine await move_guardrails_to_metadata( diff --git a/litellm/utils.py b/litellm/utils.py index 4687e4fa27a..897138a940c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1580,6 +1580,22 @@ def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: return kwargs.get("timeout") if stream_timeout is None else stream_timeout +def _runs_as_async_success_callback(callback: str | Callable[..., object] | CustomLogger) -> bool: + get_coroutine_checker: Final = litellm_utils.get_coroutine_checker + return get_coroutine_checker().is_async_callable(callback) or callback in ("dynamodb", "s3") + + +def split_dynamic_success_callbacks( + callbacks: Sequence[str | Callable[..., object] | CustomLogger], +) -> tuple[ + tuple[str | Callable[..., object] | CustomLogger, ...], tuple[str | Callable[..., object] | CustomLogger, ...] +]: + return ( + tuple(callback for callback in callbacks if not _runs_as_async_success_callback(callback)), + tuple(callback for callback in callbacks if _runs_as_async_success_callback(callback)), + ) + + def check_coroutine(value) -> bool: get_coroutine_checker: Final = litellm_utils.get_coroutine_checker return get_coroutine_checker().is_async_callable(value) diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index 6d481a98fd0..1b7fc975217 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -1,8 +1,10 @@ import asyncio +import base64 import json import uuid from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor +from itertools import chain from typing import Final import httpx @@ -10,8 +12,12 @@ import openai import pytest from integration._support.client import Gateway, Scenario, eventually, held, list_value, object_value, string_value from integration._support.database import read_rows, write_rows +from integration._support.otlp_sink import decode_spans from integration._support.upstream import delete_scenario, register_scenario +from integration._support.wire import Request, wire_server from integration.caching.response_cache_case import entry_key, redis_store +from integration.security._callback_traffic import langfuse_sink +from integration.security._sinks import Recorder from integration.spend._background_responses import ( ATTRIBUTION_COLUMNS, CONCURRENT_READS, @@ -60,7 +66,7 @@ from integration.spend._background_responses import ( ) from openai._legacy_response import LegacyAPIResponse from openai.types.responses import Response -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN @@ -823,3 +829,78 @@ def test_a_client_cannot_make_a_background_create_free_by_forging_the_managed_ma scenario.cleanups.callback(forget_row, response_id) assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == free_reads("completed", 2) assert_charged_exactly(lambda: key_spend(key), PRICE) + + +LANGFUSE_PUBLIC_KEY: Final = "pk-lf-background" +LANGFUSE_SECRET_KEY: Final = "sk-lf-background" +LANGFUSE_TRACES_PATH: Final = "/api/public/otel/v1/traces" +LANGFUSE_COST: Final = "langfuse.observation.cost_details" +_LANGFUSE_COST_DETAILS: Final = TypeAdapter(dict[str, float]) + + +def _langfuse_callback(host: str) -> dict[str, JsonValue]: + return { + "callback_name": "langfuse", + "callback_type": "success", + "callback_vars": { + "langfuse_public_key": LANGFUSE_PUBLIC_KEY, + "langfuse_secret_key": LANGFUSE_SECRET_KEY, + "langfuse_host": host, + }, + } + + +def _key_with_its_own_logging(scenario: Scenario, langfuse_url: str) -> str: + return scenario.key(metadata={"logging": [_langfuse_callback(langfuse_url)]}) + + +def _key_in_a_team_with_logging(scenario: Scenario, langfuse_url: str) -> str: + team_id: Final = scenario.team() + added: Final = scenario.gateway.request("POST", f"/team/{team_id}/callback", _langfuse_callback(langfuse_url)) + assert added.status_code == 200, added.text + return scenario.key(team_id=team_id) + + +LOGGING_ON: Final[Mapping[str, Callable[[Scenario, str], str]]] = { + "the_key": _key_with_its_own_logging, + "the_team": _key_in_a_team_with_logging, +} + + +def _trace_exports(langfuse: Recorder) -> tuple[Request, ...]: + return tuple(request for request in langfuse.requests() if request.target.startswith(LANGFUSE_TRACES_PATH)) + + +def _exported_costs(langfuse: Recorder) -> tuple[float, ...]: + spans: Final = tuple( + chain.from_iterable( + decode_spans(export.body, export.headers.get("content-type", "")) for export in _trace_exports(langfuse) + ) + ) + return tuple( + _LANGFUSE_COST_DETAILS.validate_json(string_value(span["attributes"][LANGFUSE_COST]))["total"] + for span in spans + if LANGFUSE_COST in span["attributes"] + ) + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("logging_on", LOGGING_ON.keys()) +def test_the_billed_cost_reaches_the_langfuse_project_set_on_the_creating_key_or_team( + gateway: Gateway, logging_on: str +) -> None: + with wire_server(langfuse_sink) as wire, gateway.scenario() as scenario: + langfuse: Final = Recorder(wire) + key: Final = LOGGING_ON[logging_on](scenario, wire.url) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, key, upstream.model) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + exported: Final = eventually( + lambda: _exported_costs(langfuse), + lambda costs: pytest.approx(PRICE) in costs, + seconds=SPEND_LANDING_SECONDS, + ) + assert [cost for cost in exported if cost > 0] == [pytest.approx(PRICE)] + credentials: Final = base64.b64encode(f"{LANGFUSE_PUBLIC_KEY}:{LANGFUSE_SECRET_KEY}".encode()).decode() + assert {export.headers.get("authorization") for export in _trace_exports(langfuse)} == {f"Basic {credentials}"} diff --git a/tests/unit/integrations/test_responses_background_cost.py b/tests/unit/integrations/test_responses_background_cost.py index 1bf95e35c4d..7d266c94700 100644 --- a/tests/unit/integrations/test_responses_background_cost.py +++ b/tests/unit/integrations/test_responses_background_cost.py @@ -282,6 +282,12 @@ class TestCheckResponsesCost: client = MagicMock() client.db = MagicMock() client.db.litellm_managedobjecttable = MagicMock() + client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + client.db.litellm_deprecatedverificationtoken.find_unique = AsyncMock(return_value=None) + client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None) + client.db.litellm_deletedteamtable.find_first = AsyncMock(return_value=None) return client @pytest.fixture diff --git a/tests/unit/proxy/common_utils/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py index d78c3ef6185..b8d5ddef3d2 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -4,6 +4,8 @@ Unit tests for CheckResponsesCost class import asyncio import dataclasses +import hashlib +import json from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime @@ -127,10 +129,21 @@ class _InMemoryManagedObjectTable: return len(matched) +def _prisma_with_no_creator_rows() -> MagicMock: + prisma_client: Final = MagicMock() + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_deprecatedverificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_deletedteamtable.find_first = AsyncMock(return_value=None) + return prisma_client + + def _poller_over(table: _InMemoryManagedObjectTable, billed_logger: _RecordingBilledLogger): from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost - prisma_client: Final = MagicMock() + prisma_client: Final = _prisma_with_no_creator_rows() prisma_client.db.litellm_managedobjecttable = table poller: Final = CheckResponsesCost( proxy_logging_obj=MagicMock(), @@ -174,8 +187,7 @@ class TestCheckResponsesCost: @pytest.fixture def mock_prisma_client(self): """Create a mock Prisma client""" - client = MagicMock() - client.db = MagicMock() + client = _prisma_with_no_creator_rows() client.db.litellm_managedobjecttable = MagicMock() return client @@ -1702,26 +1714,60 @@ class _PricedBilling: table: _InMemoryManagedObjectTable response_id: str cost: float | None + logging_obj: object + + +@dataclass(frozen=True) +class _PricedPoller: + poller: object + prisma_client: MagicMock + table: _InMemoryManagedObjectTable + response_id: str + + async def bill(self, read: ResponsesAPIResponse) -> object: + from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + assert isinstance(self.poller, CheckResponsesCost) + charge: Final = await self.poller.bill_if_finished(response_id=self.response_id, response=read) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + return charge _PRICED_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType( {"model": "openai/gpt-5.6", "api_key": "sk-unused", "input_cost_per_token": 0.001, "output_cost_per_token": 0.002} ) +_CREATING_KEY_HASH: Final = hashlib.sha256(b"sk-alice").hexdigest() -async def _bill_through_a_priced_deployment( - read: ResponsesAPIResponse, +def _team_alpha_when_its_rows_exist( + team_row: object, deleted_team_row: object, default_team_settings: tuple[Mapping[str, object], ...] +) -> str | None: + return "team-alpha" if team_row is not None or deleted_team_row is not None or default_team_settings else None + + +def _priced_poller( litellm_params: Mapping[str, object] = _PRICED_LITELLM_PARAMS, model_info: Mapping[str, object] = MappingProxyType({}), + key_row: object = None, + team_row: object = None, + key_lookup_error: Exception | None = None, + default_team_settings: tuple[Mapping[str, object], ...] = (), + api_key: str = _CREATING_KEY_HASH, + rotated_to: Mapping[str, object] = MappingProxyType({}), + deleted_key_row: object = None, + deleted_team_row: object = None, model_name: str = "priced-gpt", requested_model: str = "priced-gpt", model_group_alias: Mapping[str, str] = MappingProxyType({}), team_id: str | None = None, -) -> _PricedBilling: +) -> _PricedPoller: + from types import SimpleNamespace + import litellm from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.responses.utils import ResponsesAPIRequestUtils response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( @@ -1733,17 +1779,31 @@ async def _bill_through_a_priced_deployment( id="job-priced", unified_object_id=response_id, status="queued", - api_key="hash-alice", + api_key=api_key, file_object={"model": requested_model}, - team_id=team_id, + team_id=team_id or _team_alpha_when_its_rows_exist(team_row, deleted_team_row, default_team_settings), ), ) ) - prisma_client: Final = MagicMock() + key_rows: Final = {api_key: key_row, **{token: row for token, row in rotated_to.items()}} + prisma_client: Final = _prisma_with_no_creator_rows() prisma_client.db.litellm_managedobjecttable = table prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) - prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) - prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=key_lookup_error or (lambda where, **_: key_rows.get(where["token"])) + ) + rotation_chain: Final = (api_key, *rotated_to) + rotated_from: Final = dict(zip(rotation_chain, rotation_chain[1:])) + prisma_client.db.litellm_deprecatedverificationtoken.find_unique = AsyncMock( + side_effect=lambda where: ( + SimpleNamespace(token=where["token"], active_token_id=rotated_from[where["token"]]) + if where["token"] in rotated_from + else None + ) + ) + prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=deleted_key_row) + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + prisma_client.db.litellm_deletedteamtable.find_first = AsyncMock(return_value=deleted_team_row) router: Final = litellm.Router( model_list=[ { @@ -1754,11 +1814,71 @@ async def _bill_through_a_priced_deployment( ], model_group_alias=dict(model_group_alias), ) - charge: Final = await CheckResponsesCost( - proxy_logging_obj=MagicMock(), prisma_client=prisma_client, llm_router=router - ).bill_if_finished(response_id=response_id, response=read) - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) - return _PricedBilling(table=table, response_id=response_id, cost=charge.cost if charge is not None else None) + return _PricedPoller( + poller=CheckResponsesCost( + proxy_logging_obj=MagicMock(), + prisma_client=prisma_client, + llm_router=router, + proxy_config=_proxy_config_with(default_team_settings), + ), + prisma_client=prisma_client, + table=table, + response_id=response_id, + ) + + +async def _bill_through_a_priced_deployment( + read: ResponsesAPIResponse, + litellm_params: Mapping[str, object] = _PRICED_LITELLM_PARAMS, + model_info: Mapping[str, object] = MappingProxyType({}), + key_row: object = None, + team_row: object = None, + key_lookup_error: Exception | None = None, + default_team_settings: tuple[Mapping[str, object], ...] = (), + api_key: str = _CREATING_KEY_HASH, + rotated_to: Mapping[str, object] = MappingProxyType({}), + deleted_key_row: object = None, + deleted_team_row: object = None, + model_name: str = "priced-gpt", + requested_model: str = "priced-gpt", + model_group_alias: Mapping[str, str] = MappingProxyType({}), + team_id: str | None = None, +) -> _PricedBilling: + priced: Final = _priced_poller( + litellm_params=litellm_params, + model_info=model_info, + key_row=key_row, + team_row=team_row, + key_lookup_error=key_lookup_error, + default_team_settings=default_team_settings, + api_key=api_key, + rotated_to=rotated_to, + deleted_key_row=deleted_key_row, + deleted_team_row=deleted_team_row, + model_name=model_name, + requested_model=requested_model, + model_group_alias=model_group_alias, + team_id=team_id, + ) + charge: Final = await priced.bill(read) + return _PricedBilling( + table=priced.table, + response_id=priced.response_id, + cost=getattr(charge, "cost", None), + logging_obj=getattr(charge, "logging_obj", None), + ) + + +def _proxy_config_with(default_team_settings: tuple[Mapping[str, object], ...]): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config: Final = ProxyConfig() + proxy_config.config = ( + {"litellm_settings": {"default_team_settings": [dict(settings) for settings in default_team_settings]}} + if default_team_settings + else {} + ) + return proxy_config @pytest.mark.asyncio @@ -1780,7 +1900,7 @@ async def test_the_billing_event_prices_the_fetched_usage_at_the_deployment_rate assert billing.cost == pytest.approx(_PRICED_COST) assert payload["model_id"] == "deployment-custom" assert payload["model_group"] == "priced-gpt" - assert payload["metadata"]["user_api_key_hash"] == "hash-alice" + assert payload["metadata"]["user_api_key_hash"] == _CREATING_KEY_HASH assert read.hidden_params == {"response_cost": 0.0} assert billing.table.statuses == {"job-priced": "completed"} @@ -1945,3 +2065,428 @@ async def test_the_billing_event_keys_its_spend_row_apart_from_the_read_it_bills assert billed_request_id != billing.response_id assert billed_request_id.startswith(billing.response_id) + +_CREATOR_LANGFUSE_VARS: Final = MappingProxyType( + { + "langfuse_public_key": "pk-lf-creator", + "langfuse_secret_key": "sk-lf-creator", + "langfuse_host": "http://127.0.0.1:9", + } +) + + +def _key_row(metadata: Mapping[str, object]) -> object: + from types import SimpleNamespace + + return SimpleNamespace(key_alias=None, organization_id=None, spend=0.0, max_budget=None, metadata=dict(metadata)) + + +def _team_row(metadata: Mapping[str, object]) -> object: + from types import SimpleNamespace + + return SimpleNamespace(team_alias=None, organization_id=None, spend=0.0, max_budget=None, metadata=dict(metadata)) + + +def _langfuse_callback_settings(callback_vars: Mapping[str, object]) -> Mapping[str, object]: + return {"callback_settings": {"success_callback": ["langfuse"], "callback_vars": dict(callback_vars)}} + + +@dataclass(frozen=True) +class _CreatorLogging: + key_row: object + team_row: object + default_team_settings: tuple[Mapping[str, object], ...] + + +def _creator_logging(logging_on: str) -> _CreatorLogging: + logging: Final = { + "logging": [ + {"callback_name": "langfuse", "callback_type": "success", "callback_vars": dict(_CREATOR_LANGFUSE_VARS)} + ] + } + team_metadata: Final = { + "team logging": logging, + "team callback_settings": _langfuse_callback_settings(_CREATOR_LANGFUSE_VARS), + } + return _CreatorLogging( + key_row=_key_row(logging if logging_on == "key logging" else {}), + team_row=_team_row(team_metadata.get(logging_on, {})), + default_team_settings=( + ({"team_id": "team-alpha", "success_callback": ["langfuse"], **_CREATOR_LANGFUSE_VARS},) + if logging_on == "config team settings" + else () + ), + ) + + +def _recording_langfuse(): + from litellm.integrations.custom_logger import CustomLogger + from litellm.integrations.langfuse.langfuse_prompt_management import LangfusePromptManagement + + class _RecordingLangfuse(LangfusePromptManagement): + def __init__(self) -> None: + CustomLogger.__init__(self) + self.events: tuple[tuple[object, object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + params: Final = kwargs["standard_callback_dynamic_params"] + assert isinstance(params, Mapping) + self.events = (*self.events, (kwargs.get("response_cost"), params.get("langfuse_public_key"))) + + return _RecordingLangfuse() + + +async def _log_a_request_from(creator: _CreatorLogging) -> None: + import litellm + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields + + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello", + **key_and_team_logging_fields( + user_api_key_dict=UserAPIKeyAuth( + metadata=getattr(creator.key_row, "metadata"), + team_id="team-alpha", + team_metadata=getattr(creator.team_row, "metadata"), + ), + proxy_config=_proxy_config_with(creator.default_team_settings), + ), + ) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_logs_to_langfuse", [False, True]) +@pytest.mark.parametrize( + "logging_on", ["key logging", "team logging", "team callback_settings", "config team settings"] +) +async def test_the_billing_event_reaches_the_langfuse_projects_the_creators_own_requests_log_to( + monkeypatch: pytest.MonkeyPatch, logging_on: str, proxy_logs_to_langfuse: bool +): + """A key or team that logs to its own Langfuse project must see the billed charge in the same projects + its own requests log to, wherever that logging is configured and whether or not the proxy logs to + Langfuse too, or the charge only shows in the proxy's spend logs.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + + langfuse: Final = _recording_langfuse() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [langfuse] if proxy_logs_to_langfuse else []) + creator: Final = _creator_logging(logging_on) + + await _log_a_request_from(creator) + request_projects: Final = tuple(public_key for _, public_key in langfuse.events) + await _bill_through_a_priced_deployment( + _terminal_response(), + key_row=creator.key_row, + team_row=creator.team_row, + default_team_settings=creator.default_team_settings, + ) + billed_events: Final = langfuse.events[len(request_projects) :] + + assert _CREATOR_LANGFUSE_VARS["langfuse_public_key"] in request_projects + assert tuple(public_key for _, public_key in billed_events) == request_projects + assert all(cost == pytest.approx(_PRICED_COST) for cost, _ in billed_events) + + +def _rotations_for(retired: str, key_logging: object) -> Mapping[str, object]: + rotated_once: Final = hashlib.sha256(b"sk-alice-rotated").hexdigest() + if retired == "rotated key": + return {rotated_once: key_logging} + if retired == "key rotated twice": + return {rotated_once: None, hashlib.sha256(b"sk-alice-rotated-again").hexdigest(): key_logging} + return {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retired", ["rotated key", "key rotated twice", "deleted key", "deleted team"]) +async def test_a_creator_rotated_or_deleted_before_the_bill_still_gets_the_charge_in_its_own_langfuse_project( + monkeypatch: pytest.MonkeyPatch, retired: str +): + """A key rotated or deleted, or a team deleted, while its background response ran logged that response's + create to its own project, so the charge follows the settings kept on the key it was rotated to or on + its deleted row instead of leaving that project for the proxy's callbacks alone.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + + langfuse: Final = _recording_langfuse() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + key_logging: Final = _creator_logging("key logging").key_row + + await _bill_through_a_priced_deployment( + _terminal_response(), + rotated_to=_rotations_for(retired, key_logging), + deleted_key_row=key_logging if retired == "deleted key" else SimpleNamespace(metadata={}), + deleted_team_row=_creator_logging("team callback_settings").team_row if retired == "deleted team" else None, + ) + + assert langfuse.events == ((pytest.approx(_PRICED_COST), _CREATOR_LANGFUSE_VARS["langfuse_public_key"]),) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failing_lookup, live_key_logging, team_logging", + [ + ("litellm_verificationtoken.find_unique", "key logging", "none"), + ("litellm_teamtable.find_unique", "none", "team callback_settings"), + ("litellm_deprecatedverificationtoken.find_unique", "none", "none"), + ("litellm_deletedverificationtoken.find_first", "none", "none"), + ], + ids=["live key", "live team", "rotated key", "deleted key"], +) +async def test_a_creator_lookup_that_keeps_failing_still_bills_the_response_once_without_the_creators_settings( + monkeypatch: pytest.MonkeyPatch, failing_lookup: str, live_key_logging: str, team_logging: str +): + """A creator lookup that fails on every read must not leave the response unbilled. It bills once under the + key and team the managed row stores, through the proxy's callbacks only, with the text left out since the + creator's redaction setting could not be read.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + + langfuse: Final = _recording_langfuse() + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + priced: Final = _priced_poller( + key_row=_creator_logging(live_key_logging).key_row if live_key_logging != "none" else None, + team_row=_creator_logging(team_logging).team_row if team_logging != "none" else None, + ) + table_name, method = failing_lookup.split(".") + setattr( + getattr(priced.prisma_client.db, table_name), + method, + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + + charge: Final = await priced.bill(_read_with_text("the stored answer")) + await priced.bill(_read_with_text("the stored answer")) + + assert getattr(charge, "cost", None) == pytest.approx(_PRICED_COST) + assert priced.table.statuses == {"job-priced": "completed"} + assert langfuse.events == () + (payload,) = capture.payloads + assert payload["response_cost"] == pytest.approx(_PRICED_COST) + assert "the stored answer" not in json.dumps(payload, default=str) + metadata: Final = payload["metadata"] + assert isinstance(metadata, Mapping) + assert metadata["user_api_key_hash"] == _CREATING_KEY_HASH + assert metadata["user_api_key_team_id"] == ("team-alpha" if team_logging != "none" else None) + + +_TEAM_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 2.0, "budget_duration": "7d"}} +_LOADED_KEY_ATTRIBUTION: Final = {"user_api_key_alias": "prod-key", "user_api_key_max_budget": 5.0} +_LOADED_TEAM_ATTRIBUTION: Final = { + "user_api_key_team_alias": "Team Alpha", + "user_api_key_team_max_budget": 50.0, + "user_api_key_team_model_max_budget": _TEAM_MODEL_BUDGET, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failing_lookup, key_is_live, loaded_attribution", + [ + ("litellm_verificationtoken.find_unique", True, _LOADED_TEAM_ATTRIBUTION), + ("litellm_teamtable.find_unique", True, _LOADED_KEY_ATTRIBUTION), + ("litellm_deprecatedverificationtoken.find_unique", False, _LOADED_TEAM_ATTRIBUTION), + ("litellm_deletedverificationtoken.find_first", False, _LOADED_TEAM_ATTRIBUTION), + ], + ids=["live key", "live team", "rotated key", "deleted key"], +) +async def test_a_creator_lookup_that_fails_keeps_the_aliases_and_budgets_of_the_rows_that_did_load( + failing_lookup: str, key_is_live: bool, loaded_attribution: Mapping[str, object] +): + """One failed creator lookup must not strip what the other rows give the charge: the team's per-model budget + only counts a billed charge whose event carries it, and the aliases name who spent it.""" + table: Final = _InMemoryManagedObjectTable( + rows=( + _StoredResponse( + id="job-attributed", + unified_object_id="resp_attributed", + status="queued", + api_key=_CREATING_KEY_HASH, + team_id="team-alpha", + ), + ) + ) + billed_logger: Final = _RecordingBilledLogger() + poller: Final = _poller_over(table, billed_logger) + poller.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=( + SimpleNamespace(key_alias="prod-key", organization_id=None, spend=4.0, max_budget=5.0, metadata={}) + if key_is_live + else None + ) + ) + poller.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=SimpleNamespace( + team_alias="Team Alpha", + organization_id=None, + spend=40.0, + max_budget=50.0, + model_max_budget=_TEAM_MODEL_BUDGET, + metadata={}, + ) + ) + table_name, method = failing_lookup.split(".") + setattr( + getattr(poller.prisma_client.db, table_name), + method, + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + + await poller.bill_if_finished(response_id="resp_attributed", response=_terminal_response()) + + assert table.statuses == {"job-attributed": "completed"} + (billed,) = billed_logger.billed + assert {field: billed.metadata.get(field) for field in loaded_attribution} == loaded_attribution + + +@pytest.mark.asyncio +async def test_a_team_logging_setting_the_logger_rejects_still_bills_the_response_once(monkeypatch: pytest.MonkeyPatch): + """A team whose stored callback settings the logger rejects, such as an os.environ reference it no + longer resolves, must not stop the bill: the charge reaches the proxy's callbacks and the response is + not put back for every later poll to fail on again.""" + import litellm + + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + + billing: Final = await _bill_through_a_priced_deployment( + _terminal_response(), + key_row=_key_row({}), + team_row=_team_row(_langfuse_callback_settings({"langfuse_public_key": "os.environ/QA_UNSET_LANGFUSE_KEY"})), + ) + + (payload,) = capture.payloads + assert payload["response_cost"] == pytest.approx(_PRICED_COST) + assert billing.cost == pytest.approx(_PRICED_COST) + assert billing.table.statuses == {"job-priced": "completed"} + + +def _read_with_text(text: str) -> ResponsesAPIResponse: + return _terminal_response().model_copy( + update={ + "output": [ + { + "type": "message", + "id": "msg_stored", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ] + } + ) + + +def _key_logging_to_langfuse(callback_vars: Mapping[str, object]) -> object: + return _key_row( + {"logging": [{"callback_name": "langfuse", "callback_type": "success", "callback_vars": dict(callback_vars)}]} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("key_row", "deleted_key_row", "text_logged"), + [ + (_key_logging_to_langfuse(_CREATOR_LANGFUSE_VARS), None, True), + (_key_logging_to_langfuse({**_CREATOR_LANGFUSE_VARS, "turn_off_message_logging": "true"}), None, False), + (None, _key_logging_to_langfuse({**_CREATOR_LANGFUSE_VARS, "turn_off_message_logging": "true"}), False), + (None, None, True), + ], + ids=[ + "creator keeps message logging", + "creator turns message logging off", + "deleted creator turned message logging off", + "creator has no key row anywhere", + ], +) +async def test_the_billing_event_logs_the_response_text_only_as_the_creators_logging_settings_allow( + monkeypatch: pytest.MonkeyPatch, key_row: object, deleted_key_row: object, text_logged: bool +): + """The billed event follows the creating key's or team's message logging setting, a deleted key's included, + which falls back to the proxy's the way the batch poll decides it, so a key with no stored settings + anywhere logs the text exactly as the proxy is set.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + + capture: Final = _spend_capture_logger() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [_recording_langfuse()]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [capture]) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + + await _bill_through_a_priced_deployment( + _read_with_text("the stored answer"), key_row=key_row, deleted_key_row=deleted_key_row + ) + + (payload,) = capture.payloads + assert payload["response_cost"] == pytest.approx(_PRICED_COST) + assert ("the stored answer" in json.dumps(payload, default=str)) is text_logged + + +@pytest.mark.asyncio +@pytest.mark.parametrize("caller_token", ["oauth2-raw-access-token", "sk-custom-auth-1234"]) +async def test_a_response_created_without_a_virtual_key_bills_into_its_teams_langfuse_project( + monkeypatch: pytest.MonkeyPatch, caller_token: str +): + """Under OAuth2 the row keeps the caller's raw token, and under custom auth the hash of an sk- token that + has no key row, so the charge follows the team's logging settings like that caller's own requests.""" + from litellm.proxy._types import UserAPIKeyAuth + + import litellm + from litellm.litellm_core_utils import litellm_logging + + langfuse: Final = _recording_langfuse() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + + await _bill_through_a_priced_deployment( + _terminal_response(), + api_key=UserAPIKeyAuth(api_key=caller_token).api_key or caller_token, + team_row=_creator_logging("team callback_settings").team_row, + ) + + assert langfuse.events == ((pytest.approx(_PRICED_COST), _CREATOR_LANGFUSE_VARS["langfuse_public_key"]),) + + +@pytest.mark.asyncio +async def test_a_key_that_disabled_a_callback_has_its_charge_logged_wherever_its_own_requests_are( + monkeypatch: pytest.MonkeyPatch, +): + """A key's stored list of disabled callbacks must treat the billed charge like that key's own requests, or + a project that shows the key's requests at $0 never shows the charge they add up to.""" + import litellm + from litellm.litellm_core_utils import litellm_logging + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", True) + langfuse: Final = _recording_langfuse() + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [langfuse]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", [langfuse]) + creator: Final = _CreatorLogging( + key_row=_key_row({"litellm_disabled_callbacks": ["langfuse"]}), team_row=_team_row({}), default_team_settings=() + ) + + await _log_a_request_from(creator) + request_events: Final = langfuse.events + await _bill_through_a_priced_deployment(_terminal_response(), key_row=creator.key_row) + billed_events: Final = langfuse.events[len(request_events) :] + + assert len(billed_events) == len(request_events) + assert all(cost == pytest.approx(_PRICED_COST) for cost, _ in billed_events) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 32b77c3b564..abc1fe88c7a 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1079,6 +1079,9 @@ def _managed_files_over(table: _UniqueManagedObjectTable) -> tuple[MagicMock, Ma prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_deprecatedverificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_deletedteamtable.find_first = AsyncMock(return_value=None) proxy_logging: Final = MagicMock() proxy_logging.get_proxy_hook = MagicMock( return_value=_PROXY_LiteLLMManagedFiles(internal_usage_cache=DualCache(), prisma_client=prisma_client)