mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(responses): send the billed background response cost to the creating key's and team's loggers (#47506)
* test(responses): the billed background response cost reaches the creating key's and team's loggers * fix(responses): send the billed background response cost to the creating key's and team's loggers * fix(responses): bill an unreadable or rejected creator through proxy loggers, redact and honor disabled callbacks on the billed event * fix(responses): follow the creator's message logging setting on the billed event and read team logging for non-key creators The billed event no longer forces redaction. It follows the creating key's or team's turn_off_message_logging, else the proxy's, and redacts when the creator cannot be read or its settings are rejected. A creator without a virtual key hash bills through its team's settings. function_setup keeps its original callback split, and the billed path uses a separate non-mutating helper. * fix(responses): bill rotated, deleted, and custom auth creators through their own logging and retry failed creator lookups * fix(responses): follow a key's whole grace-period rotation chain to the live key for its billed charge's logging * fix(responses): bill a background response whose creator cannot be read, without the creator's settings * docs(enterprise): say the background response row is claimed before its creator is read * fix(responses): keep the loaded creator rows' budgets and aliases when one creator lookup fails
This commit is contained in:
Mateo Wang
GitHub
parent
7c39dbfab8
commit
480ee653c5
9 changed files with 936 additions and 87 deletions
|
|
@ -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 {}
|
||||
|
|
|
|||
|
|
@ -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": "<retrieve_response>"}],
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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}"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue