scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 138

fix(responses): bill a background response once, on the first read or poll that sees it finished (#45703)

Browse source 
* fix(responses): bill a background response once, from the poll job

Reading a completed background response with GET /v1/responses/{id} wrote a
full-cost spend row on every read, on top of the row the background cost poll
writes. Reads of a Responses object (GET, DELETE, cancel) now never bill, and
the poll job is the single biller.

The poll job now bills with the creating request's attribution (key hash, team,
org, user, aliases, tags) read off the managed object row, through a helper it
shares with the batch cost poll. The background create persists the creating
key and request tags so that attribution survives until the poll runs.

* test(proxy): a background response read reports zero cost

* fix(responses): bill a background response once, on the first read or poll that sees it finished

A conditional update moves the stored row out of queued or in_progress so only one worker or pod bills it, through one re-read stamped with the creator's attribution. A failed billing read puts the row back for the next poll

* fix(responses): keep the caller's response id out of the billing failure warning

* fix(responses): keep the creating client's User-Agent tags on the one background charge

* fix(responses): keep managed background reads free and bill each response once

The managed files hook marks reads and cancels of proxy-issued response ids, so only those stay free. SDK and database-less reads of a background response still bill, extra_body background creates are tracked, and the poll's status read no longer cools down a deployment on a 404.

* fix(responses): bill finished background responses from the fetched body and store rows Postgres accepts

* fix(responses): give the background billing event its own spend log row and bill responses the create left untracked

* fix(responses): narrow the managed files hook with a protocol instead of an enterprise import

* fix(responses): leave a read the caller retargets with client-side credentials to the cost poll

* fix(responses): bill background responses on the creator's end user and report the charge on the settling read

The read that settles the bill returns the charge in x-litellm-response-cost, the billed event
carries the creator's end user and the key, team, and user budget fields, and its spend row is
named after the id the client got. Override reads, cancels, and deletes settle from the
configured deployment, a streamed create marks its row billed, stored rows keep no model
output, and model insights counts the billed row's spend without a request.

* test(responses): cover background response billing across proxies, reads and chaos

* fix(responses): settle background billing only for callers allowed to address the response

A delete now settles the bill only after the response id access check passes, a streamed
create is marked billed when it sets background inside extra_body or its completed event
arrives after the last chunk, the billing event keeps the deployment's litellm_params
base_model, a key with no user and no team bills no user instead of its owner scope, and the
enterprise cost check import is optional

* fix(responses): record the bound on the pending status IN filters

* fix(responses): send the billed cost split on the settling read and keep the batch cost event as it was

The read that settles a background response now sends the whole cost header
family from the billed event, so the breakdown headers add up to the charge.
Batch cost events go back to main's attribution (creator id as the user, no
budget fields); only background responses get the end user and budgets.

Integration tests expect the charge on the settling read (or none when the
poll billed first), one charge for a streamed create, and SDK flows that poll
with the id each read returns.

* fix(responses): bill each background response once across upgrades, streams, and failed settles

Rows stored before this billing model are found by the issued id and rekeyed by the
provider id, and the cost poll retires one a read already billed. A streamed background
create is free on its own and bills the finished response through the same managed row
before the client sees response.completed, so a racing read cannot bill it again. A
settle that fails retries in the background, a retargeted read with no row tracks and
bills from the configured deployment, and a delete keeps its billed cost header.
Reverts the unrelated json.dumps narrowing in the completion transformation.

* fix(responses): bill cache-polled background streams and match poll rows by the issued id

* test(responses): cover deletes, failed settles, budgets, base models, and cache-polled background billing

Adds integration cells for the background response billing paths the audit
inventory lists: a DELETE that settles the bill and refuses a stranger, a
settle the database refuses that retries once it recovers, the Prometheus
remaining budget, base_model pricing from model_info and litellm_params, the
daily model usage rollup, the end user, the settling read's cost split, stored
rows without model output, a forged managed marker on the create, streamed
creates with background in extra_body, and the cache-polled create

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-15 11:00:14 -07:00 • committed by GitHub
parent c9ce9a733d
commit fbc2e358bc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
41 changed files with 5110 additions and 549 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

202
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -8,21 +8,20 @@ from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
)
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
CLI_SESSION_KEY_PREFIX,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
)
from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
@ -47,29 +46,6 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
)
class _ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@property
def unified_object_id(self) -> str: ...
@property
def created_by(self) -> str | None: ...
@property
def file_object(self) -> object: ...
@property
def org_id(self) -> str | None: ...
@property
def api_key(self) -> str | None: ...
@property
def team_id(self) -> str | None: ...
class _ReadableFileContent(Protocol):
async def read(self) -> bytes: ...
@ -87,22 +63,10 @@ async def _file_content_bytes(file_content: object) -> bytes:
return cast(bytes, file_content)
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
return ManagedObjectRepository(prisma_client).table
def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]":
return UserRepository(prisma_client).table
def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
return VerificationTokenRepository(prisma_client).table
def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]":
return TeamRepository(prisma_client).table
class CheckBatchCost:
def __init__(
self,
@ -118,6 +82,7 @@ class CheckBatchCost:
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost")
# Cached after the first poll cycle. Once we know the column is absent we skip
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
@ -149,133 +114,6 @@ class CheckBatchCost:
return
self.batch_processed_support_confirmed = True
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]:
"""
Look up user email and key alias by user_id for enriching the S3 callback metadata.
Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
Returns an empty dict when user_id is None: batches created by a team or service
account key carry no user id, and find_unique(where={"user_id": None}) raises.
"""
if not user_id:
return {}
try:
user_row: prisma_models.LiteLLM_UserTable | None = await _user_table(self.prisma_client).find_unique(
where={"user_id": user_id}
)
if user_row is None:
return {}
return {
"user_api_key_user_email": getattr(user_row, "user_email", None),
"user_api_key_alias": getattr(user_row, "user_alias", None),
}
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
return {}
async def _get_key_alias(self, batch_id: str, api_key: str | None, created_by: str | None) -> str | None:
"""Resolve the creating virtual key's alias from its hashed token."""
if not api_key:
return None
if created_by and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}":
return api_key
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
self.prisma_client
).find_unique(where={"token": api_key})
return getattr(key_row, "key_alias", None) if key_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
return None
async def _get_team_alias(self, team_id: str | None) -> str | None:
"""Resolve a team's alias from its id."""
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
where={"team_id": team_id}
)
return getattr(team_row, "team_alias", None) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
return None
async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | None:
org_id: Final = getattr(job, "org_id", None)
if org_id:
return org_id
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
if api_key:
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
self.prisma_client
).find_unique(where={"token": api_key})
key_org_id: Final = (
cast(str | None, getattr(key_row, "organization_id", None))
if key_row is not None
else None
)
if key_org_id:
return key_org_id
except Exception as e:
verbose_proxy_logger.error(
f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
f"still trying the team's: {e}"
)
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
where={"team_id": team_id}
)
return cast(str | None, getattr(team_row, "organization_id", None)) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
return None
async def _build_creator_attribution_metadata(
self, job: "_ManagedObjectRow", batch_id: str
) -> dict[str, object]:
"""
Rebuild the spend-tracking metadata for the key, team, and tags that created the
batch so the batch-cost spend log is attributed the same way a non-batch request
is. Rows created before api_key and request_tags were persisted carry only
created_by and team_id, and fall back to those. A named creating key owns
user_api_key_alias; when it has no alias, or the key has since been rotated or
deleted, the field keeps the creating user's alias that _get_user_info filled in,
because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id must be resolved here too: the spend update writer reads it
off this metadata to increment organization spend, so leaving it out silently
drops batch cost from org accounting for keys and teams that belong to one.
"""
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
request_tags = getattr(job, "request_tags", None)
metadata: dict[str, object] = {
"user_api_key_user_id": job.created_by,
"user_api_key": api_key,
"user_api_key_hash": api_key,
"user_api_key_team_id": team_id,
**(await self._get_user_info(batch_id, job.created_by)),
}
key_alias = await self._get_key_alias(batch_id, api_key, job.created_by)
if key_alias is not None:
metadata["user_api_key_alias"] = key_alias
team_alias = await self._get_team_alias(team_id)
if team_alias is not None:
metadata["user_api_key_team_alias"] = team_alias
org_id: Final = await self._get_org_id(job, batch_id)
if org_id is not None:
metadata["user_api_key_org_id"] = org_id
if isinstance(request_tags, list) and request_tags:
metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
return metadata
async def _cleanup_stale_managed_objects(self) -> None:
"""
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
@ -317,7 +155,7 @@ class CheckBatchCost:
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]":
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await _managed_object_table(self.prisma_client).find_many(
where={
@ -337,7 +175,7 @@ class CheckBatchCost:
order={"created_at": "asc"},
)
async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None:
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
@ -364,7 +202,7 @@ class CheckBatchCost:
"so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "_ManagedObjectRow") -> bool:
async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool:
"""
Atomically flip batch_processed from false to true, returning whether this pod won
the row. Every pod and uvicorn worker schedules its own poller against the shared
@ -390,7 +228,7 @@ class CheckBatchCost:
return False
return claimed > 0
async def _release_job_claim(self, job: "_ManagedObjectRow") -> None:
async def _release_job_claim(self, job: "ManagedObjectRow") -> None:
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
Safe to match on batch_processed=True: while this poller is active the retrieve
@ -411,7 +249,7 @@ class CheckBatchCost:
)
@staticmethod
def _has_unified_id_without_model(job: "_ManagedObjectRow") -> bool:
def _has_unified_id_without_model(job: "ManagedObjectRow") -> bool:
"""A unified id that decodes but carries no model_id can never be routed."""
from litellm.proxy.openai_files_endpoints.common_utils import (
convert_b64_uid_to_unified_uid,
@ -460,7 +298,7 @@ class CheckBatchCost:
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
async def _finalize_unbilled_terminal_job(
self, job: "_ManagedObjectRow", response: "LiteLLMBatch"
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
) -> None:
"""Persist a terminal batch that has nothing billable, converting any raw
provider file ids to managed ids, and take it out of the poll page."""
@ -505,7 +343,7 @@ class CheckBatchCost:
def _resolve_job_routing(
self,
job: "_ManagedObjectRow",
job: "ManagedObjectRow",
prom_logger: Optional["PrometheusLogger"],
) -> Optional[Tuple[str, str]]:
"""
@ -582,7 +420,7 @@ class CheckBatchCost:
def _resolve_unmanaged_provider_routing(
self,
job: "_ManagedObjectRow",
job: "ManagedObjectRow",
prom_logger: Optional["PrometheusLogger"],
llm_provider: str,
bare_model_name: str,
@ -678,7 +516,7 @@ class CheckBatchCost:
@classmethod
def _get_managed_file_model_name(
cls,
job: "_ManagedObjectRow",
job: "ManagedObjectRow",
deployment_info: "Deployment",
) -> Optional[str]:
"""
@ -698,7 +536,7 @@ class CheckBatchCost:
)
@staticmethod
def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
import json
from litellm.types.utils import LiteLLMBatch
@ -719,7 +557,7 @@ class CheckBatchCost:
async def _track_completed_batch_cost(
self,
job: "_ManagedObjectRow",
job: "ManagedObjectRow",
response: "LiteLLMBatch",
model_id: str,
batch_id: str,
@ -735,9 +573,9 @@ class CheckBatchCost:
"""
from litellm.batches.batch_utils import (
calculate_batch_cost_and_usage,
count_error_file_failed_requests,
get_file_content_as_dictionary,
calculate_batch_cost_and_usage,
)
from litellm.files.main import afile_content
from litellm.files.types import FileContentCallOptions
@ -911,7 +749,7 @@ class CheckBatchCost:
},
**({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
"metadata": {
**(await self._build_creator_attribution_metadata(job, batch_id)),
**(await self._creator_attribution.batch_metadata(job, batch_id)),
# spend logs read the deployment identity off these metadata keys, so
# without them the batch cost row carries no model_id or model_group
"model_info": {"id": model_id},

409
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -1,12 +1,36 @@
"""
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
Cost tracking is handled by the get-responses call, which prices normally only because the
poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
same route are non-inference and free.
A background response is billed exactly once: whichever client read or poll first sees it
finished moves its row out of a pending status, then logs the body that read already fetched as
one spend event stamped with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN and the creating key's
attribution, priced against the deployment that served it. The poll's status checks are stamped
with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
managed files hook, so both stay free. The spend event carries the id the client holds, the row's
unified_object_id, so its spend log is found under that id with the _background_cost suffix.
A read that finds no row, because the create could not store one, stores it under the reading key
and bills it the same way, so a response the create left untracked is still billed once. A streamed
background create is free on its own; its stream stores the row under the creating key and bills
the finished response through it before the client sees response.completed.
A read the caller retargets with client-side credentials, and a delete of a still-pending response,
settle the bill from a fetch through the configured deployment instead: the client controls what a
retargeted read returns, and a deleted response leaves the poll nothing to fetch.
A row stored before this billing model keeps the encrypted create id as its model_object_id. The
first lookup by that id rekeys it by the provider id, and the poll retires such a row without
billing it once a read has already tracked the response under the provider id.
"""
import uuid
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Dict, Final, Protocol, cast
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, cast
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
)
from pydantic import BaseModel, ConfigDict, TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -14,36 +38,124 @@ from litellm.constants import (
INTERNAL_CALL_ORIGIN_METADATA_KEY,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
PENDING_RESPONSE_STATUSES,
STALE_OBJECT_CLEANUP_BATCH_SIZE,
TERMINAL_RESPONSE_STATUSES,
)
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
from litellm.types.router import Deployment
from litellm.types.utils import (
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import Router
TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
class _ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@property
def unified_object_id(self) -> str: ...
@property
def created_by(self) -> str | None: ...
@property
def file_object(self) -> object: ...
@dataclass(frozen=True, slots=True)
class _PollTarget:
response_id: str
model_metadata: Mapping[str, object]
deployment: Deployment | None
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
@dataclass(frozen=True, slots=True)
class BilledBackgroundResponse:
response: ResponsesAPIResponse
model: str
custom_llm_provider: str | None
litellm_params: Mapping[str, object]
metadata: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class BilledCharge:
cost: float | None
logging_obj: "LiteLLMLogging"
BilledResponseLogger = Callable[[BilledBackgroundResponse], Awaitable[BilledCharge]]
UntrackedResponseTracker = Callable[[ResponsesAPIResponse], Awaitable[bool]]
def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
unpriced: Final = response.model_copy(deep=True)
unpriced.hidden_params = {key: value for key, value in response.hidden_params.items() if key != "response_cost"}
return unpriced
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj: Final = LiteLLMLogging(
model=billed.model,
messages=[{"role": "user", "content": "<retrieve_response>"}],
stream=False,
call_type=_RETRIEVE_RESPONSE_CALL_TYPE,
start_time=datetime.now(),
litellm_call_id=str(uuid.uuid4()),
function_id=str(uuid.uuid4()),
)
logging_obj.update_environment_variables(
litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
optional_params={},
custom_llm_provider=billed.custom_llm_provider,
)
unpriced: Final = _without_read_price(billed.response)
cost: Final = logging_obj.response_cost_calculator(result=unpriced)
end_time: Final = datetime.now()
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=logging_obj.async_success_handler(
result=unpriced, start_time=logging_obj.start_time, end_time=end_time
)
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=unpriced, start_time=logging_obj.start_time, end_time=end_time
)
return BilledCharge(cost=cost, logging_obj=logging_obj)
class _DeploymentBaseModel(BaseModel):
model_config = ConfigDict(frozen=True)
base_model: str | None = None
def _deployment_litellm_params(deployment: Deployment) -> Mapping[str, object]:
from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
api_base: Final = deployment.litellm_params.api_base
base_model: Final = _DeploymentBaseModel.model_validate(deployment.litellm_params, from_attributes=True).base_model
return MappingProxyType(
{
**({"api_base": mask_api_base_credentials(api_base)} if api_base else {}),
**({"base_model": base_model} if base_model else {}),
}
)
def _deployment_metadata(deployment: Deployment) -> Mapping[str, object]:
return MappingProxyType(
{
"model_info": deployment.model_info.model_dump(exclude_none=True),
"model_group": deployment.model_name,
"deployment": deployment.litellm_params.model,
}
)
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
return ManagedObjectRepository(prisma_client).table
@ -51,12 +163,17 @@ def _is_response_gone_at_provider(error: Exception, provider_response_id: str) -
return getattr(error, "status_code", None) == 404 and provider_response_id in str(error)
def _keyed_by_issued_id(job: ManagedObjectRow) -> bool:
return decrypted_response_id(job.model_object_id) != job.model_object_id
class CheckResponsesCost:
def __init__(
self,
proxy_logging_obj: "ProxyLogging",
prisma_client: "PrismaClient",
llm_router: "Router",
log_billed_response: BilledResponseLogger = log_billed_background_response,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -64,15 +181,17 @@ class CheckResponsesCost:
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
self._log_billed_response: Final = log_billed_response
def _resolve_deployment(self, response_id: str) -> bool:
def _resolve_deployment(self, response_id: str) -> Deployment | None:
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
return model_id is not None and self.llm_router.get_deployment(model_id=model_id) is not None
return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None
async def _get_response(
self,
response_id: str,
litellm_metadata: Dict[str, str],
litellm_metadata: dict[str, object],
via_router: bool,
) -> ResponsesAPIResponse:
"""Fetch the upstream response, using deployment credentials when available.
@ -91,9 +210,17 @@ class CheckResponsesCost:
)
return cast(ResponsesAPIResponse, router_response)
async def _expire_stale_rows(
self, cutoff: datetime, batch_size: int
) -> int:
async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse:
return await self._get_response(
response_id=target.response_id,
litellm_metadata={
**target.model_metadata,
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
},
via_router=target.deployment is not None,
)
async def _expire_stale_rows(self, cutoff: datetime, batch_size: int) -> int:
"""Execute the bounded UPDATE that marks stale rows as 'stale_expired'.
Isolated so it can be swapped / mocked in tests without touching the
@ -140,14 +267,154 @@ class CheckResponsesCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
def _target(self, response_id: str, model_name: str | None) -> _PollTarget:
return _PollTarget(
response_id=response_id,
model_metadata=MappingProxyType({"model": model_name, "model_group": model_name} if model_name else {}),
deployment=self._resolve_deployment(response_id),
)
def _poll_target(self, job: ManagedObjectRow) -> _PollTarget:
stored_model: Final = _STORED_RESPONSE_ADAPTER.validate_python(job.file_object).get("model")
return self._target(
decrypted_response_id(job.unified_object_id), stored_model if isinstance(stored_model, str) else None
)
async def _claim_for_billing(self, job: ManagedObjectRow) -> bool:
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
# bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
where={"id": job.id, "status": {"in": list(PENDING_RESPONSE_STATUSES)}},
data={"status": "completed"},
)
return claimed > 0
async def _release_billing_claim(self, job: ManagedObjectRow) -> None:
await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "status": "completed"},
data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
)
async def _billed_response(
self,
job: ManagedObjectRow,
target: _PollTarget,
response: ResponsesAPIResponse,
) -> BilledBackgroundResponse:
deployment: Final = target.deployment
return BilledBackgroundResponse(
response=response.model_copy(update={"id": job.unified_object_id}),
model=deployment.litellm_params.model if deployment is not None else (response.model or ""),
custom_llm_provider=ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
"custom_llm_provider"
),
litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
metadata=MappingProxyType(
{
**(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
**target.model_metadata,
**(_deployment_metadata(deployment) if deployment is not None else {}),
}
),
)
async def _bill_once(
self,
job: ManagedObjectRow,
target: _PollTarget,
response: ResponsesAPIResponse,
) -> BilledCharge | None:
if not await self._claim_for_billing(job):
verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
return None
try:
charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
except BaseException:
await self._release_billing_claim(job)
raise
verbose_proxy_logger.info(f"Response {job.unified_object_id} reached a terminal status and was billed")
return charge
async def _job_for_provider_id(self, provider_response_id: str) -> ManagedObjectRow | None:
return await _managed_object_table(self.prisma_client).find_first(
where={"model_object_id": provider_response_id, "file_purpose": "response"}
)
async def _rekey_by_provider_id(self, job: ManagedObjectRow) -> None:
await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id}, data={"model_object_id": decrypted_response_id(job.model_object_id)}
)
async def _tracked_job(self, response_id: str) -> ManagedObjectRow | None:
if tracked := await self._job_for_provider_id(decrypted_response_id(response_id)):
return tracked
stored_by_issued_id: Final = await _managed_object_table(self.prisma_client).find_first(
where={"unified_object_id": response_id, "file_purpose": "response"}
)
if stored_by_issued_id is not None and _keyed_by_issued_id(stored_by_issued_id):
await self._rekey_by_provider_id(stored_by_issued_id)
return stored_by_issued_id
async def _superseded_by_a_provider_id_row(self, job: ManagedObjectRow) -> bool:
if not _keyed_by_issued_id(job):
return False
if await self._job_for_provider_id(decrypted_response_id(job.model_object_id)) is None:
await self._rekey_by_provider_id(job)
return False
await self._claim_for_billing(job)
return True
async def _track_from_deployment(
self, response_id: str, track_untracked: UntrackedResponseTracker | None
) -> BilledCharge | None:
target: Final = self._target(decrypted_response_id(response_id), None)
if track_untracked is None or target.deployment is None:
return None
return await self.bill_if_finished(response_id, await self._fetch_unbilled(target), track_untracked)
async def _start_tracking(
self, response_id: str, response: ResponsesAPIResponse, track_untracked: UntrackedResponseTracker | None
) -> ManagedObjectRow | None:
if track_untracked is None or not await track_untracked(response):
return None
return await self._tracked_job(response_id)
async def bill_if_finished(
self,
response_id: str,
response: ResponsesAPIResponse,
track_untracked: UntrackedResponseTracker | None = None,
) -> BilledCharge | None:
if response.status not in TERMINAL_RESPONSE_STATUSES:
return None
job: Final = await self._tracked_job(response_id) or await self._start_tracking(
response_id, response, track_untracked
)
if job is None or job.status not in PENDING_RESPONSE_STATUSES:
return None
return await self._bill_once(job, self._poll_target(job), response)
async def bill_from_deployment(
self, response_id: str, track_untracked: UntrackedResponseTracker | None = None
) -> BilledCharge | None:
job: Final = await self._tracked_job(response_id)
if job is None:
return await self._track_from_deployment(response_id, track_untracked)
if job.status not in PENDING_RESPONSE_STATUSES:
return None
target: Final = self._poll_target(job)
response: Final = await self._fetch_unbilled(target)
if response.status not in TERMINAL_RESPONSE_STATUSES:
return None
return await self._bill_once(job, target, response)
async def check_responses_cost(self):
"""
Check if background responses are complete and track their cost.
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
- Query the provider to check if response is complete
- Cost is tracked by the get-responses call, billed because the poll is stamped
with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
- Mark responses in a terminal state as complete in the database
- Query the provider, unbilled, to check if the response is complete
- Claim each terminal job by moving its row to completed, then log the body that check fetched
as its one billed event, so two pollers racing over the same job bill it once
- Mark responses the provider no longer has (404 through a resolved
deployment) as stale_expired
"""
@ -160,99 +427,53 @@ class CheckResponsesCost:
jobs = await _managed_object_table(self.prisma_client).find_many(
where={
"status": {"in": ["queued", "in_progress"]},
# bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
"status": {"in": list(PENDING_RESPONSE_STATUSES)},
"file_purpose": "response",
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
completed_jobs: Final[list[_ManagedObjectRow]] = []
expired_jobs: Final[list[_ManagedObjectRow]] = []
expired_jobs: Final[list[ManagedObjectRow]] = []
for job in jobs:
unified_object_id = job.unified_object_id
try:
from litellm.proxy.hooks.responses_id_security import (
ResponsesIDSecurity,
)
# Get the stored response object to extract model information
stored_response = job.file_object
model_name = stored_response.get("model", None)
# Decrypt the response ID
responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
# Prepare metadata with model information for cost tracking
litellm_metadata = {
"user_api_key_user_id": job.created_by or "default-user-id",
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
}
# Add model information if available
if model_name:
litellm_metadata["model"] = model_name
litellm_metadata["model_group"] = model_name # Use same value for model_group
via_router = self._resolve_deployment(responses_id_security)
if await self._superseded_by_a_provider_id_row(job):
verbose_proxy_logger.info(f"Response {unified_object_id} is billed under its provider id row")
continue
target = self._poll_target(job)
except Exception as e:
verbose_proxy_logger.warning(
f"Skipping job {unified_object_id} due to error: {e}"
)
verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
continue
provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(
responses_id_security
).get("response_id", responses_id_security)
provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
"response_id", target.response_id
)
try:
response = await self._get_response(
response_id=responses_id_security,
litellm_metadata=litellm_metadata,
via_router=via_router,
)
response = await self._fetch_unbilled(target)
verbose_proxy_logger.debug(
f"Response {unified_object_id} status: {response.status}, model: {model_name}"
)
verbose_proxy_logger.debug(f"Response {unified_object_id} status: {response.status}")
if response.status in TERMINAL_RESPONSE_STATUSES:
await self._bill_once(job, target, response)
except Exception as e:
if via_router and _is_response_gone_at_provider(e, provider_response_id):
if target.deployment is not None and _is_response_gone_at_provider(e, provider_response_id):
verbose_proxy_logger.info(
f"Response {unified_object_id} no longer available at provider (404), marking stale_expired: {e}"
)
expired_jobs.append(job)
continue
verbose_proxy_logger.warning(
f"Skipping job {unified_object_id} due to error: {e}"
)
verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
continue
if response.status in TERMINAL_RESPONSE_STATUSES:
verbose_proxy_logger.info(
f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
)
completed_jobs.append(job)
# Mark completed jobs in the database
if len(completed_jobs) > 0:
await _managed_object_table(self.prisma_client).update_many(
# bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
where={"id": {"in": [job.id for job in completed_jobs]}},
data={"status": "completed"},
)
verbose_proxy_logger.info(
f"Marked {len(completed_jobs)} response jobs as completed"
)
if len(expired_jobs) > 0:
await _managed_object_table(self.prisma_client).update_many(
# bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
where={"id": {"in": [job.id for job in expired_jobs]}},
data={"status": "stale_expired"},
)
verbose_proxy_logger.info(
f"Marked {len(expired_jobs)} response jobs as stale_expired"
)
verbose_proxy_logger.info(f"Marked {len(expired_jobs)} response jobs as stale_expired")

218
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py Normal file
View file

@ -0,0 +1,218 @@
"""Creator attribution for the spend row a managed-object cost poll writes.
A batch or a background response is billed by a poll job that runs long after the creating
request returned, so the managed object row's creator columns are the only record of which
key, team, organization, user, end user, and tags carry the spend.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol, cast
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import CLI_SESSION_KEY_PREFIX
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.proxy.utils import PrismaClient
from litellm.repositories.prisma_protocols import TableActions
class ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@property
def unified_object_id(self) -> str: ...
@property
def model_object_id(self) -> str: ...
@property
def created_by(self) -> str | None: ...
@property
def status(self) -> str | None: ...
@property
def file_object(self) -> object: ...
@property
def org_id(self) -> str | None: ...
@property
def api_key(self) -> str | None: ...
@property
def team_id(self) -> str | None: ...
@property
def request_tags(self) -> object: ...
@property
def end_user_id(self) -> str | None: ...
_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
def _string_tags(request_tags: object) -> list[str] | None:
if not isinstance(request_tags, list):
return None
return [tag for tag in _TAG_LIST_ADAPTER.validate_python(request_tags) if isinstance(tag, str)] or None
def _creator_user_id(created_by: str | None) -> str | None:
if created_by is None or created_by.startswith(_KEY_OWNER_SCOPE_PREFIX):
return None
return created_by
def _user_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_UserTable]:
return UserRepository(prisma_client).table
def _token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_VerificationToken]:
return VerificationTokenRepository(prisma_client).table
def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_TeamTable]:
return TeamRepository(prisma_client).table
@dataclass(frozen=True, slots=True)
class _CreatorRows:
key: prisma_models.LiteLLM_VerificationToken | None
team: prisma_models.LiteLLM_TeamTable | None
user: prisma_models.LiteLLM_UserTable | None
class ManagedObjectCreatorAttribution:
"""Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
job's spend row is attributed like the creating request: key hash, team, organization,
user, aliases, and tags.
Rows created before api_key and request_tags were persisted carry only created_by and
team_id, and fall back to those. A named creating key owns user_api_key_alias; when it
has no alias, or the key has since been rotated or deleted, the field keeps the creating
user's alias, because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
to increment organization spend, so leaving it out drops the cost from org accounting.
A background response also carries its end user and the key, team, and user budgets, the
way the read that billed it before did, and never names a key-scoped creator as its user.
"""
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
self._prisma_client: Final = prisma_client
self._poller_name: Final = poller_name
async def batch_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
rows: Final = await self._rows(job, object_id)
return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
rows: Final = await self._rows(job, object_id)
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
**_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
**_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
**_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
}
async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
return _CreatorRows(
key=await self._key_row(object_id, getattr(job, "api_key", None), job.created_by),
team=await self._team_row(object_id, getattr(job, "team_id", None)),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
async def _user_row(self, object_id: str, user_id: str | None) -> prisma_models.LiteLLM_UserTable | None:
"""Objects created by a team or service account key carry no user id, and
find_unique(where={"user_id": None}) raises, so those resolve no row."""
if not user_id:
return None
try:
return await _user_table(self._prisma_client).find_unique(where={"user_id": user_id})
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up user {user_id} for {object_id}: {e}")
return None
async def _key_row(
self, object_id: str, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
try:
return await _token_table(self._prisma_client).find_unique(where={"token": api_key})
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
return None
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
return None
try:
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
return None
def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
key_alias: Final = _key_alias(api_key, job.created_by, rows.key)
team_alias: Final = getattr(rows.team, "team_alias", None)
org_id: Final = getattr(job, "org_id", None) or _organization_id(rows.key) or _organization_id(rows.team)
tags: Final = _string_tags(getattr(job, "request_tags", None))
return {
"user_api_key": api_key,
"user_api_key_hash": api_key,
"user_api_key_team_id": team_id,
**(
{
"user_api_key_user_email": getattr(rows.user, "user_email", None),
"user_api_key_alias": getattr(rows.user, "user_alias", None),
}
if rows.user is not None
else {}
),
**({"user_api_key_alias": key_alias} if key_alias is not None else {}),
**({"user_api_key_team_alias": team_alias} if team_alias is not None else {}),
**({"user_api_key_org_id": org_id} if org_id is not None else {}),
**({"tags": tags} if tags else {}),
}
def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"
def _key_alias(
api_key: str | None, created_by: str | None, key_row: prisma_models.LiteLLM_VerificationToken | None
) -> str | None:
if api_key and _is_cli_session_key(api_key, created_by):
return api_key
return getattr(key_row, "key_alias", None) if key_row is not None else None
def _organization_id(row: object) -> str | None:
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
if row is None:
return {}
return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}

153
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -24,7 +24,9 @@ from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import BaseModel, ValidationError
from openai import APIConnectionError, AsyncOpenAI
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, Unpack
import litellm
@ -35,17 +37,17 @@ from litellm.caching.caching import DualCache
from litellm.constants import (
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
MAX_FILE_LIST_LIMIT,
PENDING_RESPONSE_STATUSES,
)
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.litellm_core_utils.internal_call_metadata import requests_background
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_metadata,
)
from openai import APIConnectionError, AsyncOpenAI
from openai.types.file_deleted import FileDeleted
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.managed_resources.isolation import (
@ -88,6 +90,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
request_tags_from_metadata,
)
from litellm.repositories.managed_file_repository import ManagedFileRepository
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue]
AllMessageValues,
AsyncCursorPage,
@ -100,6 +103,7 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess
ResponsesAPIResponse,
)
from litellm.types.utils import (
ADDRESSED_RESPONSE_ID_FIELD,
CallTypesLiteral,
LiteLLMBatch,
LiteLLMFineTuningJob,
@ -346,6 +350,8 @@ class _ManagedObjectTableActions(Protocol):
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
async def create(self, data: Mapping[str, object]) -> "PrismaManagedObjectRow": ...
class _ManagedResourceDatabase(Protocol):
@property
@ -392,6 +398,39 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
return prisma_client.db.litellm_managedobjecttable
_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
def _storable_text(text: str) -> str:
return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "")
def _storable_json_value(value: JsonValue) -> JsonValue:
if isinstance(value, str):
return _storable_text(value)
if isinstance(value, list):
return [_storable_json_value(item) for item in value]
if isinstance(value, dict):
return {_storable_text(key): _storable_json_value(item) for key, item in value.items()}
return value
def _postgres_storable_json(file_object: BaseModel) -> str:
dumped: Final = file_object.model_dump_json()
if "\\u0000" not in dumped:
return dumped
return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode()
def _without_model_output(
file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse,
) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse:
return (
file_object.model_copy(update={"output": []}) if isinstance(file_object, ResponsesAPIResponse) else file_object
)
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
@ -405,6 +444,40 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
)
_MANAGED_RESPONSE_READ_CALL_TYPES: Final = frozenset({"aget_responses", "acancel_responses"})
_REQUEST_FIELDS: Final = TypeAdapter(dict[str, object])
def _addresses_a_proxy_issued_response(data: Mapping[str, object]) -> bool:
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
addressed_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) or data.get("response_id")
return (
isinstance(addressed_id, str)
and ResponsesAPIRequestUtils.get_model_id_from_response_id(decrypted_response_id(addressed_id)) is not None
)
def _is_streamed_background_create(request: Mapping[str, object]) -> bool:
return (
request.get("stream") is True
and isinstance(request.get("litellm_metadata"), dict)
and requests_background(request)
)
def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str, object]:
metadata_field: Final = "litellm_metadata" if isinstance(request.get("litellm_metadata"), dict) else "metadata"
existing_metadata: Final = request.get(metadata_field)
return {
**request,
metadata_field: {
**(_REQUEST_FIELDS.validate_python(existing_metadata) if isinstance(existing_metadata, dict) else {}),
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True,
},
}
_MANAGED_FILES_TARGET: Final = "managed_files"
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
@ -529,11 +602,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
) -> None:
"""Persist a managed object row, caching it and upserting it in the DB.
persist_attribution is set only by the batch create, which is the one caller
that can speak for the creator; it gates the api_key and request_tags columns
that CheckBatchCost bills against, so a later poll or retrieve of the same
batch cannot record itself as the paying key. Like created_by and team_id,
both are written only in the upsert create branch, never on update.
persist_attribution is set only by the batch and background response creates,
the callers that can speak for the creator; it gates the api_key, request_tags, and,
for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or
retrieve of the same object cannot record itself as the paying key. Like created_by and team_id,
they are written only in the upsert create branch, never on update.
A stored response keeps no model output: billing reads only its status and model from
the row, so the row never holds generated content.
create_if_missing is cleared by callers that observe a batch they did not
create, such as a poll. They still refresh status and file_object, but a
@ -546,11 +622,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
It is written only in the upsert create branch.
"""
verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
stored_object: Final = _without_model_output(file_object)
litellm_managed_object = LiteLLM_ManagedObjectTable(
unified_object_id=unified_object_id,
model_object_id=model_object_id,
file_purpose=file_purpose,
file_object=file_object,
file_object=stored_object,
)
await self.internal_usage_cache.async_set_cache(
key=unified_object_id,
@ -561,32 +638,35 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
from prisma import Json
api_key = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
end_user_id: Final = user_api_key_dict.end_user_id if file_purpose == "response" else None
attribution_columns = (
{
**({"api_key": api_key} if api_key is not None else {}),
**({"request_tags": Json(list(request_tags))} if request_tags else {}),
**({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
**({"end_user_id": end_user_id} if end_user_id else {}),
}
if persist_attribution
else {}
)
# FIX: Update status and file_object on every operation to keep state in sync
update_columns: Final = {
"file_object": file_object.model_dump_json(),
"file_object": _postgres_storable_json(stored_object),
"status": file_object.status,
"updated_by": user_api_key_dict.user_id,
}
managed_object_table: Final = _managed_object_table(self.prisma_client)
if not create_if_missing:
await _managed_object_table(self.prisma_client).update_many(
await managed_object_table.update_many(
where={"unified_object_id": unified_object_id},
data=update_columns,
)
return
await _managed_object_table(self.prisma_client).upsert(
await managed_object_table.upsert(
where={"unified_object_id": unified_object_id},
data={
"create": {
"unified_object_id": unified_object_id,
"file_object": file_object.model_dump_json(),
"file_object": _postgres_storable_json(stored_object),
"model_object_id": model_object_id,
"file_purpose": file_purpose,
"created_by": resolve_resource_owner_id(user_api_key_dict),
@ -601,6 +681,42 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
},
)
async def track_untracked_background_response(
self,
response_id: str,
response: ResponsesAPIResponse,
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None = None,
) -> bool:
if not _addresses_a_proxy_issued_response({"response_id": response_id}):
return False
from prisma import Json
from prisma.errors import UniqueViolationError
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
try:
await _managed_object_table(self.prisma_client).create(
data={
"unified_object_id": response_id,
"model_object_id": decrypted_response_id(response_id),
"file_purpose": "response",
"file_object": _postgres_storable_json(_without_model_output(response)),
"status": PENDING_RESPONSE_STATUSES[0],
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"org_id": await self._resolve_creator_org_id(user_api_key_dict),
"updated_by": user_api_key_dict.user_id,
**({"api_key": api_key} if api_key is not None else {}),
**({"end_user_id": user_api_key_dict.end_user_id} if user_api_key_dict.end_user_id else {}),
**({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
}
)
except UniqueViolationError:
return False
return True
@with_service_target(_MANAGED_FILES_TARGET)
async def get_unified_file_id(
self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
@ -1375,6 +1491,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools)
if unified_vs_ids:
await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict)
create_request: Final = _REQUEST_FIELDS.validate_python(data)
if _is_streamed_background_create(create_request):
return _with_managed_background_marker(create_request)
elif call_type == CallTypes.afile_content.value:
retrieve_file_id = cast(Optional[str], data.get("file_id"))
potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
@ -1432,6 +1551,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id)
elif retrieve_object_id and accessor_key == "batch_id":
await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict)
elif call_type in _MANAGED_RESPONSE_READ_CALL_TYPES:
request: Final = _REQUEST_FIELDS.validate_python(data)
if _addresses_a_proxy_issued_response(request):
return _with_managed_background_marker(request)
elif call_type == CallTypes.acreate_fine_tuning_job.value:
fine_tuning_file_ids: Final = tuple(
file_id

4
litellm-proxy-extras/litellm_proxy_extras/migrations/20261015000000_managed_object_end_user_id/migration.sql Normal file
View file

@ -0,0 +1,4 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "end_user_id" TEXT;

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
end_user_id String?
updated_at DateTime @updatedAt
updated_by String?

3
litellm/constants.py
View file

@ -1622,6 +1622,7 @@ CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
@ -1862,6 +1863,8 @@ RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
PROXY_BATCH_POLLING_INTERVAL: Final = get_env_int("PROXY_BATCH_POLLING_INTERVAL", 3600)
MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, get_env_int("MAX_OBJECTS_PER_POLL_CYCLE", 50))
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
PENDING_RESPONSE_STATUSES: Final = ("queued", "in_progress")
TERMINAL_RESPONSE_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "incomplete"})
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60
STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, get_env_int("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))

63
litellm/litellm_core_utils/internal_call_metadata.py
View file

@ -21,9 +21,17 @@ from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES
from litellm.constants import (
INTERNAL_CALL_ORIGIN_METADATA_KEY,
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
NON_INFERENCE_CALL_TYPES,
)
from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin
from litellm.types.utils import (
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
InternalCallOrigin,
)
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
@ -62,10 +70,32 @@ def is_background_response(response: object) -> bool:
because the callers hold it in both shapes.
"""
if isinstance(response, Mapping):
return response.get("background") is True
return "background" in response and response["background"] is True
return getattr(response, "background", None) is True
def requests_background(request_body: Mapping[str, object]) -> bool:
return bool(request_body.get("background")) or is_background_response(request_body.get("extra_body"))
_RESPONSE_CREATE_CALL_TYPES: Final = frozenset({"responses", "aresponses"})
def _finished_response(response: object) -> object:
wrapped: Final = response.get("response") if isinstance(response, Mapping) else getattr(response, "response", None)
return response if wrapped is None else wrapped
def _may_be_unbilled(call_type: str | None, response: object) -> bool:
if call_type in _RESPONSE_CREATE_CALL_TYPES:
return is_background_response(_finished_response(response))
return call_type in NON_INFERENCE_CALL_TYPES
def _is_managed_background_call(metadata: Mapping[str, object] | None) -> bool:
return metadata is not None and metadata.get(MANAGED_BACKGROUND_RESPONSE_METADATA_KEY) is True
def is_unbilled_non_inference_call(
call_type: str | None,
metadata: Mapping[str, object] | None,
@ -74,18 +104,25 @@ def is_unbilled_non_inference_call(
"""A read/management route priced at zero, because the usage it reports belongs to the
call that created the object it just read.
Retrieving a background response is the exception, and the enterprise cost poller's read
is the same exception seen from the other side: that job's create billed nothing, so its
retrieval is the only place the spend is ever visible. Pricing those at zero would lose
the spend rather than deduplicate it.
A background response is the exception, since its create billed nothing. When the proxy
tracks it as a managed object, the read is marked with
``MANAGED_BACKGROUND_RESPONSE_METADATA_KEY`` and stays free, because the one billing event
the proxy logs from the finished response, stamped with
``BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN``, bills the job. A streamed background create
the proxy tracks carries the same marker and is free for the same reason. Any other
read of a background response (the SDK, a proxy without a database) is the only place its
spend is ever visible, so it prices normally.
"""
if call_type not in NON_INFERENCE_CALL_TYPES:
if not _may_be_unbilled(call_type, response):
return False
if is_background_response(response):
if call_type in _RESPONSE_CREATE_CALL_TYPES:
return _is_managed_background_call(metadata)
origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) if metadata is not None else None
if origin == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN:
return False
if metadata is None:
if not is_background_response(response):
return True
return metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
return origin == BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN or _is_managed_background_call(metadata)
def is_unbilled_non_inference_call_from_params(
@ -95,10 +132,10 @@ def is_unbilled_non_inference_call_from_params(
) -> bool:
""":func:`is_unbilled_non_inference_call` for callers holding raw ``litellm_params``.
The call-type membership test runs first so that inference traffic, which is every
The call-type and background tests run first so that inference traffic, which is every
request in a normal workload, never pays for the metadata merge behind it.
"""
if call_type not in NON_INFERENCE_CALL_TYPES:
if not _may_be_unbilled(call_type, response):
return False
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup

4
litellm/litellm_core_utils/litellm_logging.py
View file

@ -5856,7 +5856,7 @@ class StandardLoggingPayloadSetup:
return messages
@staticmethod
def merge_litellm_metadata(litellm_params: Mapping[str, object]) -> dict:
def merge_litellm_metadata(litellm_params: Mapping[str, object]) -> dict[str, object]:
"""
Merge both litellm_metadata and metadata from litellm_params.
@ -6483,7 +6483,7 @@ class StandardLoggingPayloadSetup:
@staticmethod
def get_request_tags(
litellm_params: dict[str, object],
proxy_server_request: dict[str, object],
proxy_server_request: Mapping[str, object],
) -> list[str]:
metadata: Final = cast( # cast-ok: request metadata is caller-provided
Mapping[str, object], litellm_params.get("metadata") or {}

9
litellm/proxy/auth/auth_utils.py
View file

@ -674,6 +674,15 @@ def is_request_body_safe(
return True
def carries_client_side_overrides(request_body: object) -> bool:
if not isinstance(request_body, dict):
return False
try:
return not is_request_body_safe(request_body=request_body, general_settings={}, llm_router=None, model="")
except ValueError:
return True
def _coerce_metadata_to_dict(value: object) -> dict[str, object] | None:
"""Return ``value`` as a dict, parsing it from JSON if delivered as a string.

13
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -12,7 +12,6 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -68,7 +67,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 #
update_batch_in_database,
validate_managed_id_requirement,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_metadata
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_request_data
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy, is_known_model
from litellm.repositories.managed_batch_repository import ManagedBatchRepository
@ -81,7 +80,6 @@ if TYPE_CHECKING:
from prisma.models import LiteLLM_ManagedObjectTable
router: Final = APIRouter()
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _hidden_param_string(hidden_params: Mapping[str, object], key: str) -> str:
@ -89,13 +87,6 @@ def _hidden_param_string(hidden_params: Mapping[str, object], key: str) -> str:
return value if isinstance(value, str) else ""
def _request_tags(data: Mapping[str, object]) -> tuple[str, ...] | None:
metadata: Final = data.get("litellm_metadata")
if metadata is None:
return None
return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata))
def _require_batch_response(response: LLMResponseTypes) -> LiteLLMBatch:
if not isinstance(response, LiteLLMBatch):
raise TypeError("Batch endpoint received a non-batch response")
@ -444,7 +435,7 @@ async def create_batch(
model=model,
provider=executed_provider,
user_api_key_dict=user_api_key_dict,
request_tags=_request_tags(_create_batch_data),
request_tags=request_tags_from_request_data(_create_batch_data),
)
if executed_provider is not None
else await _create_provider_batch_for_managed_file(

69
litellm/proxy/common_request_processing.py
View file

@ -1605,6 +1605,31 @@ def _get_cost_breakdown_from_logging_obj(
)
def _header_value(value: float | None) -> str | None:
return str(value) if value is not None else None
def cost_headers(
response_cost: float | str | None, litellm_logging_obj: LiteLLMLoggingObj | None
) -> Mapping[str, str | None]:
cost_breakdown: Final = _get_cost_breakdown_from_logging_obj(
litellm_logging_obj=litellm_logging_obj, response_cost=response_cost
)
return {
"x-litellm-response-cost": str(response_cost),
"x-litellm-response-cost-original": _header_value(cost_breakdown.original_cost),
"x-litellm-response-cost-discount-amount": _header_value(cost_breakdown.discount_amount),
"x-litellm-response-cost-margin-amount": _header_value(cost_breakdown.margin_total_amount),
"x-litellm-response-cost-margin-percent": _header_value(cost_breakdown.margin_percent),
"x-litellm-response-cost-input": _header_value(cost_breakdown.input_cost),
"x-litellm-response-cost-output": _header_value(cost_breakdown.output_cost),
"x-litellm-response-cost-cache-read": _header_value(cost_breakdown.cache_read_cost),
"x-litellm-response-cost-cache-creation": _header_value(cost_breakdown.cache_creation_cost),
"x-litellm-response-cost-reasoning": _header_value(cost_breakdown.reasoning_cost),
"x-litellm-response-cost-tool-usage": _header_value(cost_breakdown.tool_usage_cost),
}
def _classifier_cost_from_request_data(request_data: Mapping[str, object] | None) -> float | None:
"""Cost of the auto-router's LLM classifier call, read from the request's routing_decision.
@ -1878,10 +1903,6 @@ class ProxyBaseLLMRequestProcessing:
use_logging_obj=read_timing_from_logging_obj,
)
cost_breakdown: Final = _get_cost_breakdown_from_logging_obj(
litellm_logging_obj=litellm_logging_obj, response_cost=response_cost
)
# Calculate updated spend for header (include current response_cost)
current_spend: Final = user_api_key_dict.spend or 0.0
updated_spend = current_spend
@ -1908,37 +1929,7 @@ class ProxyBaseLLMRequestProcessing:
), # don't include query params, risk of leaking sensitive info
"x-litellm-version": version,
"x-litellm-model-region": model_region,
"x-litellm-response-cost": str(response_cost),
"x-litellm-response-cost-original": (
str(cost_breakdown.original_cost) if cost_breakdown.original_cost is not None else None
),
"x-litellm-response-cost-discount-amount": (
str(cost_breakdown.discount_amount) if cost_breakdown.discount_amount is not None else None
),
"x-litellm-response-cost-margin-amount": (
str(cost_breakdown.margin_total_amount) if cost_breakdown.margin_total_amount is not None else None
),
"x-litellm-response-cost-margin-percent": (
str(cost_breakdown.margin_percent) if cost_breakdown.margin_percent is not None else None
),
"x-litellm-response-cost-input": (
str(cost_breakdown.input_cost) if cost_breakdown.input_cost is not None else None
),
"x-litellm-response-cost-output": (
str(cost_breakdown.output_cost) if cost_breakdown.output_cost is not None else None
),
"x-litellm-response-cost-cache-read": (
str(cost_breakdown.cache_read_cost) if cost_breakdown.cache_read_cost is not None else None
),
"x-litellm-response-cost-cache-creation": (
str(cost_breakdown.cache_creation_cost) if cost_breakdown.cache_creation_cost is not None else None
),
"x-litellm-response-cost-reasoning": (
str(cost_breakdown.reasoning_cost) if cost_breakdown.reasoning_cost is not None else None
),
"x-litellm-response-cost-tool-usage": (
str(cost_breakdown.tool_usage_cost) if cost_breakdown.tool_usage_cost is not None else None
),
**cost_headers(response_cost, litellm_logging_obj),
"x-litellm-classifier-cost": (str(classifier_cost) if classifier_cost is not None else None),
"x-litellm-key-tpm-limit": str(user_api_key_dict.tpm_limit),
"x-litellm-key-rpm-limit": str(user_api_key_dict.rpm_limit),
@ -3164,7 +3155,7 @@ class ProxyBaseLLMRequestProcessing:
)
@staticmethod
def _extract_completed_responses_response(stream_response: object) -> object:
def extract_completed_responses_response(stream_response: object) -> object:
"""Pull the assembled ``ResponsesAPIResponse`` off a streaming iterator.
``ResponsesAPIStreamingIterator`` stores the terminal stream event
@ -3200,9 +3191,7 @@ class ProxyBaseLLMRequestProcessing:
appears, ahead of the chunk carrying ``response.completed``.
"""
async for chunk in wrapped_generator:
completed_obj = ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
original_stream_response
)
completed_obj = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(original_stream_response)
if completed_obj is None:
yield chunk
continue
@ -3214,7 +3203,7 @@ class ProxyBaseLLMRequestProcessing:
async for remaining_chunk in wrapped_generator:
yield remaining_chunk
return
late_completed_obj: Final = ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
late_completed_obj: Final = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(
original_stream_response
)
if late_completed_obj is not None:

21
litellm/proxy/db/model_usage_rollup.py
View file

@ -19,6 +19,7 @@ from litellm.constants import (
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES, SpendLogsPayload
from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks
from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL, rollup_lock_timeout_setting
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
@ -59,6 +60,7 @@ class ModelUsageTransaction:
prompt_tokens: int
completion_tokens: int
successful: bool
counts_as_request: bool = True
def model_usage_task_type(request_tags: str) -> str:
@ -78,12 +80,12 @@ def model_usage_task_type(request_tags: str) -> str:
)
def _is_internal_call(metadata: str) -> bool:
def _internal_call_origin(metadata: str) -> object:
try:
decoded: Final = _METADATA.validate_json(metadata)
except ValidationError:
return False
return bool(decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY))
return None
return decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)
def _date_from_start_time(start_time: datetime | str) -> str | None:
@ -94,7 +96,8 @@ def _date_from_start_time(start_time: datetime | str) -> str | None:
def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransaction | None:
date: Final = _date_from_start_time(payload["startTime"])
if date is None or _is_internal_call(payload["metadata"]):
origin: Final = _internal_call_origin(payload["metadata"])
if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN):
return None
model: Final = payload["model"] or "unknown"
return ModelUsageTransaction(
@ -109,6 +112,7 @@ def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransa
prompt_tokens=payload["prompt_tokens"],
completion_tokens=payload["completion_tokens"],
successful=payload["status"] == "success",
counts_as_request=not origin,
)
@ -142,8 +146,9 @@ async def flush_model_usage_transactions(
spend = sum(entry.spend for entry in entries)
prompt_tokens = sum(entry.prompt_tokens for entry in entries)
completion_tokens = sum(entry.completion_tokens for entry in entries)
successful = sum(1 for entry in entries if entry.successful)
failed = len(entries) - successful
requests = tuple(entry for entry in entries if entry.counts_as_request)
successful = sum(1 for entry in requests if entry.successful)
failed = len(requests) - successful
key_fields = {
"date": key.date,
"model_group": key.model_group,
@ -159,7 +164,7 @@ async def flush_model_usage_transactions(
"spend": spend,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"request_count": len(entries),
"request_count": len(requests),
"successful_requests": successful,
"failed_requests": failed,
},
@ -167,7 +172,7 @@ async def flush_model_usage_transactions(
"spend": {"increment": spend},
"prompt_tokens": {"increment": prompt_tokens},
"completion_tokens": {"increment": completion_tokens},
"request_count": {"increment": len(entries)},
"request_count": {"increment": len(requests)},
"successful_requests": {"increment": successful},
"failed_requests": {"increment": failed},
},

10
litellm/proxy/hooks/responses_id_security.py
View file

@ -140,12 +140,12 @@ class ResponsesIDSecurity(CustomLogger):
)
if not isinstance(addressed_id, str) or not addressed_id:
return data
authorized_id: Final = self._authorize_response_id(addressed_id, user_api_key_dict)
authorized_id: Final = self.authorize_response_id(addressed_id, user_api_key_dict)
data[addressed_id_field] = authorized_id
data[ADDRESSED_RESPONSE_ID_FIELD] = addressed_id
return data
def _authorize_response_id(
def authorize_response_id(
self,
response_id: str,
user_api_key_dict: "UserAPIKeyAuth",
@ -344,3 +344,9 @@ class ResponsesIDSecurity(CustomLogger):
):
chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache)
yield chunk
def decrypted_response_id(response_id: str) -> str:
"""The id a client-facing response id stands for, the same for every encryption of one response."""
original_response_id, _, _ = ResponsesIDSecurity()._decrypt_response_id(response_id)
return original_response_id

2
litellm/proxy/litellm_pre_call_utils.py
View file

@ -31,6 +31,7 @@ from litellm.constants import (
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
OTEL_SERVICE_NAME_METADATA_KEYS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
@ -378,6 +379,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
CONSUMED_REQUEST_TAGS_METADATA_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
"standard_logging_object",
"litellm_roi_estimator",
"proxy_server_request",

13
litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py
View file

@ -9,6 +9,8 @@ import asyncio
from collections.abc import Mapping, Sequence
from typing import Final
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import strip_null_bytes
@ -45,6 +47,17 @@ def request_tags_from_metadata(request_metadata: Mapping[str, object]) -> tuple[
return None
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def request_tags_from_request_data(data: Mapping[str, object]) -> tuple[str, ...] | None:
"""Tags of a request whose metadata travels as litellm_metadata (batches, responses)"""
metadata: Final = data.get("litellm_metadata")
if metadata is None:
return None
return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata))
def log_batch_registration_result(
finished: asyncio.Task[None],
provider: str,

390
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -1,11 +1,23 @@
import asyncio
import contextlib
import json
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Mapping, Sequence
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from enum import Enum
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeAlias, cast, get_args
from typing import (
TYPE_CHECKING,
Any,
Final,
Literal,
NamedTuple,
Protocol,
TypeAlias,
cast,
get_args,
runtime_checkable,
)
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, Response
@ -17,25 +29,32 @@ from starlette.websockets import WebSocket, WebSocketDisconnect
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import EMPTY_MAPPING
from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
build_blocked_response,
)
from litellm.proxy._types import *
from litellm.proxy.auth.auth_utils import carries_client_side_overrides
from litellm.proxy.auth.user_api_key_auth import (
UserAPIKeyAuth,
user_api_key_auth,
user_api_key_auth_websocket,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, create_response
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, cost_headers, create_response
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_safe_set_request_parsed_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
read_request_body,
safe_set_request_parsed_body,
)
from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
request_tags_from_request_data,
)
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import (
@ -47,10 +66,13 @@ from litellm.types.responses.main import DeleteResponseResult
from litellm.types.utils import TokenCountResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
router: Final = APIRouter()
_RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None)
_PROXY_SERVER_REQUEST_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg
@ -344,7 +366,7 @@ async def responses_api(
asyncio.create_task(
background_streaming_task(
polling_id=polling_id,
data=data.copy(),
data=dict(_without_managed_background_marker(data)),
polling_handler=polling_handler,
request=request,
fastapi_response=fastapi_response,
@ -379,7 +401,11 @@ async def responses_api(
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=native_data_generator,
select_data_generator=(
partial(_streamed_background_data_generator, native_data_generator, proxy_logging_obj)
if requests_background(_PROXY_SERVER_REQUEST_ADAPTER.validate_python(data))
else native_data_generator
),
model=None,
user_model=user_model,
user_temperature=user_temperature,
@ -389,49 +415,38 @@ async def responses_api(
version=version,
)
# Store in managed objects table if background mode is enabled
if data.get("background") and isinstance(response, ResponsesAPIResponse):
if response.status in ["queued", "in_progress"]:
from litellm_enterprise.proxy.hooks.managed_files import (
PROXY_LiteLLMManagedFiles,
)
if isinstance(response, ResponsesAPIResponse) and (data.get("background") or is_background_response(response)):
managed_files_obj: Final = _managed_files_hook(proxy_logging_obj)
managed_files_obj: Final = cast(
PROXY_LiteLLMManagedFiles | None,
proxy_logging_obj.get_proxy_hook("managed_files"),
)
if managed_files_obj and llm_router:
try:
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
model_id: Final = hidden_params.get("model_id", None)
if managed_files_obj and llm_router:
try:
# Get the actual deployment model_id from hidden params
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
model_id: Final = hidden_params.get("model_id", None)
if not model_id:
verbose_proxy_logger.warning(
"No model_id found in response hidden params for response %s, skipping managed object storage",
response.id,
)
raise Exception("No model_id found in response hidden params")
# Store in managed objects table
await managed_files_obj.store_unified_object_id(
unified_object_id=response.id,
file_object=response,
litellm_parent_otel_span=None,
model_object_id=response.id,
file_purpose="response",
user_api_key_dict=user_api_key_dict,
)
verbose_proxy_logger.info(
"Stored background response %s in managed objects table with unified_id=%s",
response.id,
if not model_id:
verbose_proxy_logger.warning(
"No model_id found in response hidden params for response %s, skipping managed object storage",
response.id,
)
except Exception as e:
verbose_proxy_logger.error(
"Failed to store background response in managed objects table: %s", e
)
raise Exception("No model_id found in response hidden params")
await managed_files_obj.store_unified_object_id(
unified_object_id=response.id,
file_object=response,
litellm_parent_otel_span=None,
model_object_id=decrypted_response_id(response.id),
file_purpose="response",
user_api_key_dict=user_api_key_dict,
request_tags=_background_response_spend_tags(processor.data),
persist_attribution=True,
)
verbose_proxy_logger.info(
"Stored background response %s in managed objects table with unified_id=%s",
response.id,
response.id,
)
except Exception as e:
verbose_proxy_logger.error("Failed to store background response in managed objects table: %s", e)
return response
except ModifyResponseException as e:
@ -764,10 +779,11 @@ async def get_response(
# Normal provider response flow
data: Final = await read_request_body(request=request)
read_settles_billing: Final = not carries_client_side_overrides(data)
data["response_id"] = response_id
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
response: Final[object] = await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
@ -792,6 +808,266 @@ async def get_response(
proxy_logging_obj=proxy_logging_obj,
version=version,
)
await _bill_finished_background_response(
response=response,
response_id=response_id,
user_api_key_dict=user_api_key_dict,
fastapi_response=fastapi_response,
read_settles_billing=read_settles_billing,
)
return response
def _background_response_spend_tags(data: Mapping[str, object]) -> tuple[str, ...] | None:
proxy_server_request: Final = data.get("proxy_server_request")
client_tags: Final = (
StandardLoggingPayloadSetup.get_request_tags(
{}, _PROXY_SERVER_REQUEST_ADAPTER.validate_python(proxy_server_request)
)
if isinstance(proxy_server_request, dict)
else []
)
tags: Final = (*(request_tags_from_request_data(data) or ()), *client_tags)
return tags or None
@runtime_checkable
class _BackgroundResponseStore(Protocol):
async def store_unified_object_id(
self,
unified_object_id: str,
file_object: ResponsesAPIResponse,
litellm_parent_otel_span: None,
model_object_id: str,
file_purpose: Literal["response"],
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None = None,
persist_attribution: bool = False,
) -> None: ...
async def track_untracked_background_response(
self,
response_id: str,
response: ResponsesAPIResponse,
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None = None,
) -> bool: ...
def _managed_files_hook(proxy_logging_obj: "ProxyLogging") -> _BackgroundResponseStore | None:
hook: Final = proxy_logging_obj.get_proxy_hook("managed_files")
return hook if isinstance(hook, _BackgroundResponseStore) else None
async def _bill_finished_background_response(
response: object,
response_id: str,
user_api_key_dict: UserAPIKeyAuth,
fastapi_response: Response,
read_settles_billing: bool,
) -> None:
from litellm.proxy.proxy_server import proxy_logging_obj
if not isinstance(response, ResponsesAPIResponse) or not is_background_response(response):
return
if response.status not in TERMINAL_RESPONSE_STATUSES:
return
managed_files_hook: Final = _managed_files_hook(proxy_logging_obj)
track_untracked: Final = (
partial(
managed_files_hook.track_untracked_background_response,
response_id,
user_api_key_dict=user_api_key_dict,
)
if managed_files_hook is not None
else None
)
charge: Final = await _settle_background_billing(
(lambda cost_check: cost_check.bill_if_finished(response_id, response, track_untracked))
if read_settles_billing
else (lambda cost_check: cost_check.bill_from_deployment(response_id, track_untracked))
)
_report_charge(charge, fastapi_response)
class _BilledCharge(Protocol):
@property
def cost(self) -> float | None: ...
@property
def logging_obj(self) -> "LiteLLMLoggingObj": ...
_UntrackedResponseTracker: TypeAlias = Callable[[ResponsesAPIResponse], Awaitable[bool]]
class _BackgroundCostCheck(Protocol):
async def bill_from_deployment(
self, response_id: str, track_untracked: _UntrackedResponseTracker | None = None
) -> _BilledCharge | None: ...
async def bill_if_finished(
self,
response_id: str,
response: ResponsesAPIResponse,
track_untracked: _UntrackedResponseTracker | None = None,
) -> _BilledCharge | None: ...
def _background_cost_check() -> _BackgroundCostCheck | None:
from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_logging_obj
if prisma_client is None or llm_router is None:
return None
try:
from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost
except ImportError:
return None
return CheckResponsesCost( # pyright: ignore[reportUnknownVariableType] # litellm_enterprise is absent from the type-check venv
proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, llm_router=llm_router
)
_BillingSettlement: TypeAlias = Callable[[_BackgroundCostCheck], Awaitable[_BilledCharge | None]]
_BILLING_RETRY_DELAYS_SECONDS: Final = (5.0, 30.0, 120.0)
_billing_retries: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: the event loop keeps only weak task refs
@dataclass(frozen=True, slots=True)
class _BillingFailed:
error: Exception
async def _attempt_billing(
settle: _BillingSettlement, cost_check: _BackgroundCostCheck
) -> _BilledCharge | None | _BillingFailed:
try:
return await settle(cost_check)
except Exception as e:
return _BillingFailed(e)
async def _retry_background_billing(
settle: _BillingSettlement, cost_check: _BackgroundCostCheck, retry_delays: Sequence[float]
) -> None:
for delay in retry_delays:
await asyncio.sleep(delay)
outcome = await _attempt_billing(settle, cost_check)
if not isinstance(outcome, _BillingFailed):
return
verbose_proxy_logger.warning("Retrying the bill for a finished background response failed: %s", outcome.error)
verbose_proxy_logger.error("Gave up billing a finished background response; the cost poll bills it when enabled")
def _schedule_billing_retry(
settle: _BillingSettlement, cost_check: _BackgroundCostCheck, retry_delays: Sequence[float]
) -> None:
retry: Final = asyncio.create_task(_retry_background_billing(settle, cost_check, retry_delays))
_billing_retries.add(retry)
retry.add_done_callback(_billing_retries.discard)
async def _settle_background_billing(
settle: _BillingSettlement, retry_delays: Sequence[float] = _BILLING_RETRY_DELAYS_SECONDS
) -> _BilledCharge | None:
cost_check: Final = _background_cost_check()
if cost_check is None:
return None
outcome: Final = await _attempt_billing(settle, cost_check)
if not isinstance(outcome, _BillingFailed):
return outcome
verbose_proxy_logger.warning(
"Billing a finished background response failed, retrying in background: %s", outcome.error
)
_schedule_billing_retry(settle, cost_check, retry_delays)
return None
def _report_charge(charge: _BilledCharge | None, fastapi_response: Response) -> None:
if charge is None or charge.cost is None:
return
billed_headers: Final = cost_headers(charge.cost, charge.logging_obj)
fastapi_response.headers.update({name: value for name, value in billed_headers.items() if value is not None})
def _caller_may_address_response(
response_id: str, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: "ProxyLogging"
) -> bool:
security_hook: Final = proxy_logging_obj.get_proxy_hook("responses_id_security")
if not isinstance(security_hook, ResponsesIDSecurity):
return True
try:
security_hook.authorize_response_id(response_id, user_api_key_dict)
except HTTPException:
return False
return True
def _without_managed_background_marker(request_data: Mapping[str, object]) -> Mapping[str, object]:
metadata: Final = request_data.get("litellm_metadata")
if not isinstance(metadata, Mapping):
return dict(request_data)
return {
**request_data,
"litellm_metadata": {
key: value
for key, value in _PROXY_SERVER_REQUEST_ADAPTER.validate_python(metadata).items()
if key != MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
},
}
def _streamed_background_data_generator(
select_data_generator: Callable[..., AsyncIterator[str | bytes]],
proxy_logging_obj: "ProxyLogging",
response: object,
user_api_key_dict: UserAPIKeyAuth,
request_data: dict[str, object],
request: Request | None = None,
) -> AsyncIterator[str | bytes]:
chunks: Final = select_data_generator(
response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, request=request
)
managed_files_hook: Final = _managed_files_hook(proxy_logging_obj)
if managed_files_hook is None:
return chunks
track_untracked: Final = partial(
managed_files_hook.track_untracked_background_response,
user_api_key_dict=user_api_key_dict,
request_tags=_background_response_spend_tags(request_data),
)
return _bill_streamed_background_response(response, chunks, track_untracked)
async def _bill_streamed_background_response(
stream: object,
chunks: AsyncIterator[str | bytes],
track_untracked: Callable[[str, ResponsesAPIResponse], Awaitable[bool]],
) -> AsyncIterator[str | bytes]:
async for chunk in chunks:
completed: object = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream)
if not isinstance(completed, ResponsesAPIResponse):
yield chunk
continue
await _bill_finished_stream(completed, track_untracked)
yield chunk
async for remaining_chunk in chunks:
yield remaining_chunk
return
late_completed: Final = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream)
if isinstance(late_completed, ResponsesAPIResponse):
await _bill_finished_stream(late_completed, track_untracked)
async def _bill_finished_stream(
response: ResponsesAPIResponse, track_untracked: Callable[[str, ResponsesAPIResponse], Awaitable[bool]]
) -> None:
if not is_background_response(response):
return
await _settle_background_billing(
lambda cost_check: cost_check.bill_if_finished(response.id, response, partial(track_untracked, response.id))
)
@router.delete(
@ -874,9 +1150,14 @@ async def delete_response(
# Normal provider response flow
data: Final = await read_request_body(request=request)
data["response_id"] = response_id
charge: Final = (
await _settle_background_billing(lambda cost_check: cost_check.bill_from_deployment(response_id))
if _caller_may_address_response(response_id, user_api_key_dict, proxy_logging_obj)
else None
)
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
deleted: Final = await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
@ -901,6 +1182,8 @@ async def delete_response(
proxy_logging_obj=proxy_logging_obj,
version=version,
)
_report_charge(charge, fastapi_response)
return deleted
@router.get(
@ -1287,10 +1570,11 @@ async def cancel_response(
# Normal provider response flow
data: Final = await read_request_body(request=request)
read_settles_billing: Final = not carries_client_side_overrides(data)
data["response_id"] = response_id
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
response: Final[object] = await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
@ -1315,6 +1599,14 @@ async def cancel_response(
proxy_logging_obj=proxy_logging_obj,
version=version,
)
await _bill_finished_background_response(
response=response,
response_id=response_id,
user_api_key_dict=user_api_key_dict,
fastapi_response=fastapi_response,
read_settles_billing=read_settles_billing,
)
return response
def _resolve_responses_ws_session_limit_seconds() -> float:
@ -1578,7 +1870,7 @@ async def _responses_websocket_session(
}
)
)
except Exception:
except (WebSocketDisconnect, RuntimeError, OSError):
pass
await websocket.close(code=1008, reason="Pre-call error")
return

1
litellm/proxy/schema.prisma
View file

@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
end_user_id String?
updated_at DateTime @updatedAt
updated_by String?

27
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -17,6 +17,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.constants import (
CLI_SESSION_KEY_PREFIX,
EMPTY_MAPPING,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
LITELLM_TRUNCATED_PAYLOAD_FIELD,
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
@ -24,6 +25,7 @@ from litellm.constants import (
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
MAX_SPEND_LOG_MODEL_NAME_LENGTH,
MCP_SPEND_LOG_MODEL_PREFIX,
NON_INFERENCE_CALL_TYPES,
REDACTED_BY_LITELM_STRING,
SESSION_ID_OMITTED_METADATA_KEY,
UNKNOWN_MODEL_SPEND_LOG_MODEL,
@ -63,6 +65,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.proxy.utils import PrismaClient, hash_token
from litellm.types.router import DeploymentTypedDict, LiteLLM_Params
from litellm.types.utils import (
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
PROMPT_CARRYING_GUARDRAIL_FIELDS,
AzureSpillover,
CallTypes,
@ -292,9 +295,23 @@ def _get_spend_logs_metadata(
BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost"
BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost"
def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | None:
def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object]) -> str:
if call_type == CallTypes.aretrieve_batch.value:
return BATCH_COST_REQUEST_ID_SUFFIX
if (
call_type in NON_INFERENCE_CALL_TYPES
and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
):
return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX
return ""
def get_spend_logs_id(
call_type: str, response_obj: dict, kwargs: dict, metadata: Mapping[str, object] = EMPTY_MAPPING
) -> str | None:
standard_logging_payload = kwargs.get("standard_logging_object")
candidate_ids: Final = (
response_obj.get("id"),
@ -304,9 +321,9 @@ def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str |
resolved_id: Final = next(
(candidate for candidate in candidate_ids if isinstance(candidate, str) and candidate), None
)
if resolved_id is not None and call_type == CallTypes.aretrieve_batch.value:
return f"{resolved_id}{BATCH_COST_REQUEST_ID_SUFFIX}"
return resolved_id
if resolved_id is None:
return None
return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata)}"
_MISSING_ATTRIBUTE: Final = object()
@ -560,7 +577,7 @@ def get_logging_payload(
if not usage and isinstance(_combined_usage, litellm.Usage):
usage = _combined_usage.model_dump()
id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs)
id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs, metadata)
standard_logging_payload: Final = cast(StandardLoggingPayload | None, kwargs.get("standard_logging_object", None))
end_user_id = get_end_user_id_for_cost_tracking(litellm_params)

8
litellm/router.py
View file

@ -214,7 +214,7 @@ from litellm.router_utils.cooldown_handlers import (
async_get_cooldown_deployments_with_debug_info,
get_cooldown_deployments,
is_advisor_orchestration_failure,
is_background_response_cost_poll_not_found,
is_background_response_poll_not_found,
is_caller_scoped_auth_failure,
is_caller_timeout_408,
set_cooldown_deployments,
@ -3617,7 +3617,7 @@ class Router:
# fallback_iterator, or any future wrapper). Without this
# the proxy's container-ownership hook (which reads
# ``getattr(stream_response, "completed_response", None)``
# via _extract_completed_responses_response) silently
# via extract_completed_responses_response) silently
# records nothing on streaming /v1/responses calls — every
# follow-up /v1/containers/<id>/files call then 403s for
# the very key that created the container (#30210).
@ -8743,10 +8743,10 @@ class Router:
litellm_params: Final = kwargs.get("litellm_params", {})
_model_info: Final = litellm_params.get("model_info", {})
if is_background_response_cost_poll_not_found(exception, litellm_params):
if is_background_response_poll_not_found(exception, litellm_params):
verbose_router_logger.debug(
"Router: Exiting 'deployment_callback_on_failure' without cooldown. "
"Provider 404 came from the background response cost poll, not the deployment's health."
"Provider 404 came from the background response poll, not the deployment's health."
)
return False

17
litellm/router_utils/cooldown_handlers.py
View file

@ -30,7 +30,11 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCacheValue
from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback
from litellm.types.router import Deployment
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
from litellm.types.utils import (
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
InternalCallOrigin,
)
from .router_callbacks.track_deployment_metrics import (
get_deployment_failures_for_current_minute,
@ -71,11 +75,16 @@ def is_advisor_orchestration_failure(exception: BaseException | None) -> bool:
return bool(getattr(exception, _ADVISOR_ORCHESTRATION_FAILURE_ATTR, False))
def is_background_response_cost_poll_not_found(exception: Exception, litellm_params: Mapping[str, object]) -> bool:
"""Whether a background response cost poll failed with a provider 404."""
_BACKGROUND_RESPONSE_POLL_CALL_ORIGINS: Final[frozenset[InternalCallOrigin]] = frozenset(
{BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN}
)
def is_background_response_poll_not_found(exception: Exception, litellm_params: Mapping[str, object]) -> bool:
"""Whether a background response poll, a status check or the billed read, failed with a provider 404."""
return getattr(exception, "status_code", None) == 404 and any(
isinstance(candidate, Mapping)
and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) in _BACKGROUND_RESPONSE_POLL_CALL_ORIGINS
for candidate in (litellm_params.get("metadata"), litellm_params.get("litellm_metadata"))
)

2
litellm/types/utils.py
View file

@ -3164,6 +3164,7 @@ InternalCallOrigin = Literal[
"shadow_eval_judge",
"llm_as_a_judge_guardrail",
"background_response_cost_poll",
"background_response_status_poll",
]
"""Which internal litellm feature originated a billed sub-call, so a spend log row
records that it is not traffic the caller sent."""
@ -3173,6 +3174,7 @@ SHADOW_EVAL_ROUTER_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_router"
SHADOW_EVAL_JUDGE_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_judge"
LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN: Final[InternalCallOrigin] = "llm_as_a_judge_guardrail"
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_cost_poll"
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_status_poll"
class StandardLoggingHeuristicV2Forecast(TypedDict):

1
schema.prisma
View file

@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
end_user_id String?
updated_at DateTime @updatedAt
updated_by String?

327
tests/integration/spend/_background_responses.py Normal file
View file

@ -0,0 +1,327 @@
import itertools
import json
import math
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from hashlib import sha256
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, Scenario, eventually, held, list_value, object_value, string_value
from integration._support.database import read_rows, write_rows
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, SseResponse
from pydantic import JsonValue
INPUT_COST_PER_TOKEN: Final = 0.001
OUTPUT_COST_PER_TOKEN: Final = 0.002
INPUT_TOKENS: Final = 19
OUTPUT_TOKENS: Final = 7
PRICE: Final = INPUT_TOKENS * INPUT_COST_PER_TOKEN + OUTPUT_TOKENS * OUTPUT_COST_PER_TOKEN
SCHEDULER_PERIOD_CEILING_SECONDS: Final = 31
ONE_POLL_SECONDS: Final = 2 * SCHEDULER_PERIOD_CEILING_SECONDS + 8
POLL_WINDOW_SECONDS: Final = SCHEDULER_PERIOD_CEILING_SECONDS + 9
SPEND_LANDING_SECONDS: Final = 60
PROVIDER_ID: Final = "resp_$REQUEST_ID"
PROMPT: Final = "bill me once"
READS: Final = 3
CONCURRENT_READS: Final = 12
FINISHED_STATUSES: Final = frozenset({"completed", "incomplete", "cancelled"})
MANAGED_READ_MARKER: Final = "_managed_background_response"
ATTRIBUTION_COLUMNS: Final = ("api_key", "team_id", "organization_id", "user")
def response_body(status: str, *, background: bool = True, **fields: JsonValue) -> dict[str, JsonValue]:
finished: Final = status in FINISHED_STATUSES
return {
"id": PROVIDER_ID,
"object": "response",
"created_at": 1,
"status": status,
"background": background,
"store": False,
"error": None,
"incomplete_details": {"reason": "max_output_tokens"} if status == "incomplete" else None,
"model": "gpt-4o-mini",
"output": (
[
{
"id": "msg_$REQUEST_ID",
"type": "message",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "pong", "annotations": []}],
}
]
if finished
else []
),
"parallel_tool_calls": True,
"temperature": 1.0,
"tool_choice": "auto",
"tools": [],
"top_p": 1.0,
"truncation": "disabled",
"usage": (
{
"input_tokens": INPUT_TOKENS,
"output_tokens": OUTPUT_TOKENS,
"total_tokens": INPUT_TOKENS + OUTPUT_TOKENS,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
}
if finished
else None
),
"metadata": {},
**fields,
}
def json_reply(body: dict[str, JsonValue], status: int = 200) -> JsonResponse:
return JsonResponse(content_type="application/json", body=body, status=status)
def scripted_routes(
create: JsonResponse, read: JsonResponse, *, provider_id: str = PROVIDER_ID, **extra: JsonResponse
) -> RoutedResponse:
return RoutedResponse(
content_type="application/x-routed",
routes={
"POST /responses": create,
f"GET /responses/{provider_id}": read,
**{f"POST /responses/{provider_id}/{action}": response for action, response in extra.items()},
},
)
def background_routes(read_status: str, **fields: JsonValue) -> RoutedResponse:
return scripted_routes(
json_reply(response_body("queued", **fields)), json_reply(response_body(read_status, **fields))
)
def sse_frame(event: str, sequence_number: int, response: dict[str, JsonValue]) -> str:
return f"event: {event}\ndata: " + json.dumps(
{"type": event, "sequence_number": sequence_number, "response": response}, separators=(",", ":")
)
def streamed_reply(*, background: bool = True) -> SseResponse:
return SseResponse(
content_type="text/event-stream",
frames=(
sse_frame("response.created", 0, response_body("queued", background=background)),
sse_frame("response.completed", 1, response_body("completed", background=background)),
),
)
@dataclass(frozen=True, slots=True)
class Upstream:
model: str
model_id: str
handle: ScenarioHandle
def script(self, response: RoutedResponse | SseResponse) -> None:
register_scenario(self.handle.scenario_id, response, control_url=self.handle.control_url)
def scripted_upstream(
scenario: Scenario,
response: RoutedResponse | SseResponse,
*,
model_info: Mapping[str, JsonValue] | None = None,
**litellm_params: JsonValue,
) -> Upstream:
scenario_id: Final = f"bgonce-{uuid.uuid4().hex[:12]}"
handle: Final = register_scenario(scenario_id, response)
scenario.cleanups.callback(delete_scenario, handle)
created: Final = scenario.gateway.post(
"/model/new",
{
"model_name": scenario_id,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "sk-scripted-provider",
"api_base": handle.api_base(),
"input_cost_per_token": INPUT_COST_PER_TOKEN,
"output_cost_per_token": OUTPUT_COST_PER_TOKEN,
**litellm_params,
},
**({"model_info": dict(model_info)} if model_info else {}),
},
)
model_id: Final = string_value(object_value(created["model_info"])["id"])
scenario.cleanups.callback(scenario.delete_model, model_id)
return Upstream(scenario_id, model_id, handle)
def forget_row(response_id: str) -> None:
write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', (response_id,))
def create_background_response(
scenario: Scenario,
key: str,
model: str,
*,
body: Mapping[str, JsonValue] | None = None,
headers: Mapping[str, str] | None = None,
) -> str:
response: Final = scenario.gateway.request(
"POST",
"/v1/responses",
body or {"model": model, "input": PROMPT, "background": True, "store": False},
key=key,
headers=headers,
)
assert response.status_code == 200, response.text
response_id: Final = string_value(object_value(response.json())["id"])
scenario.cleanups.callback(forget_row, response_id)
return response_id
@dataclass(frozen=True, slots=True)
class Read:
status: str
text: str
cost: float | None
def reported_cost(headers: httpx.Headers) -> float | None:
reported: Final = headers.get("x-litellm-response-cost")
return None if reported is None else float(reported)
def read_of(response: httpx.Response) -> Read:
assert response.status_code == 200, response.text
body: Final = object_value(response.json())
output: Final = list_value(body["output"])
text: Final = (
string_value(object_value(list_value(object_value(output[0])["content"])[0])["text"]) if output else ""
)
return Read(string_value(body["status"]), text, reported_cost(response.headers))
def read_response(gateway: Gateway, key: str, response_id: str, body: Mapping[str, JsonValue] | None = None) -> Read:
return read_of(gateway.request("GET", f"/v1/responses/{response_id}", body, key=key))
def free_reads(status: str, count: int) -> tuple[Read, ...]:
return (Read(status, "pong", 0.0),) * count
def settling_reads(status: str, count: int) -> tuple[Read, ...]:
return (Read(status, "pong", pytest.approx(PRICE)), *free_reads(status, count - 1))
def hashed(key: str) -> str:
return sha256(key.encode()).hexdigest()
def number(value: JsonValue) -> float:
assert isinstance(value, int | float), value
return float(value)
def key_spend(key: str) -> float:
rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed(key),))
return number(rows[0]["spend"])
def team_spend(team_id: str) -> float:
rows: Final = read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
return number(rows[0]["spend"])
def user_spend(user_id: str) -> float:
rows: Final = read_rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = %s', (user_id,))
return number(rows[0]["spend"])
def org_spend(organization_id: str) -> float:
rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (organization_id,)
)
return number(rows[0]["spend"])
def tag_spend(tag: str) -> float:
rows: Final = read_rows(
'SELECT COALESCE(SUM(spend), 0) AS spend FROM "LiteLLM_DailyTagSpend" WHERE tag = %s', (tag,)
)
return number(rows[0]["spend"])
def costs(expected: float) -> Callable[[float], bool]:
return lambda spend: math.isclose(spend, expected, abs_tol=1e-9)
def assert_charged_exactly(read_spend: Callable[[], float], expected: float) -> None:
held(read_spend, costs(expected), holding=POLL_WINDOW_SECONDS, seconds=POLL_WINDOW_SECONDS + SPEND_LANDING_SECONDS)
def charged_rows(model_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT request_id, call_type, api_key, team_id, organization_id, "user", request_tags, spend, '
'prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_id = %s AND spend > 0',
(model_id,),
)
def create_rows(model_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT api_key, team_id, organization_id, "user", request_tags FROM "LiteLLM_SpendLogs" '
"WHERE model_id = %s AND call_type = 'aresponses'",
(model_id,),
)
def one_charged_row(model_id: str) -> dict[str, JsonValue]:
rows: Final = eventually(lambda: charged_rows(model_id), lambda found: len(found) >= 1, seconds=30)
assert len(rows) == 1, rows
return rows[0]
def managed_rows(response_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT status, api_key, created_by, team_id FROM "LiteLLM_ManagedObjectTable" '
"WHERE unified_object_id = %s OR model_object_id = %s",
(response_id, response_id),
)
def assert_billed_like_its_creator(
row: dict[str, JsonValue], key: str, team_id: str | None, user_id: str | None
) -> None:
assert string_value(row["request_id"]).endswith("_background_cost"), row
assert row["call_type"] == "aget_responses"
assert row["api_key"] == hashed(key)
assert row["team_id"] == team_id
assert row["user"] == user_id
assert row["prompt_tokens"] == INPUT_TOKENS
assert row["completion_tokens"] == OUTPUT_TOKENS
assert number(row["spend"]) == pytest.approx(PRICE)
def observed_requests(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
return tuple(map(object_value, list_value(upstream.get("/__observations").json()["requests"])))
def calls_to(requests: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> tuple[dict[str, JsonValue], ...]:
return tuple(request for request in requests if string_value(request["path"]).startswith(f"/{handle.scenario_id}/"))
def upstream_log(gateway: Gateway, handle: ScenarioHandle) -> Iterator[tuple[dict[str, JsonValue], ...]]:
fresh: Final = (calls_to(observed_requests(gateway), handle) for _ in itertools.count())
return itertools.accumulate(fresh, lambda seen, calls: (*seen, *calls))
def submissions(log: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> tuple[JsonValue, ...]:
return tuple(
call["body"] for call in log if call["method"] == "POST" and call["path"] == f"/{handle.scenario_id}/responses"
)

771
tests/integration/spend/test_background_response_billed_once.py Normal file
View file

@ -0,0 +1,771 @@
import asyncio
import json
import uuid
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from typing import Final
import openai
import pytest
from integration._support.client import Gateway, Scenario, eventually, held, list_value, object_value, string_value
from integration._support.database import read_rows, write_rows
from integration._support.upstream import delete_scenario, register_scenario
from integration.caching.response_cache_case import entry_key, redis_store
from integration.spend._background_responses import (
ATTRIBUTION_COLUMNS,
CONCURRENT_READS,
INPUT_COST_PER_TOKEN,
INPUT_TOKENS,
MANAGED_READ_MARKER,
ONE_POLL_SECONDS,
OUTPUT_COST_PER_TOKEN,
OUTPUT_TOKENS,
POLL_WINDOW_SECONDS,
PRICE,
PROMPT,
READS,
SPEND_LANDING_SECONDS,
Read,
Upstream,
assert_billed_like_its_creator,
assert_charged_exactly,
background_routes,
charged_rows,
costs,
create_background_response,
create_rows,
forget_row,
free_reads,
hashed,
json_reply,
key_spend,
managed_rows,
number,
one_charged_row,
org_spend,
read_of,
read_response,
reported_cost,
response_body,
scripted_routes,
scripted_upstream,
settling_reads,
streamed_reply,
submissions,
tag_spend,
team_spend,
upstream_log,
user_spend,
)
from openai._legacy_response import LegacyAPIResponse
from openai.types.responses import Response
from pydantic import JsonValue
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
def _sdk_read(raw: LegacyAPIResponse[Response]) -> Read:
parsed: Final = raw.parse()
return Read(parsed.status or "", parsed.output_text, reported_cost(raw.headers))
def _httpx_flow(scenario: Scenario, key: str, model: str) -> tuple[str, tuple[Read, ...]]:
response_id: Final = create_background_response(scenario, key, model)
return response_id, tuple(read_response(scenario.gateway, key, response_id) for _ in range(READS))
def _forged_metadata_flow(scenario: Scenario, key: str, model: str) -> tuple[str, tuple[Read, ...]]:
forged: Final[dict[str, JsonValue]] = {
"internal_call_origin": BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
MANAGED_READ_MARKER: False,
}
response_id: Final = create_background_response(scenario, key, model)
return response_id, tuple(
read_response(scenario.gateway, key, response_id, {"metadata": forged, "litellm_metadata": forged})
for _ in range(READS)
)
def _sdk_poll(client: openai.OpenAI, response_id: str, remaining: int) -> tuple[Read, ...]:
if remaining == 0:
return ()
raw: Final = client.responses.with_raw_response.retrieve(response_id)
return (_sdk_read(raw), *_sdk_poll(client, raw.parse().id, remaining - 1))
def _sdk_sync_flow(scenario: Scenario, key: str, model: str) -> tuple[str, tuple[Read, ...]]:
with openai.OpenAI(base_url=f"{scenario.gateway.client.base_url}/v1", api_key=key, max_retries=0) as client:
created: Final = client.responses.create(model=model, input=PROMPT, background=True, store=False)
scenario.cleanups.callback(forget_row, created.id)
return created.id, _sdk_poll(client, created.id, READS)
async def _sdk_async_poll(client: openai.AsyncOpenAI, response_id: str, remaining: int) -> tuple[Read, ...]:
if remaining == 0:
return ()
raw: Final = await client.responses.with_raw_response.retrieve(response_id)
return (_sdk_read(raw), *await _sdk_async_poll(client, raw.parse().id, remaining - 1))
async def _sdk_async_reads(base_url: str, key: str, model: str) -> tuple[str, tuple[Read, ...]]:
async with openai.AsyncOpenAI(base_url=base_url, api_key=key, max_retries=0) as client:
created: Final = await client.responses.create(model=model, input=PROMPT, background=True, store=False)
return created.id, await _sdk_async_poll(client, created.id, READS)
def _sdk_async_flow(scenario: Scenario, key: str, model: str) -> tuple[str, tuple[Read, ...]]:
response_id, reads = asyncio.run(_sdk_async_reads(f"{scenario.gateway.client.base_url}/v1", key, model))
scenario.cleanups.callback(forget_row, response_id)
return response_id, reads
Flow = Callable[[Scenario, str, str], tuple[str, tuple[Read, ...]]]
FLOWS: Final[Mapping[str, Flow]] = {
"httpx": _httpx_flow,
"httpx_forged_metadata": _forged_metadata_flow,
"openai_sync": _sdk_sync_flow,
"openai_async": _sdk_async_flow,
}
def _team_member_key(scenario: Scenario, **fields: JsonValue) -> tuple[str, str, str]:
team_id: Final = scenario.team(**fields)
user_id: Final = scenario.member(team_id)
return scenario.key(team_id=team_id, user_id=user_id), team_id, user_id
@pytest.mark.timeout(240)
@pytest.mark.parametrize(
("client", "terminal_status"),
[
("httpx", "completed"),
("httpx", "incomplete"),
("httpx_forged_metadata", "completed"),
("openai_sync", "completed"),
("openai_async", "completed"),
],
)
def test_a_finished_background_response_bills_its_creator_once_however_often_it_is_read(
gateway: Gateway, client: str, terminal_status: str
) -> None:
with gateway.scenario() as scenario:
key, team_id, user_id = _team_member_key(scenario)
upstream: Final = scripted_upstream(scenario, background_routes(terminal_status))
log: Final = upstream_log(gateway, upstream.handle)
response_id, reads = FLOWS[client](scenario, key, upstream.model)
assert reads in (settling_reads(terminal_status, READS), free_reads(terminal_status, READS))
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert_billed_like_its_creator(one_charged_row(upstream.model_id), key, team_id, user_id)
assert team_spend(team_id) == pytest.approx(PRICE)
assert submissions(next(log), upstream.handle) == (
{"model": "gpt-4o-mini", "input": PROMPT, "background": True, "store": False},
)
assert [row["status"] for row in managed_rows(response_id)] == ["completed"]
def _org_team_member_key(scenario: Scenario) -> tuple[str, str, str, str]:
organization_id: Final = scenario.organization()
team_id: Final = scenario.team(organization_id=organization_id)
user_id: Final = scenario.org_member(organization_id, "internal_user")
scenario.gateway.post("/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}})
return scenario.key(team_id=team_id, user_id=user_id), team_id, user_id, organization_id
def _billed_by_the_poll_alone(
scenario: Scenario, key: str, tags: tuple[str, str]
) -> tuple[Upstream, dict[str, JsonValue]]:
body_tag, header_tag = tags
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
response_id: Final = create_background_response(
scenario,
key,
upstream.model,
body={
"model": upstream.model,
"input": PROMPT,
"background": True,
"store": False,
"litellm_metadata": {"tags": [body_tag]},
},
headers={"x-litellm-tags": header_tag},
)
eventually(
lambda: [row["status"] for row in managed_rows(response_id)],
lambda statuses: statuses == ["completed"],
seconds=ONE_POLL_SECONDS,
)
held(lambda: key_spend(key), costs(PRICE), holding=5, seconds=SPEND_LANDING_SECONDS)
charged: Final = one_charged_row(upstream.model_id)
create_row: Final = eventually(lambda: create_rows(upstream.model_id), lambda rows: len(rows) == 1)[0]
assert {column: charged[column] for column in ATTRIBUTION_COLUMNS} == {
column: create_row[column] for column in ATTRIBUTION_COLUMNS
}
assert set(list_value(charged["request_tags"])) == set(list_value(create_row["request_tags"])) >= set(tags)
return upstream, charged
def _unique_tags() -> tuple[str, str]:
return f"audit-body-{uuid.uuid4().hex[:8]}", f"audit-header-{uuid.uuid4().hex[:8]}"
@pytest.mark.timeout(240)
def test_a_background_response_nobody_reads_is_billed_once_to_its_creating_key_team_org_and_tags(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
key, team_id, user_id, organization_id = _org_team_member_key(scenario)
tags: Final = _unique_tags()
_, charged = _billed_by_the_poll_alone(scenario, key, tags)
assert_billed_like_its_creator(charged, key, team_id, user_id)
assert charged["organization_id"] == organization_id
eventually(lambda: tuple(map(tag_spend, tags)), lambda spent: all(map(costs(PRICE), spent)))
eventually(lambda: org_spend(organization_id), costs(PRICE))
listed: Final = gateway.request(
"GET", "/spend/logs", params={"request_id": string_value(charged["request_id"])}
)
assert listed.status_code == 200, listed.text
assert [number(object_value(row)["spend"]) for row in list_value(listed.json())] == [pytest.approx(PRICE)]
assert number(object_value(gateway.get("/key/info", {"key": key})["info"])["spend"]) == pytest.approx(PRICE)
team_info: Final = object_value(gateway.get("/team/info", {"team_id": team_id})["team_info"])
assert number(team_info["spend"]) == pytest.approx(PRICE)
@pytest.mark.timeout(240)
@pytest.mark.parametrize("in_a_team", [False, True], ids=["no_team", "team"])
def test_a_background_response_nobody_reads_on_a_key_with_no_user_is_billed_once_to_the_key(
gateway: Gateway, in_a_team: bool
) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key(team_id=scenario.team()) if in_a_team else scenario.key()
_, charged = _billed_by_the_poll_alone(scenario, key, _unique_tags())
assert string_value(charged["request_id"]).endswith("_background_cost"), charged
assert charged["api_key"] == hashed(key)
assert number(charged["spend"]) == pytest.approx(PRICE)
@pytest.mark.timeout(240)
def test_a_background_response_is_free_until_it_finishes_and_a_failed_read_bills_nothing(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream: Final = scripted_upstream(scenario, background_routes("in_progress"))
response_id: Final = create_background_response(scenario, key, upstream.model)
assert read_response(gateway, key, response_id) == Read("in_progress", "", 0.0)
outage: Final = json_reply({"error": {"message": "The server had an error.", "type": "server_error"}}, 500)
upstream.script(scripted_routes(json_reply(response_body("queued")), outage))
failed: Final = gateway.request("GET", f"/v1/responses/{response_id}", key=key)
assert failed.status_code == 500, failed.text
assert "The server had an error." in failed.text
assert [row["status"] for row in managed_rows(response_id)] == ["queued"]
upstream.script(background_routes("completed"))
assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == settling_reads("completed", 2)
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert [row["status"] for row in managed_rows(response_id)] == ["completed"]
@pytest.mark.timeout(240)
def test_cancelling_a_background_response_bills_it_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
cancelled: Final = json_reply(response_body("cancelled"))
upstream: Final = scripted_upstream(
scenario, scripted_routes(json_reply(response_body("queued")), cancelled, cancel=cancelled)
)
response_id: Final = create_background_response(scenario, key, upstream.model)
cancel: Final = gateway.request("POST", f"/v1/responses/{response_id}/cancel", key=key)
reads: Final = (read_of(cancel), read_response(gateway, key, response_id))
assert reads in (settling_reads("cancelled", 2), free_reads("cancelled", 2))
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert number(one_charged_row(upstream.model_id)["spend"]) == pytest.approx(PRICE)
@pytest.mark.timeout(120)
def test_a_foreground_response_bills_its_create_once_and_its_read_is_free(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
foreground: Final = json_reply(response_body("completed", background=False))
upstream: Final = scripted_upstream(scenario, scripted_routes(foreground, foreground))
created: Final = gateway.request(
"POST", "/v1/responses", {"model": upstream.model, "input": PROMPT, "store": False}, key=key
)
assert read_of(created) == Read("completed", "pong", pytest.approx(PRICE))
response_id: Final = string_value(object_value(created.json())["id"])
assert read_response(gateway, key, response_id) == Read("completed", "pong", 0.0)
held(lambda: key_spend(key), costs(PRICE), holding=10, seconds=SPEND_LANDING_SECONDS)
assert managed_rows(response_id) == []
STREAMED_BACKGROUND_FIELDS: Final[Mapping[str, dict[str, JsonValue]]] = {
"top_level": {"background": True},
"extra_body": {"extra_body": {"background": True}},
}
def _streamed_response_id(gateway: Gateway, key: str, model: str, background: Mapping[str, JsonValue]) -> str:
with gateway.client.stream(
"POST",
"/v1/responses",
json={"model": model, "input": PROMPT, "store": False, "stream": True, **background},
headers={"Authorization": f"Bearer {key}"},
) as stream:
assert stream.status_code == 200, stream.read()
events: Final = tuple(
object_value(json.loads(line.removeprefix("data: ")))
for line in stream.iter_lines()
if line.startswith("data: {")
)
completed: Final = [event for event in events if event["type"] == "response.completed"]
assert len(completed) == 1, events
return string_value(object_value(completed[0]["response"])["id"])
@pytest.mark.timeout(240)
@pytest.mark.parametrize("background_in", STREAMED_BACKGROUND_FIELDS.keys())
def test_a_streamed_background_response_read_later_bills_its_stream_alone(gateway: Gateway, background_in: str) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream: Final = scripted_upstream(scenario, streamed_reply())
response_id: Final = _streamed_response_id(
gateway, key, upstream.model, STREAMED_BACKGROUND_FIELDS[background_in]
)
scenario.cleanups.callback(forget_row, response_id)
upstream.script(background_routes("completed"))
assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == free_reads("completed", 2)
assert_charged_exactly(lambda: key_spend(key), PRICE)
def _finished(scenario: Scenario, key: str) -> tuple[Upstream, str]:
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
return upstream, create_background_response(scenario, key, upstream.model)
def _lose_the_row_then_finish(scenario: Scenario, key: str) -> tuple[Upstream, str]:
upstream: Final = scripted_upstream(scenario, background_routes("queued"))
response_id: Final = create_background_response(scenario, key, upstream.model)
forget_row(response_id)
upstream.script(background_routes("completed"))
return upstream, response_id
@pytest.mark.timeout(240)
def test_a_finished_response_the_proxy_lost_track_of_bills_its_reader_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key, team_id, user_id = _team_member_key(scenario)
upstream, response_id = _lose_the_row_then_finish(scenario, key)
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS))
assert reads == settling_reads("completed", READS)
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert_billed_like_its_creator(one_charged_row(upstream.model_id), key, team_id, user_id)
assert [(row["status"], row["api_key"]) for row in managed_rows(response_id)] == [("completed", hashed(key))]
@pytest.mark.timeout(240)
@pytest.mark.parametrize("tracked", [True, False], ids=["tracked", "untracked"])
def test_concurrent_reads_of_a_finished_response_bill_it_once(gateway: Gateway, tracked: bool) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream, response_id = _finished(scenario, key) if tracked else _lose_the_row_then_finish(scenario, key)
with ThreadPoolExecutor(max_workers=CONCURRENT_READS) as pool:
reads: Final = tuple(pool.map(lambda _: read_response(gateway, key, response_id), range(CONCURRENT_READS)))
assert [(read.status, read.text) for read in reads] == [("completed", "pong")] * CONCURRENT_READS
assert [read.cost for read in reads if read.cost != 0.0] in ([], [pytest.approx(PRICE)])
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert len(charged_rows(upstream.model_id)) == 1
assert [row["status"] for row in managed_rows(response_id)] == ["completed"]
@pytest.mark.timeout(240)
def test_another_key_of_the_same_owner_reads_for_free_and_the_creating_key_pays(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
creator, team_id, user_id = _team_member_key(scenario)
reader: Final = scenario.key(team_id=team_id, user_id=user_id)
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
response_id: Final = create_background_response(scenario, creator, upstream.model)
reads: Final = tuple(read_response(gateway, reader, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
assert_charged_exactly(lambda: key_spend(creator), PRICE)
assert key_spend(reader) == 0
assert_billed_like_its_creator(one_charged_row(upstream.model_id), creator, team_id, user_id)
@pytest.mark.timeout(240)
def test_a_create_that_asks_for_background_inside_extra_body_is_billed_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
log: Final = upstream_log(gateway, upstream.handle)
response_id: Final = create_background_response(
scenario,
key,
upstream.model,
body={"model": upstream.model, "input": PROMPT, "store": False, "extra_body": {"background": True}},
)
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert [object_value(body).get("background") for body in submissions(next(log), upstream.handle)] == [True]
@pytest.mark.timeout(240)
def test_a_response_whose_instructions_carry_a_nul_is_tracked_and_billed_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream: Final = scripted_upstream(scenario, background_routes("completed", instructions="be\u0000brief"))
response_id: Final = create_background_response(scenario, key, upstream.model)
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
assert_charged_exactly(lambda: key_spend(key), PRICE)
stored: Final = read_rows(
"SELECT file_object->>'instructions' AS instructions FROM \"LiteLLM_ManagedObjectTable\" "
"WHERE unified_object_id = %s",
(response_id,),
)
assert stored == [{"instructions": "bebrief"}]
@pytest.mark.timeout(240)
def test_a_request_tag_carrying_a_nul_is_stored_without_it_and_lands_on_the_billed_row(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
suffix: Final = uuid.uuid4().hex[:8]
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
response_id: Final = create_background_response(
scenario,
key,
upstream.model,
body={
"model": upstream.model,
"input": PROMPT,
"background": True,
"store": False,
"litellm_metadata": {"tags": [f"audit-nul\u0000{suffix}"]},
},
)
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(2))
assert reads in (settling_reads("completed", 2), free_reads("completed", 2))
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert f"audit-nul{suffix}" in list_value(one_charged_row(upstream.model_id)["request_tags"])
@pytest.mark.timeout(240)
@pytest.mark.parametrize("managed_row", ["tracked", "lost"])
def test_a_read_retargeted_with_a_client_side_api_base_bills_the_real_usage_once(
gateway: Gateway, managed_row: str
) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream: Final = scripted_upstream(
scenario, background_routes("queued"), configurable_clientside_auth_params=["api_base"]
)
provider_id: Final = f"resp_{upstream.model}"
cheap_body: Final = {
**response_body("completed"),
"id": provider_id,
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
}
decoy: Final = register_scenario(
f"bgonce-decoy-{uuid.uuid4().hex[:12]}",
scripted_routes(json_reply(cheap_body), json_reply(cheap_body), provider_id=provider_id),
)
scenario.cleanups.callback(delete_scenario, decoy)
decoy_log: Final = upstream_log(gateway, decoy)
response_id: Final = create_background_response(scenario, key, upstream.model)
if managed_row == "lost":
forget_row(response_id)
upstream.script(background_routes("completed"))
retargeted: Final = read_response(
gateway, key, response_id, {"model": upstream.model, "api_base": decoy.api_base()}
)
assert (retargeted,) in (settling_reads("completed", 1), free_reads("completed", 1))
assert [call["method"] for call in next(decoy_log)] == ["GET"]
eventually(
lambda: [row["status"] for row in managed_rows(response_id)],
lambda statuses: statuses == ["completed"],
seconds=ONE_POLL_SECONDS,
)
assert_charged_exactly(lambda: key_spend(key), PRICE)
charged: Final = one_charged_row(upstream.model_id)
assert (charged["prompt_tokens"], charged["completion_tokens"]) == (INPUT_TOKENS, OUTPUT_TOKENS)
def _id_a_read_returned(gateway: Gateway, key: str, response_id: str) -> str:
read: Final = gateway.request("GET", f"/v1/responses/{response_id}", None, key=key)
assert read.status_code == 200, read.text
return string_value(object_value(read.json())["id"])
@pytest.mark.timeout(240 + ONE_POLL_SECONDS)
@pytest.mark.parametrize("read_by", ["create_id", "id_a_read_returned"])
def test_a_row_stored_before_the_upgrade_bills_the_creating_team_and_user_once(gateway: Gateway, read_by: str) -> None:
with gateway.scenario() as scenario:
key, team_id, user_id = _team_member_key(scenario)
upstream: Final = scripted_upstream(scenario, background_routes("queued"))
created_id: Final = create_background_response(scenario, key, upstream.model)
response_id: Final = created_id if read_by == "create_id" else _id_a_read_returned(gateway, key, created_id)
write_rows(
'UPDATE "LiteLLM_ManagedObjectTable" SET model_object_id = unified_object_id, api_key = NULL, '
"request_tags = NULL WHERE unified_object_id = %s",
(created_id,),
)
upstream.script(background_routes("completed"))
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
held(lambda: team_spend(team_id), costs(PRICE), holding=ONE_POLL_SECONDS, seconds=SPEND_LANDING_SECONDS)
assert user_spend(user_id) == pytest.approx(PRICE)
charged: Final = one_charged_row(upstream.model_id)
assert (charged["team_id"], charged["user"]) == (team_id, user_id)
@pytest.mark.timeout(240)
@pytest.mark.parametrize("settled_status", ["completed", "stale_expired"])
def test_reads_of_a_response_whose_row_already_settled_are_free(gateway: Gateway, settled_status: str) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream: Final = scripted_upstream(scenario, background_routes("queued"))
response_id: Final = create_background_response(scenario, key, upstream.model)
write_rows(
'UPDATE "LiteLLM_ManagedObjectTable" SET status = %s WHERE unified_object_id = %s',
(settled_status, response_id),
)
upstream.script(background_routes("completed"))
assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == free_reads("completed", 2)
assert_charged_exactly(lambda: key_spend(key), 0.0)
assert charged_rows(upstream.model_id) == []
FOREIGN_IDS: Final[Mapping[str, Callable[[Upstream], str]]] = {
"unknown": lambda _: f"resp_{uuid.uuid4().hex}",
"five_kilobytes": lambda _: "resp_" + "a" * 5120,
"raw_vendor_id": lambda upstream: f"resp_{upstream.model}",
}
@pytest.mark.timeout(240)
@pytest.mark.parametrize("foreign", FOREIGN_IDS.keys())
def test_a_read_of_an_id_the_proxy_never_issued_fails_without_tracking_or_billing_it(
gateway: Gateway, foreign: str
) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
response_id: Final = create_background_response(scenario, key, upstream.model)
foreign_id: Final = FOREIGN_IDS[foreign](upstream)
refused: Final = gateway.request("GET", f"/v1/responses/{foreign_id}", key=key)
assert refused.status_code >= 400, refused.text
assert managed_rows(foreign_id) == []
assert (read_response(gateway, key, response_id),) in (
settling_reads("completed", 1),
free_reads("completed", 1),
)
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert len(charged_rows(upstream.model_id)) == 1
@pytest.mark.timeout(240)
def test_a_background_create_answered_at_once_and_its_cached_twin_bill_the_create_alone(gateway: Gateway) -> None:
store: Final = redis_store()
with gateway.scenario() as scenario:
key: Final = scenario.key()
answered: Final = json_reply(response_body("completed"))
upstream: Final = scripted_upstream(scenario, scripted_routes(answered, answered))
log: Final = upstream_log(gateway, upstream.handle)
first: Final = create_background_response(scenario, key, upstream.model)
eventually(
lambda: entry_key(store, f"msg_{upstream.model}"),
lambda found: found is not None,
)
second: Final = create_background_response(scenario, key, upstream.model)
assert len(submissions(next(log), upstream.handle)) == 1
reads: Final = [read_response(gateway, key, response_id) for response_id in (first, second, first, second)]
assert tuple(reads) == free_reads("completed", 4)
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert len(charged_rows(upstream.model_id)) == 1
def _end_user_spend(end_user: str) -> float:
rows: Final = read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id = %s', (end_user,))
return number(rows[0]["spend"]) if rows else 0.0
def _charged_end_users(model_id: str) -> list[JsonValue]:
rows: Final = read_rows('SELECT end_user FROM "LiteLLM_SpendLogs" WHERE model_id = %s AND spend > 0', (model_id,))
return [row["end_user"] for row in rows]
@pytest.mark.timeout(240)
def test_a_background_response_bills_the_end_user_who_created_it_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
end_user: Final = f"audit-end-user-{uuid.uuid4().hex}"
scenario.cleanups.callback(gateway.request, "POST", "/end_user/delete", {"user_ids": [end_user]})
upstream: Final = scripted_upstream(scenario, background_routes("completed"))
response_id: Final = create_background_response(
scenario,
key,
upstream.model,
body={"model": upstream.model, "input": PROMPT, "background": True, "store": False, "user": end_user},
)
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
assert_charged_exactly(lambda: _end_user_spend(end_user), PRICE)
assert key_spend(key) == pytest.approx(PRICE)
assert _charged_end_users(upstream.model_id) == [end_user]
@pytest.mark.timeout(240)
def test_the_read_that_settles_a_background_response_splits_its_charge_into_input_and_output(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
_, response_id = _lose_the_row_then_finish(scenario, key)
settling: Final = gateway.request("GET", f"/v1/responses/{response_id}", key=key)
later: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS - 1))
assert read_of(settling) == Read("completed", "pong", pytest.approx(PRICE))
split: Final = tuple(float(settling.headers[f"x-litellm-response-cost-{part}"]) for part in ("input", "output"))
assert split == pytest.approx((INPUT_TOKENS * INPUT_COST_PER_TOKEN, OUTPUT_TOKENS * OUTPUT_COST_PER_TOKEN))
assert later == free_reads("completed", READS - 1)
assert_charged_exactly(lambda: key_spend(key), PRICE)
def _answered_at_once(scenario: Scenario, key: str) -> tuple[Upstream, str]:
answered: Final = json_reply(response_body("completed"))
upstream: Final = scripted_upstream(scenario, scripted_routes(answered, answered))
return upstream, create_background_response(scenario, key, upstream.model)
def _stored_output(response_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
"SELECT file_object->'output' AS output, position('pong' in file_object::text) > 0 AS holds_text "
'FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
(response_id,),
)
STORED_BY: Final[Mapping[str, Callable[[Scenario, str], tuple[Upstream, str]]]] = {
"the_create": _answered_at_once,
"a_read": _lose_the_row_then_finish,
}
@pytest.mark.timeout(240)
@pytest.mark.parametrize("stored_by", STORED_BY.keys())
def test_a_stored_background_response_keeps_none_of_the_model_output(gateway: Gateway, stored_by: str) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
_, response_id = STORED_BY[stored_by](scenario, key)
assert read_response(gateway, key, response_id).text == "pong"
assert _stored_output(response_id) == [{"output": [], "holds_text": False}]
def _model_usage(upstream: Upstream) -> tuple[float, int]:
rows: Final = read_rows(
"SELECT COALESCE(SUM(spend), 0) AS spend, COALESCE(SUM(request_count), 0)::bigint AS requests "
'FROM "LiteLLM_DailyModelUsage" WHERE model_group IN (%s, %s)',
(upstream.model, upstream.model_id),
)
requests: Final = rows[0]["requests"]
assert isinstance(requests, int), rows
return number(rows[0]["spend"]), requests
@pytest.mark.timeout(240)
def test_model_usage_counts_a_background_response_once_in_spend_and_only_client_calls_as_requests(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
upstream, response_id = _finished(scenario, key)
reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS))
assert reads in (settling_reads("completed", READS), free_reads("completed", READS))
assert_charged_exactly(lambda: key_spend(key), PRICE)
client_calls: Final = 1 + READS
held(
lambda: _model_usage(upstream),
lambda usage: costs(PRICE)(usage[0]) and usage[1] == client_calls,
holding=POLL_WINDOW_SECONDS,
seconds=POLL_WINDOW_SECONDS + SPEND_LANDING_SECONDS,
)
BASE_MODEL: Final = "gpt-4o-mini"
BASE_MODEL_ONLY: Final[Mapping[str, JsonValue]] = {"base_model": BASE_MODEL}
BASE_MODEL_IN: Final[Mapping[str, tuple[Mapping[str, JsonValue] | None, Mapping[str, JsonValue]]]] = {
"model_info": (BASE_MODEL_ONLY, {}),
"litellm_params": (None, BASE_MODEL_ONLY),
}
def _unpriced_upstream(
scenario: Scenario,
unpriced: str,
model_info: Mapping[str, JsonValue] | None,
litellm_params: Mapping[str, JsonValue],
) -> Upstream:
foreground: Final = json_reply(response_body("completed", background=False, model=unpriced))
return scripted_upstream(
scenario,
scripted_routes(foreground, foreground),
model_info=model_info,
model=f"openai/{unpriced}",
input_cost_per_token=None,
output_cost_per_token=None,
**litellm_params,
)
@pytest.mark.timeout(240)
@pytest.mark.parametrize("base_model_in", BASE_MODEL_IN.keys())
def test_a_background_response_on_a_base_model_deployment_bills_what_a_foreground_response_costs(
gateway: Gateway, base_model_in: str
) -> None:
with gateway.scenario() as scenario:
foreground_key: Final = scenario.key()
background_key: Final = scenario.key()
unpriced: Final = f"audit-unpriced-{uuid.uuid4().hex[:12]}"
reference: Final = _unpriced_upstream(scenario, unpriced, BASE_MODEL_ONLY, {})
created: Final = gateway.request(
"POST", "/v1/responses", {"model": reference.model, "input": PROMPT, "store": False}, key=foreground_key
)
foreground_cost: Final = read_of(created).cost
assert foreground_cost is not None and foreground_cost > 0, created.headers
upstream: Final = _unpriced_upstream(scenario, unpriced, *BASE_MODEL_IN[base_model_in])
upstream.script(background_routes("completed", model=unpriced))
response_id: Final = create_background_response(scenario, background_key, upstream.model)
reads: Final = tuple(read_response(gateway, background_key, response_id) for _ in range(READS))
settled: Final = (
Read("completed", "pong", pytest.approx(foreground_cost)),
*free_reads("completed", READS - 1),
)
assert reads in (settled, free_reads("completed", READS))
assert_charged_exactly(lambda: key_spend(background_key), foreground_cost)
background_rows: Final = [
row for row in charged_rows(upstream.model_id) if row["api_key"] == hashed(background_key)
]
assert [number(row["spend"]) for row in background_rows] == [pytest.approx(foreground_cost)]
@pytest.mark.timeout(240)
def test_a_client_cannot_make_a_background_create_free_by_forging_the_managed_marker(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
answered: Final = json_reply(response_body("completed"))
upstream: Final = scripted_upstream(scenario, scripted_routes(answered, answered))
forged: Final[dict[str, JsonValue]] = {MANAGED_READ_MARKER: True}
created: Final = gateway.request(
"POST",
"/v1/responses",
{
"model": upstream.model,
"input": PROMPT,
"background": True,
"store": False,
"metadata": forged,
"litellm_metadata": forged,
},
key=key,
)
assert read_of(created) == Read("completed", "pong", pytest.approx(PRICE))
response_id: Final = string_value(object_value(created.json())["id"])
scenario.cleanups.callback(forget_row, response_id)
assert tuple(read_response(gateway, key, response_id) for _ in range(2)) == free_reads("completed", 2)
assert_charged_exactly(lambda: key_spend(key), PRICE)

569
tests/integration/spend/test_background_response_billing_proxies.py Normal file
View file

@ -0,0 +1,569 @@
import itertools
import uuid
from collections.abc import Iterator, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack, contextmanager
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, held, object_value, string_value
from integration._support.database import read_rows, scratch_database, write_rows
from integration._support.otlp_sink import Span, configure_sink, owned_sinks, recorded_spans
from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process
from integration._support.prometheus_series import Sample, scrape
from integration._support.timing import assert_not_delayed_by, timed
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import RoutedResponse, SseResponse
from integration.spend._background_responses import (
INPUT_COST_PER_TOKEN,
MANAGED_READ_MARKER,
ONE_POLL_SECONDS,
OUTPUT_COST_PER_TOKEN,
PRICE,
PROMPT,
PROVIDER_ID,
READS,
Read,
assert_charged_exactly,
background_routes,
costs,
create_background_response,
free_reads,
hashed,
json_reply,
key_spend,
managed_rows,
number,
read_response,
reported_cost,
response_body,
scripted_routes,
settling_reads,
streamed_reply,
)
from pydantic import JsonValue, TypeAdapter
CONFIGURATION: Final = TypeAdapter(dict[str, JsonValue])
OWNED_PROXY_SECONDS: Final = 2 * graceful_stop_seconds() + 240
NO_DATABASE: Final = ("DATABASE_URL", "DATABASE_URL_READ_REPLICA")
POLL_OFF: Final = {"PROXY_BATCH_POLLING_ENABLED": "false"}
SINK_DELAY_SECONDS: Final = 4.0
SPEND_RETRY_SECONDS: Final = 60
BURST_RESPONSES: Final = 6
RESTART_RESPONSES: Final = 3
COOLDOWN_SECONDS: Final = 300
CHAT_COMPLETION: Final[dict[str, JsonValue]] = {
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
}
@contextmanager
def _scripted(response: RoutedResponse | SseResponse, count: int = 1) -> Iterator[tuple[ScenarioHandle, ...]]:
with ExitStack() as stack:
handles: Final = tuple(register_scenario(f"bgonce-{uuid.uuid4().hex[:12]}", response) for _ in range(count))
for handle in handles:
stack.callback(delete_scenario, handle)
yield handles
def _rescript(handles: Sequence[ScenarioHandle], response: RoutedResponse) -> None:
for handle in handles:
register_scenario(handle.scenario_id, response, control_url=handle.control_url)
def _with_chat(response: RoutedResponse) -> RoutedResponse:
return RoutedResponse(
content_type="application/x-routed",
routes={**response.routes, "POST /chat/completions": json_reply(CHAT_COMPLETION)},
)
def _deployment(handle: ScenarioHandle, model_info: Mapping[str, JsonValue] | None = None) -> dict[str, JsonValue]:
return {
"model_name": handle.scenario_id,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "sk-scripted-provider",
"api_base": handle.api_base(),
"input_cost_per_token": INPUT_COST_PER_TOKEN,
"output_cost_per_token": OUTPUT_COST_PER_TOKEN,
},
**({"model_info": dict(model_info)} if model_info else {}),
}
def _write_config(directory: Path, config: Mapping[str, JsonValue]) -> Path:
path: Final = directory / f"background-billing-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(dict(config)))
return path
def _config(
directory: Path,
handles: Sequence[ScenarioHandle],
*,
model_info: Mapping[str, JsonValue] | None = None,
**sections: Mapping[str, JsonValue],
) -> Path:
shared: Final = CONFIGURATION.validate_python(
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
)
settings: Final = {
name: {**object_value(shared[name]), **sections.get(name, {})}
for name in ("general_settings", "litellm_settings", "router_settings")
}
extra: Final = {name: dict(section) for name, section in sections.items() if name not in settings}
return _write_config(
directory,
{"model_list": [_deployment(handle, model_info) for handle in handles], **settings, **extra},
)
@contextmanager
def _scratch_proxy(
gateway: Gateway,
directory: Path,
monkeypatch: pytest.MonkeyPatch,
config: Path,
overrides: Mapping[str, str] | None = None,
) -> Iterator[Gateway]:
with scratch_database() as database_url:
monkeypatch.setenv("DATABASE_URL", database_url)
with owned_proxy(
gateway, directory, {"DATABASE_URL": database_url, **(overrides or {})}, config=config
) as proxy:
yield proxy
def _team_member_key(scenario: Scenario) -> str:
team_id: Final = scenario.team()
return scenario.key(team_id=team_id, user_id=scenario.member(team_id))
def _billed_rows_of(key: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND spend > 0', (hashed(key),)
)
def _assert_each_billed_once_as_background_cost(key: str, count: int) -> None:
rows: Final = _billed_rows_of(key)
assert len(rows) == count, rows
assert all(string_value(row["request_id"]).endswith("_background_cost") for row in rows), rows
assert len({row["request_id"] for row in rows}) == count, rows
assert [number(row["spend"]) for row in rows] == pytest.approx([PRICE] * count)
@pytest.mark.timeout(OWNED_PROXY_SECONDS)
def test_a_proxy_without_a_database_bills_every_read_of_a_finished_background_response(
gateway: Gateway, tmp_path: Path
) -> None:
forged: Final[dict[str, JsonValue]] = {
"internal_call_origin": "background_response_status_poll",
MANAGED_READ_MARKER: True,
}
with _scripted(background_routes("completed")) as (handle,):
config: Final = _write_config(
tmp_path,
{
"model_list": [_deployment(handle)],
"general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"},
},
)
overrides: Final = {"STORE_MODEL_IN_DB": "False", "DISABLE_SCHEMA_UPDATE": "true"}
with owned_proxy(
gateway, tmp_path, overrides, config=config, database_setup=(), remove_environment=NO_DATABASE
) as proxy:
created: Final = proxy.post(
"/v1/responses", {"model": handle.scenario_id, "input": PROMPT, "background": True, "store": False}
)
response_id: Final = string_value(created["id"])
reads: Final = (
read_response(proxy, proxy.key, response_id),
read_response(proxy, proxy.key, response_id),
read_response(proxy, proxy.key, response_id, {"metadata": forged, "litellm_metadata": forged}),
)
assert [(read.status, read.text) for read in reads] == [("completed", "pong")] * 3
assert [read.cost for read in reads] == pytest.approx([PRICE] * 3)
def _span_cost(span: Span) -> float:
return number(span["attributes"].get("gen_ai.cost.total_cost", 0))
def _spans_of(sink: str, key: str) -> tuple[Span, ...]:
_, spans = recorded_spans(sink)
return tuple(span for span in spans if span["attributes"].get("metadata.user_api_key_hash") == hashed(key))
@pytest.mark.timeout(OWNED_PROXY_SECONDS + 120)
def test_with_the_cost_poll_off_the_first_finished_read_bills_the_creator_once_and_tells_the_destination_once(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
with _scripted(background_routes("queued")) as (handle,), owned_sinks(tmp_path / "sinks") as sinks:
config: Final = _config(
tmp_path,
(handle,),
litellm_settings={"callbacks": ["otel"]},
callback_settings={
"otel": {"exporter": "http/json", "endpoint": sinks.operator, "use_simple_processor": True}
},
)
overrides: Final = {**POLL_OFF, "OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": sinks.operator}
with (
_scratch_proxy(gateway, tmp_path, monkeypatch, config, overrides) as proxy,
proxy.scenario() as scenario,
):
key: Final = _team_member_key(scenario)
response_id: Final = create_background_response(scenario, key, handle.scenario_id)
unfinished: Final = read_response(proxy, key, response_id)
_rescript((handle,), background_routes("completed"))
finished: Final = tuple(read_response(proxy, key, response_id) for _ in range(READS))
assert (unfinished.status, unfinished.cost) == ("queued", 0.0)
assert finished == settling_reads("completed", READS)
assert_charged_exactly(lambda: key_spend(key), PRICE)
_assert_each_billed_once_as_background_cost(key, 1)
spans: Final = eventually(
lambda: _spans_of(sinks.operator, key), lambda found: len(found) >= READS + 3, seconds=30
)
assert [_span_cost(span) for span in spans if _span_cost(span) > 0] == pytest.approx([PRICE])
@pytest.mark.timeout(OWNED_PROXY_SECONDS + 60)
def test_the_read_that_finishes_a_background_response_does_not_wait_on_a_slow_spend_destination(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
with _scripted(background_routes("queued")) as (handle,), owned_sinks(tmp_path / "sinks") as sinks:
config: Final = _config(tmp_path, (handle,), litellm_settings={"callbacks": ["openmeter"]})
overrides: Final = {
**POLL_OFF,
"OPENMETER_API_KEY": "om-scripted-key",
"OPENMETER_API_ENDPOINT": sinks.tenant,
}
with (
_scratch_proxy(gateway, tmp_path, monkeypatch, config, overrides) as proxy,
proxy.scenario() as scenario,
):
key: Final = scenario.key()
response_id: Final = create_background_response(scenario, key, handle.scenario_id)
configure_sink(sinks.tenant, delay_seconds=SINK_DELAY_SECONDS)
unfinished: Final = timed(lambda: read_response(proxy, key, response_id))
_rescript((handle,), background_routes("completed"))
first_finished: Final = timed(lambda: read_response(proxy, key, response_id))
later: Final = tuple(timed(lambda: read_response(proxy, key, response_id)) for _ in range(2))
assert first_finished.value.status == "completed"
assert_not_delayed_by(
(first_finished.seconds,),
control=(unfinished.seconds, *(read.seconds for read in later)),
injected_delay=SINK_DELAY_SECONDS,
)
def _provider_404() -> RoutedResponse:
gone: Final = json_reply(
{
"error": {
"message": f"Response with id '{PROVIDER_ID}' not found.",
"type": "invalid_request_error",
"param": None,
"code": None,
}
},
404,
)
return scripted_routes(json_reply(response_body("queued")), gone)
def _create_status(proxy: Gateway, key: str, model: str) -> int:
return proxy.request(
"POST", "/v1/responses", {"model": model, "input": PROMPT, "background": True, "store": False}, key=key
).status_code
@pytest.mark.timeout(OWNED_PROXY_SECONDS + ONE_POLL_SECONDS)
def test_neither_a_poll_404_nor_a_client_read_404_takes_the_deployment_out_of_service(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
with ExitStack() as stack:
(handle,) = stack.enter_context(_scripted(_provider_404()))
config: Final = _config(
tmp_path,
(handle,),
model_info={"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0, "NotFoundErrorAllowedFails": 0}},
router_settings={"disable_cooldowns": False, "cooldown_time": COOLDOWN_SECONDS},
)
proxy: Final = stack.enter_context(_scratch_proxy(gateway, tmp_path, monkeypatch, config))
scenario: Final = stack.enter_context(proxy.scenario())
key: Final = scenario.key()
polled: Final = create_background_response(scenario, key, handle.scenario_id)
eventually(
lambda: [row["status"] for row in managed_rows(polled)],
lambda statuses: statuses == ["stale_expired"],
seconds=ONE_POLL_SECONDS,
)
after_poll: Final = _create_status(proxy, key, handle.scenario_id)
client_read: Final = proxy.request("GET", f"/v1/responses/{polled}", key=key).status_code
after_client_read: Final = _create_status(proxy, key, handle.scenario_id)
assert (after_poll, client_read, after_client_read) == (200, 404, 200)
def _burst_call(proxy: Gateway, call: tuple[str, str, str]) -> int | None:
kind, key, target = call
try:
if kind == "read":
return proxy.request("GET", f"/v1/responses/{target}", key=key).status_code
return proxy.request(
"POST",
"/v1/chat/completions",
{"model": target, "messages": [{"role": "user", "content": PROMPT}]},
key=key,
).status_code
except httpx.TransportError:
return None
@pytest.mark.timeout(OWNED_PROXY_SECONDS + 2 * graceful_stop_seconds() + 60 + ONE_POLL_SECONDS + 100)
def test_background_responses_read_across_a_killed_worker_are_each_billed_once_to_their_creator(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
with (
_scripted(_with_chat(background_routes("queued")), BURST_RESPONSES) as handles,
scratch_database() as database_url,
):
monkeypatch.setenv("DATABASE_URL", database_url)
config: Final = _config(tmp_path, handles)
with (
owned_proxy_process(gateway, tmp_path, {"DATABASE_URL": database_url}, config=config, workers=2) as owned,
owned.gateway.scenario() as scenario,
):
key: Final = _team_member_key(scenario)
chat_key: Final = scenario.key()
response_ids: Final = tuple(
create_background_response(scenario, key, handle.scenario_id) for handle in handles
)
reads: Final = tuple(("read", key, response_id) for response_id in response_ids * 3)
chats: Final = tuple(("chat", chat_key, handle.scenario_id) for handle in handles * 2)
calls: Final = tuple(
call for call in itertools.chain.from_iterable(itertools.zip_longest(reads, chats)) if call is not None
)
with owned.replacing_worker(min(owned.worker_pids())) as kill, ThreadPoolExecutor(max_workers=10) as pool:
kill.freeze()
futures: Final = tuple(pool.submit(_burst_call, owned.gateway, call) for call in calls)
eventually(
lambda: sum(future.done() for future in futures),
lambda done: done >= len(futures) // 2,
seconds=30,
)
kill()
outcomes: Final = tuple(future.result() for future in futures)
assert set(outcomes) <= {200, None}, outcomes
assert outcomes.count(200) >= len(outcomes) // 2, outcomes
_rescript(handles, _with_chat(background_routes("completed")))
finished: Final = tuple(read_response(owned.gateway, key, response_id) for response_id in response_ids * 2)
assert [read.status for read in finished] == ["completed"] * 2 * BURST_RESPONSES
assert [read.cost in (0.0, pytest.approx(PRICE)) for read in finished] == [True] * 2 * BURST_RESPONSES
assert sum(read.cost or 0.0 for read in finished) <= BURST_RESPONSES * PRICE + 1e-9
assert_charged_exactly(lambda: key_spend(key), BURST_RESPONSES * PRICE)
_assert_each_billed_once_as_background_cost(key, BURST_RESPONSES)
def _statuses(response_ids: Sequence[str]) -> list[JsonValue]:
rows: Final = itertools.chain.from_iterable(managed_rows(response_id) for response_id in response_ids)
return [row["status"] for row in rows]
def _submit(proxy: Gateway, key: str, model: str) -> str:
response: Final = proxy.request(
"POST", "/v1/responses", {"model": model, "input": PROMPT, "background": True, "store": False}, key=key
)
assert response.status_code == 200, response.text
return string_value(object_value(response.json())["id"])
@pytest.mark.timeout(2 * OWNED_PROXY_SECONDS + ONE_POLL_SECONDS + 100)
def test_background_responses_left_pending_across_a_proxy_restart_are_billed_once_to_their_creator(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
with _scripted(background_routes("queued"), RESTART_RESPONSES) as handles, scratch_database() as database_url:
monkeypatch.setenv("DATABASE_URL", database_url)
config: Final = _config(tmp_path, handles)
with owned_proxy(
gateway, tmp_path / "before", {"DATABASE_URL": database_url, **POLL_OFF}, config=config
) as first:
key: Final = string_value(first.post("/key/generate", {})["key"])
response_ids: Final = tuple(_submit(first, key, handle.scenario_id) for handle in handles)
_rescript(handles, background_routes("completed"))
with owned_proxy(gateway, tmp_path / "after", {"DATABASE_URL": database_url}, config=config):
eventually(
lambda: _statuses(response_ids),
lambda statuses: statuses == ["completed"] * RESTART_RESPONSES,
seconds=ONE_POLL_SECONDS,
)
assert_charged_exactly(lambda: key_spend(key), RESTART_RESPONSES * PRICE)
_assert_each_billed_once_as_background_cost(key, RESTART_RESPONSES)
DELETED: Final[dict[str, JsonValue]] = {"id": PROVIDER_ID, "object": "response.deleted", "deleted": True}
REFUSAL_HOLD_SECONDS: Final = 10
def _deletable(read_status: str) -> RoutedResponse:
return RoutedResponse(
content_type="application/x-routed",
routes={**background_routes(read_status).routes, f"DELETE /responses/{PROVIDER_ID}": json_reply(DELETED)},
)
@pytest.mark.timeout(OWNED_PROXY_SECONDS + 120)
def test_a_delete_bills_a_finished_background_response_once_and_only_for_a_caller_that_may_address_it(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
with _scripted(_deletable("queued")) as (handle,):
config: Final = _config(tmp_path, (handle,))
with (
_scratch_proxy(gateway, tmp_path, monkeypatch, config, POLL_OFF) as proxy,
proxy.scenario() as scenario,
):
creator: Final = _team_member_key(scenario)
stranger: Final = _team_member_key(scenario)
response_id: Final = create_background_response(scenario, creator, handle.scenario_id)
_rescript((handle,), _deletable("completed"))
refused: Final = proxy.request("DELETE", f"/v1/responses/{response_id}", key=stranger)
assert refused.status_code == 403, refused.text
held(lambda: key_spend(creator), costs(0.0), holding=REFUSAL_HOLD_SECONDS)
assert [row["status"] for row in managed_rows(response_id)] == ["queued"]
deleted: Final = proxy.request("DELETE", f"/v1/responses/{response_id}", key=creator)
assert deleted.status_code == 200, deleted.text
assert object_value(deleted.json())["deleted"] is True
assert reported_cost(deleted.headers) == pytest.approx(PRICE)
assert_charged_exactly(lambda: key_spend(creator), PRICE)
_assert_each_billed_once_as_background_cost(creator, 1)
assert key_spend(stranger) == 0
assert [row["status"] for row in managed_rows(response_id)] == ["completed"]
@contextmanager
def _billing_claims_refused() -> Iterator[None]:
write_rows(
"CREATE FUNCTION audit_refuse_billing_claim() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN "
"IF NEW.status = 'completed' AND OLD.status IS DISTINCT FROM 'completed' THEN "
"RAISE EXCEPTION 'audit refuses the billing claim'; END IF; RETURN NEW; END $$",
(),
)
write_rows(
'CREATE TRIGGER audit_refuse_billing_claim BEFORE UPDATE ON "LiteLLM_ManagedObjectTable" '
"FOR EACH ROW EXECUTE FUNCTION audit_refuse_billing_claim()",
(),
)
try:
yield
finally:
write_rows('DROP TRIGGER audit_refuse_billing_claim ON "LiteLLM_ManagedObjectTable"', ())
write_rows("DROP FUNCTION audit_refuse_billing_claim()", ())
@pytest.mark.timeout(OWNED_PROXY_SECONDS + 180)
def test_a_finished_read_whose_bill_fails_to_settle_bills_the_creator_once_after_the_database_recovers(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
with _scripted(background_routes("queued")) as (handle,):
config: Final = _config(tmp_path, (handle,))
with (
_scratch_proxy(gateway, tmp_path, monkeypatch, config, POLL_OFF) as proxy,
proxy.scenario() as scenario,
):
key: Final = _team_member_key(scenario)
response_id: Final = create_background_response(scenario, key, handle.scenario_id)
_rescript((handle,), background_routes("completed"))
with _billing_claims_refused():
unsettled: Final = read_response(proxy, key, response_id)
statuses: Final = [row["status"] for row in managed_rows(response_id)]
assert unsettled == Read("completed", "pong", 0.0)
assert statuses == ["queued"]
eventually(lambda: key_spend(key), costs(PRICE), seconds=SPEND_RETRY_SECONDS)
assert tuple(read_response(proxy, key, response_id) for _ in range(READS)) == free_reads("completed", READS)
assert_charged_exactly(lambda: key_spend(key), PRICE)
_assert_each_billed_once_as_background_cost(key, 1)
assert [row["status"] for row in managed_rows(response_id)] == ["completed"]
REMAINING_KEY_BUDGET: Final = "litellm_remaining_api_key_budget_metric"
KEY_BUDGET: Final = 10.0
def _remaining_key_budget(samples: Sequence[Sample], key: str) -> list[float]:
return [
sample.value
for sample in samples
if sample.name == REMAINING_KEY_BUDGET and sample.labels.get("hashed_api_key") == hashed(key)
]
@pytest.mark.timeout(OWNED_PROXY_SECONDS + ONE_POLL_SECONDS + 60)
def test_a_background_response_the_poll_bills_lowers_its_creating_keys_remaining_budget_by_its_price(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
with _scripted(background_routes("completed")) as (handle,):
config: Final = _config(tmp_path, (handle,), litellm_settings={"callbacks": ["prometheus"]})
with _scratch_proxy(gateway, tmp_path, monkeypatch, config) as proxy, proxy.scenario() as scenario:
key: Final = scenario.key(max_budget=KEY_BUDGET)
response_id: Final = create_background_response(scenario, key, handle.scenario_id)
eventually(
lambda: [row["status"] for row in managed_rows(response_id)],
lambda statuses: statuses == ["completed"],
seconds=ONE_POLL_SECONDS,
)
assert_charged_exactly(lambda: key_spend(key), PRICE)
eventually(
lambda: _remaining_key_budget(scrape(proxy), key),
lambda values: values == [pytest.approx(KEY_BUDGET - PRICE)],
seconds=30,
)
CACHE_POLLING: Final[dict[str, JsonValue]] = {"responses": {"background_mode": {"polling_via_cache": "all"}}}
CACHE_POLLED_CREATES: Final[Mapping[str, tuple[dict[str, JsonValue], bool]]] = {
"streamed_with_background_in_extra_body": ({"stream": True, "extra_body": {"background": True}}, True),
"unstreamed": ({}, False),
}
def _cache_polled_create(proxy: Gateway, key: str, model: str, create: Mapping[str, JsonValue]) -> str:
created: Final = proxy.request(
"POST",
"/v1/responses",
{"model": model, "input": PROMPT, "background": True, "store": False, **create},
key=key,
)
assert created.status_code == 200, created.text
return string_value(object_value(created.json())["id"])
@pytest.mark.timeout(OWNED_PROXY_SECONDS + 120)
@pytest.mark.parametrize("create", CACHE_POLLED_CREATES.keys())
def test_a_background_response_the_proxy_polls_through_its_cache_bills_its_creator_once(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, create: str
) -> None:
body, provider_sees_background = CACHE_POLLED_CREATES[create]
with _scripted(streamed_reply(background=provider_sees_background)) as (handle,):
config: Final = _config(tmp_path, (handle,), litellm_settings=CACHE_POLLING)
with _scratch_proxy(gateway, tmp_path, monkeypatch, config, POLL_OFF) as proxy, proxy.scenario() as scenario:
key: Final = _team_member_key(scenario)
polling_id: Final = _cache_polled_create(proxy, key, handle.scenario_id, body)
finished: Final = eventually(
lambda: read_response(proxy, key, polling_id),
lambda read: read.status == "completed",
seconds=SPEND_RETRY_SECONDS,
)
assert finished.text == "pong"
assert_charged_exactly(lambda: key_spend(key), PRICE)
assert [read_response(proxy, key, polling_id).text for _ in range(READS)] == ["pong"] * READS
assert_charged_exactly(lambda: key_spend(key), PRICE)

29
tests/integration/spend/test_background_response_poll_retirement.py
View file

@ -16,7 +16,7 @@ from integration.cost_calculation.cost_tracking_case import JsonResponse, Routed
from pydantic import JsonValue
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.proxy.spend_tracking.spend_tracking_utils import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
INPUT_COST_PER_TOKEN: Final = 0.001
@ -240,30 +240,21 @@ def test_a_provider_error_other_than_404_keeps_the_row_queued_for_retry(gateway:
assert _row_status(unified_id) == [{"status": "queued"}]
def _polled_provider_id(row: dict[str, JsonValue]) -> str:
return ResponsesAPIRequestUtils.decode_responses_api_response_id(string_value(row["request_id"]))["response_id"]
def _poll_spend_rows(handle: ScenarioHandle) -> list[dict[str, JsonValue]]:
billed: Final = read_rows(
'SELECT request_id, call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
"WHERE metadata->>'internal_call_origin' = %s",
(BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,),
def _poll_spend_rows(unified_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
"WHERE request_id = %s AND metadata->>'internal_call_origin' = %s",
(f"{unified_id}{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}", BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN),
)
return [
{column: value for column, value in row.items() if column != "request_id"}
for row in billed
if _polled_provider_id(row) == f"resp_{handle.scenario_id}"
]
@pytest.mark.timeout(180)
def test_a_completed_response_is_marked_completed_and_its_poll_is_billed(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model, handle = _scripted_deployment(scenario, "done", _routes("queued", _json(_response_body("completed"))))
model, _ = _scripted_deployment(scenario, "done", _routes("queued", _json(_response_body("completed"))))
unified_id: Final = _submit_background_response(scenario, scenario.key(), model)
_await_status(unified_id, "completed", ONE_POLL_SECONDS)
spend_rows: Final = eventually(lambda: _poll_spend_rows(handle), lambda rows: len(rows) >= 1, seconds=70)
spend_rows: Final = eventually(lambda: _poll_spend_rows(unified_id), lambda rows: len(rows) >= 1, seconds=70)
assert spend_rows[0] == {
"call_type": "aget_responses",
"status": "success",
@ -305,7 +296,7 @@ def test_responses_gone_at_the_provider_do_not_starve_a_newer_response_out_of_co
)
for index in range(MAX_OBJECTS_PER_POLL_CYCLE)
)
completable_model, completable = _scripted_deployment(
completable_model, _ = _scripted_deployment(
scenario, "done", _routes("queued", _json(_response_body("completed")))
)
completable_id: Final = _submit_background_response(scenario, key, completable_model)
@ -313,7 +304,7 @@ def test_responses_gone_at_the_provider_do_not_starve_a_newer_response_out_of_co
assert [_row_status(gone_id) for gone_id in gone_ids] == [
[{"status": "stale_expired"}]
] * MAX_OBJECTS_PER_POLL_CYCLE
spend_rows: Final = eventually(lambda: _poll_spend_rows(completable), lambda rows: len(rows) >= 1, seconds=70)
spend_rows: Final = eventually(lambda: _poll_spend_rows(completable_id), lambda rows: len(rows) >= 1, seconds=70)
assert spend_rows[0]["spend"] == pytest.approx(
INPUT_TOKENS * INPUT_COST_PER_TOKEN + OUTPUT_TOKENS * OUTPUT_COST_PER_TOKEN
)

100
tests/unit/enterprise/proxy/hooks/test_managed_files.py
View file

@ -6048,3 +6048,103 @@ async def test_create_from_raw_provider_file_id_skips_managed_ownership_lookup(c
assert isinstance(result, dict) and result[file_field] == "file-raw-provider-id", result
assert not any("unified_file_id" in where for where in table.find_first_calls), table.find_first_calls
_MANAGED_READ_SALT_KEY: Final = "managed-read-unit-test-salt-key"
def _proxy_issued_response_id(encrypted: bool) -> str:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIResponse
routed_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id(
"openai", "deployment-1", "resp_provider_job"
)
if not encrypted:
return routed_id
issuer: Final = ResponsesIDSecurity(
general_settings_reader=lambda: {}, signing_key_reader=lambda: _MANAGED_READ_SALT_KEY
)
return issuer._encrypt_response_id(
ResponsesAPIResponse(id=routed_id, created_at=1, output=[], status="queued"),
UserAPIKeyAuth(user_id="alice"),
).id
@pytest.mark.asyncio
@pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"])
@pytest.mark.parametrize("encrypted", [False, True], ids=["after_id_security", "before_id_security"])
async def test_a_read_of_a_proxy_issued_response_id_is_marked_managed(
call_type: str, encrypted: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The cost poll bills a background response the proxy issued, so a client read of that id is
marked managed and priced at zero, whichever order the ID security hook runs in."""
from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
from litellm.proxy._types import UserAPIKeyAuth
monkeypatch.setenv("LITELLM_SALT_KEY", _MANAGED_READ_SALT_KEY)
addressed_id: Final = _proxy_issued_response_id(encrypted)
result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
cache=MagicMock(),
data={"response_id": addressed_id, "litellm_metadata": {"user_api_key_user_id": "alice"}},
call_type=call_type,
)
assert result == {
"response_id": addressed_id,
"litellm_metadata": {"user_api_key_user_id": "alice", MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True},
}
@pytest.mark.asyncio
async def test_a_read_of_a_raw_provider_response_id_is_not_marked_managed() -> None:
"""A raw provider id names no deployment, so the proxy never stored that response and its read
keeps pricing normally."""
from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
from litellm.proxy._types import UserAPIKeyAuth
result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
cache=MagicMock(),
data={"response_id": "resp_raw_provider_job", "metadata": {}},
call_type="aget_responses",
)
assert isinstance(result, dict)
assert MANAGED_BACKGROUND_RESPONSE_METADATA_KEY not in result["metadata"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"create_fields, marked",
[
({"stream": True, "background": True}, True),
({"stream": True, "extra_body": {"background": True}}, True),
({"stream": False, "background": True}, False),
({"stream": True}, False),
],
ids=["streamed_background", "streamed_background_in_extra_body", "not_streamed", "not_background"],
)
async def test_only_a_streamed_background_create_is_marked_for_the_managed_bill(
create_fields: Mapping[str, object], marked: bool
) -> None:
"""A streamed background create is billed once through its managed row when its stream ends,
so only that create is marked free, and only in litellm_metadata, which never reaches the
provider."""
from litellm.constants import MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
from litellm.proxy._types import UserAPIKeyAuth
result: Final = await PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()).async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
cache=MagicMock(),
data={"model": "gpt-5.6", "input": "hello", "metadata": {"run": "nightly"}, "litellm_metadata": {}, **create_fields},
call_type="aresponses",
)
assert isinstance(result, dict)
assert (MANAGED_BACKGROUND_RESPONSE_METADATA_KEY in result["litellm_metadata"]) is marked
assert result["metadata"] == {"run": "nightly"}

56
tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py
View file

@ -553,3 +553,59 @@ async def test_store_unified_object_id_omits_unset_columns():
assert "api_key" not in create_data
assert "request_tags" not in create_data
assert "unified-b" in store
@pytest.mark.asyncio
async def test_store_unified_object_id_drops_nul_characters_postgres_cannot_store():
"""Regression (spend loss): Postgres rejects a NUL inside jsonb, and a client can send one
in a background response's instructions or tags. The row the cost poll bills from must
still be written, with only the NUL characters removed."""
from litellm.types.llms.openai import ResponsesAPIResponse
instance, store = _in_memory_managed_files()
creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice")
response = ResponsesAPIResponse(
id="resp_nul",
created_at=1,
status="queued",
model="gpt-5.6",
instructions="be brief\x00",
metadata={"note\x00": "keep\x00 this"},
output=[],
)
await instance.store_unified_object_id(
unified_object_id="unified-r",
file_object=response,
litellm_parent_otel_span=None,
model_object_id="resp_nul",
file_purpose="response",
user_api_key_dict=creator,
request_tags=["env\x00:prod"],
persist_attribution=True,
)
row = store["unified-r"]
assert "\\u0000" not in row["file_object"]
stored = json.loads(row["file_object"])
assert stored == json.loads(response.model_dump_json().replace("\\u0000", ""))
assert stored["instructions"] == "be brief"
assert stored["metadata"] == {"note": "keep this"}
assert row["request_tags"].data == ["env:prod"]
@pytest.mark.asyncio
async def test_store_unified_object_id_keeps_the_provider_json_when_it_has_no_nul():
instance, store = _in_memory_managed_files()
batch = _build_batch_response(batch_id="b", status="validating")
await instance.store_unified_object_id(
unified_object_id="unified-b",
file_object=batch,
litellm_parent_otel_span=None,
model_object_id="b",
file_purpose="batch",
user_api_key_dict=UserAPIKeyAuth(user_id="alice", api_key="hash-alice"),
)
assert store["unified-b"]["file_object"] == batch.model_dump_json()

45
tests/unit/enterprise/proxy/test_managed_files_hook.py
View file

@ -1499,6 +1499,51 @@ async def test_create_still_upserts_and_claims_attribution():
assert created["team_id"] == "team-alpha"
@pytest.mark.asyncio
async def test_a_stored_background_response_names_its_creators_end_user_and_holds_none_of_its_output():
"""Billing reads only a stored response's status and model, so neither the row nor its cache
entry keeps what the model generated, and the row names the end user the create billed to."""
from litellm.types.llms.openai import ResponsesAPIResponse
managed_files, mock_prisma = _make_object_store_instance()
generated_text = "the generated answer"
response = ResponsesAPIResponse(
id="resp_unified",
object="response",
status="completed",
created_at=1,
model="gpt-5.6",
output=[
{
"type": "message",
"id": "msg_1",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": generated_text, "annotations": []}],
}
],
background=True,
)
await managed_files.store_unified_object_id(
unified_object_id="resp_unified",
file_object=response,
litellm_parent_otel_span=None,
model_object_id="resp_provider",
file_purpose="response",
user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-creator", end_user_id="customer-42", parent_otel_span=None),
persist_attribution=True,
)
written = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]
cached = managed_files.internal_usage_cache.async_set_cache.await_args.kwargs["value"]
assert generated_text not in json.dumps(written, default=str)
assert generated_text not in json.dumps(cached, default=str)
assert json.loads(written["create"]["file_object"])["status"] == "completed"
assert written["create"]["end_user_id"] == "customer-42"
assert generated_text in response.model_dump_json()
@pytest.mark.asyncio
async def test_default_callers_still_create_their_rows():
"""create_if_missing defaults to True, so the fine-tune, Responses and managed

10
tests/unit/integrations/test_responses_background_cost.py
View file

@ -360,14 +360,14 @@ class TestCheckResponsesCost:
mock_job.id = "job-123"
mock_job.unified_object_id = "resp_test_id"
mock_job.created_by = "test-user"
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_id"}
# Mock find_many to return the job
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
# Mock update_many
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
# Create a completed response
completed_response = ResponsesAPIResponse(
@ -406,7 +406,10 @@ class TestCheckResponsesCost:
if c.kwargs.get("where", {}).get("id") is not None
]
assert len(completion_calls) == 1
assert completion_calls[0].kwargs["where"]["id"]["in"] == ["job-123"]
assert completion_calls[0].kwargs["where"] == {
"id": "job-123",
"status": {"in": ["queued", "in_progress"]},
}
assert completion_calls[0].kwargs["data"]["status"] == "completed"
@pytest.mark.asyncio
@ -423,6 +426,7 @@ class TestCheckResponsesCost:
mock_job.id = "job-456"
mock_job.unified_object_id = "resp_failed"
mock_job.created_by = "test-user"
mock_job.file_object = {"model": "gpt-4o", "id": "resp_failed"}
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]

76
tests/unit/litellm_core_utils/test_internal_call_metadata.py
View file

@ -1,11 +1,17 @@
"""Unit tests for internal-call metadata forwarding: budget-reservation stripping and origin stamping."""
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
from litellm.litellm_core_utils.internal_call_metadata import (
forwarded_internal_call_metadata,
is_unbilled_non_inference_call,
is_unbilled_non_inference_call_from_params,
sanitized_forwardable_call_metadata,
)
from litellm.types.utils import SHADOW_EVAL_ROUTER_CALL_ORIGIN
from litellm.types.utils import (
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
SHADOW_EVAL_ROUTER_CALL_ORIGIN,
)
PARENT = {
"user_api_key": "sk-hash",
@ -119,3 +125,69 @@ class TestSubCallMetadataSanitization:
assert sanitized_auth.team_id == "team-1"
assert sanitized_auth.api_key == auth.api_key
assert auth.budget_reservation == {"reserved_cost": 1.0}
class TestNonInferenceCallBilling:
"""A read of a stored object replays the usage of the call that created it, so it is free. A
background create returns no usage, so its reads are where the job's spend shows up: the cost
poll's stamped read bills it, and any unmarked read (the SDK, a proxy with no database) bills
it too, because nothing else ever will."""
BACKGROUND = {"id": "resp_1", "background": True, "status": "completed"}
FOREGROUND = {"id": "resp_2", "background": False, "status": "completed"}
COST_POLL = {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}
STATUS_POLL = {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN}
MANAGED_READ = {"user_api_key": "sk-hash", MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True}
def test_unmarked_read_of_a_background_response_is_billed(self):
assert is_unbilled_non_inference_call("aget_responses", None, self.BACKGROUND) is False
assert is_unbilled_non_inference_call("aget_responses", {"user_api_key": "sk-hash"}, self.BACKGROUND) is False
assert is_unbilled_non_inference_call_from_params("aget_responses", {}, self.BACKGROUND) is False
def test_managed_read_and_status_poll_of_a_background_response_are_free(self):
assert is_unbilled_non_inference_call("aget_responses", self.MANAGED_READ, self.BACKGROUND) is True
assert is_unbilled_non_inference_call("acancel_responses", self.MANAGED_READ, self.BACKGROUND) is True
assert is_unbilled_non_inference_call("aget_responses", self.STATUS_POLL, self.BACKGROUND) is True
assert is_unbilled_non_inference_call_from_params(
"aget_responses", {"metadata": self.MANAGED_READ}, self.BACKGROUND
) is True
assert is_unbilled_non_inference_call_from_params(
"aget_responses", {"litellm_metadata": self.STATUS_POLL}, self.BACKGROUND
) is True
def test_a_marker_that_is_not_true_does_not_waive_the_charge(self):
not_true = {MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: "true"}
assert is_unbilled_non_inference_call("aget_responses", not_true, self.BACKGROUND) is False
def test_cost_poll_read_is_billed_even_when_marked(self):
assert is_unbilled_non_inference_call("aget_responses", self.COST_POLL, self.BACKGROUND) is False
assert is_unbilled_non_inference_call(
"aget_responses", {**self.MANAGED_READ, **self.COST_POLL}, self.BACKGROUND
) is False
assert is_unbilled_non_inference_call_from_params(
"aget_responses", {"litellm_metadata": self.COST_POLL}, self.FOREGROUND
) is False
def test_read_of_a_foreground_response_is_free_whatever_it_carries(self):
assert is_unbilled_non_inference_call("aget_responses", None, self.FOREGROUND) is True
assert is_unbilled_non_inference_call("aget_responses", {"user_api_key": "sk-hash"}, self.FOREGROUND) is True
assert is_unbilled_non_inference_call(
"aget_responses", {INTERNAL_CALL_ORIGIN_METADATA_KEY: SHADOW_EVAL_ROUTER_CALL_ORIGIN}, self.FOREGROUND
) is True
def test_inference_calls_are_billed_unless_the_proxy_tracks_their_background_response(self):
assert is_unbilled_non_inference_call("aresponses", None, self.BACKGROUND) is False
assert is_unbilled_non_inference_call("aresponses", self.MANAGED_READ, self.FOREGROUND) is False
assert is_unbilled_non_inference_call_from_params("aresponses", {}, self.BACKGROUND) is False
assert is_unbilled_non_inference_call_from_params("acompletion", {"metadata": self.MANAGED_READ}, None) is False
def test_a_tracked_streamed_background_create_is_free_in_every_shape_its_stream_logs(self):
from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse
finished = ResponsesAPIResponse(**{**self.BACKGROUND, "created_at": 1, "output": []})
completed_event = ResponseCompletedEvent(type="response.completed", response=finished)
for logged in (finished, completed_event, completed_event.model_dump(), self.BACKGROUND):
assert is_unbilled_non_inference_call("aresponses", self.MANAGED_READ, logged) is True
assert is_unbilled_non_inference_call_from_params(
"aresponses", {"litellm_metadata": self.MANAGED_READ}, logged
) is True

15
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -7071,6 +7071,21 @@ class TestNonInferenceCallTypesAreNotBilled:
assert payload is not None
assert payload["total_tokens"] == 6000
@pytest.mark.parametrize(
"litellm_metadata",
[
{"_managed_background_response": True},
{"internal_call_origin": "background_response_status_poll"},
],
ids=["managed_read", "status_poll"],
)
def test_managed_read_of_a_background_response_is_free(self, litellm_metadata):
"""The proxy's cost poll bills a background response it tracks once, so a read the managed files
hook marks, or the poll's own status check, must not bill the same tokens again."""
logging_obj = self._logging_obj("aget_responses", litellm_metadata=litellm_metadata)
cost = logging_obj.response_cost_calculator(result=self._retrieved_response(background=True))
assert cost == 0.0
def test_reading_a_foreground_response_is_still_free(self):
"""Guards the test above against a blanket exemption: an explicit background=false read was
already billed by its create and must stay at zero."""

83
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -2366,17 +2366,21 @@ class TestBatchCostAttribution:
"""CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so
the batch-cost log is attributed like a non-batch request."""
def _instance(self, key_row=None, team_row=None, user_row=None):
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
def _prisma(self, key_row=None, team_row=None, user_row=None):
prisma = MagicMock()
prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_row)
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
return CheckBatchCost(
proxy_logging_obj=MagicMock(),
prisma_client=prisma,
llm_router=MagicMock(),
return prisma
def _instance(self, key_row=None, team_row=None, user_row=None, prisma=None):
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
)
return ManagedObjectCreatorAttribution(
prisma_client=prisma if prisma is not None else self._prisma(key_row, team_row, user_row),
poller_name="CheckBatchCost",
)
def _job(self, **overrides):
@ -2402,7 +2406,7 @@ class TestBatchCostAttribution:
user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert metadata["user_api_key"] == "hash-alice"
assert metadata["user_api_key_hash"] == "hash-alice"
@ -2419,7 +2423,7 @@ class TestBatchCostAttribution:
instance = self._instance()
job = self._job(api_key=None, request_tags=None)
metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
metadata = await instance.batch_metadata(job, "batch-1")
assert metadata["user_api_key"] is None
assert metadata["user_api_key_user_id"] == "alice"
@ -2432,15 +2436,32 @@ class TestBatchCostAttribution:
a None user_id) and the key hash still drives key-level attribution."""
from types import SimpleNamespace
instance = self._instance(key_row=SimpleNamespace(key_alias="svc-key"))
prisma = self._prisma(key_row=SimpleNamespace(key_alias="svc-key"))
instance = self._instance(prisma=prisma)
job = self._job(created_by=None)
metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
metadata = await instance.batch_metadata(job, "batch-1")
assert metadata["user_api_key"] == "hash-alice"
assert metadata["user_api_key_user_id"] is None
assert metadata["user_api_key_alias"] == "svc-key"
instance.prisma_client.db.litellm_usertable.find_unique.assert_not_called()
prisma.db.litellm_usertable.find_unique.assert_not_called()
@pytest.mark.asyncio
async def test_batch_metadata_keeps_the_creator_column_and_adds_no_budgets(self):
"""The batch cost event keeps the user id and fields the batch poll has always written,
even though background responses read budgets and drop key-scoped creators off the same rows."""
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="svc-key", spend=4.0, max_budget=5.0),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None, spend=40.0, max_budget=50.0),
)
metadata = await instance.batch_metadata(self._job(created_by="key:hash-alice"), "batch-1")
assert metadata["user_api_key_user_id"] == "key:hash-alice"
assert [field for field in metadata if field.endswith(("_spend", "_max_budget"))] == []
@pytest.mark.asyncio
async def test_metadata_drops_non_string_tags(self):
@ -2449,7 +2470,7 @@ class TestBatchCostAttribution:
instance = self._instance()
job = self._job(request_tags=["env:prod", 7, None, "team:ml"])
metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
metadata = await instance.batch_metadata(job, "batch-1")
assert metadata["tags"] == ["env:prod", "team:ml"]
@ -2457,10 +2478,11 @@ class TestBatchCostAttribution:
async def test_key_alias_lookup_failure_does_not_break_attribution(self):
"""An alias lookup failure must not lose the spend row; the key hash and team still
attribute it."""
instance = self._instance()
instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down"))
prisma = self._prisma()
prisma.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down"))
instance = self._instance(prisma=prisma)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert metadata["user_api_key"] == "hash-alice"
assert metadata.get("user_api_key_alias") is None
@ -2469,7 +2491,7 @@ class TestBatchCostAttribution:
async def test_cli_session_batch_keeps_its_alias_without_a_key_row(self):
instance = self._instance(key_row=None)
metadata = await instance._build_creator_attribution_metadata(self._job(api_key="cli-session-alice"), "batch-1")
metadata = await instance.batch_metadata(self._job(api_key="cli-session-alice"), "batch-1")
assert metadata["user_api_key"] == "cli-session-alice"
assert metadata["user_api_key_alias"] == "cli-session-alice"
@ -2478,7 +2500,7 @@ class TestBatchCostAttribution:
async def test_raw_cli_session_token_on_a_legacy_batch_row_is_not_treated_as_the_alias(self):
instance = self._instance(key_row=None)
metadata = await instance._build_creator_attribution_metadata(
metadata = await instance.batch_metadata(
self._job(api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa"), "batch-1"
)
@ -2496,7 +2518,7 @@ class TestBatchCostAttribution:
user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert metadata["user_api_key_alias"] == "Alice Chen"
assert metadata["user_api_key"] == "hash-alice"
@ -2512,7 +2534,7 @@ class TestBatchCostAttribution:
user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert metadata["user_api_key_alias"] == "Alice Chen"
@ -2527,7 +2549,7 @@ class TestBatchCostAttribution:
user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert metadata["user_api_key_alias"] == "prod-key"
@ -2542,7 +2564,7 @@ class TestBatchCostAttribution:
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
metadata = await instance._build_creator_attribution_metadata(self._job(org_id="org-at-creation"), "batch-1")
metadata = await instance.batch_metadata(self._job(org_id="org-at-creation"), "batch-1")
assert metadata["user_api_key_org_id"] == "org-at-creation"
@ -2557,7 +2579,7 @@ class TestBatchCostAttribution:
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert metadata["user_api_key_org_id"] == "org-42"
@ -2572,7 +2594,7 @@ class TestBatchCostAttribution:
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert metadata["user_api_key_org_id"] == "org-team"
@ -2582,12 +2604,11 @@ class TestBatchCostAttribution:
organization: the two lookups fail independently, so org spend still lands."""
from types import SimpleNamespace
instance = self._instance(
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down"))
prisma = self._prisma(team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"))
prisma.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down"))
instance = self._instance(prisma=prisma)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert metadata["user_api_key_org_id"] == "org-team"
@ -2602,7 +2623,7 @@ class TestBatchCostAttribution:
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
metadata = await instance.batch_metadata(self._job(), "batch-1")
assert "user_api_key_org_id" not in metadata
@ -2625,7 +2646,7 @@ class TestBatchCostAttribution:
key_row=SimpleNamespace(key_alias="prod-key"),
user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None),
)
metadata = await instance._build_creator_attribution_metadata(self._job(api_key=token_hash), "batch-1")
metadata = await instance.batch_metadata(self._job(api_key=token_hash), "batch-1")
assert metadata["user_api_key"] == token_hash
assert metadata["user_api_key_hash"] == token_hash

788
tests/unit/proxy/common_utils/test_check_responses_cost.py
View file

@ -3,15 +3,48 @@ Unit tests for CheckResponsesCost class
"""
import asyncio
import dataclasses
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
from prisma.errors import UniqueViolationError
from litellm_enterprise.proxy.common_utils.check_responses_cost import BilledBackgroundResponse, BilledCharge
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
class _RecordingBilledLogger:
"""Stands in for the spend logger so a test reads the billing events a poller wrote, and can
make the next writes raise."""
def __init__(self, failures: tuple[BaseException, ...] = (), cost: float = 0.25) -> None:
self.billed: tuple[BilledBackgroundResponse, ...] = ()
self._failures = failures
self.charge: Final = BilledCharge(cost=cost, logging_obj=MagicMock())
async def __call__(self, billed: BilledBackgroundResponse) -> BilledCharge:
await asyncio.sleep(0)
if self._failures:
failure, *remaining = self._failures
self._failures = tuple(remaining)
raise failure
self.billed = (*self.billed, billed)
return self.charge
def _deployment(model_id: str) -> Deployment:
return Deployment(
model_name="gpt-5.6",
litellm_params=LiteLLM_Params(model="openai/gpt-5.6"),
model_info=ModelInfo(id=model_id),
)
class _RecordingManagedObjectTable:
@ -27,6 +60,114 @@ class _RecordingManagedObjectTable:
return len(self.rows)
@dataclass(frozen=True)
class _StoredResponse:
id: str
unified_object_id: str
status: str
model_object_id: str = ""
file_purpose: str = "response"
file_object: Mapping[str, object] = dataclasses.field(default_factory=lambda: {"model": "gpt-5"})
created_by: str | None = None
api_key: str | None = None
team_id: str | None = None
org_id: str | None = None
request_tags: object = None
end_user_id: str | None = None
class _InMemoryManagedObjectTable:
"""Applies where filters and status writes the way Postgres does, so two pollers sharing it
race over the same rows."""
def __init__(self, rows: tuple[_StoredResponse, ...]) -> None:
self._rows: Final = rows
self.statuses: Final = {row.id: row.status for row in rows}
self.model_object_ids: Final = {row.id: row.model_object_id for row in rows}
def _matches(self, row: _StoredResponse, where: Mapping[str, object]) -> bool:
def value_of(field: str) -> object:
if field == "status":
return self.statuses[row.id]
if field == "model_object_id":
return self.model_object_ids[row.id]
return getattr(row, field)
def holds(field: str, condition: object) -> bool:
if field == "OR" and isinstance(condition, list):
return any(self._matches(row, alternative) for alternative in condition)
if isinstance(condition, dict):
return value_of(field) in condition["in"]
return value_of(field) == condition
return all(holds(field, condition) for field, condition in where.items())
def _current(self, where: Mapping[str, object]) -> list[_StoredResponse]:
return [
dataclasses.replace(row, status=self.statuses[row.id], model_object_id=self.model_object_ids[row.id])
for row in self._rows
if self._matches(row, where)
]
async def find_many(self, where: Mapping[str, object], **_: object) -> list[_StoredResponse]:
return self._current(where)
async def find_first(self, where: Mapping[str, object]) -> _StoredResponse | None:
return next(iter(self._current(where)), None)
async def update_many(self, where: Mapping[str, object], data: Mapping[str, str]) -> int:
matched: Final = self._current(where)
if "model_object_id" in data and any(
other.id not in {row.id for row in matched} and value == data["model_object_id"]
for other, value in ((other, self.model_object_ids[other.id]) for other in self._rows)
):
raise UniqueViolationError({"user_facing_error": {"error_code": "P2002"}})
self.statuses.update({row.id: data["status"] for row in matched if "status" in data})
self.model_object_ids.update({row.id: data["model_object_id"] for row in matched if "model_object_id" in data})
return len(matched)
def _poller_over(table: _InMemoryManagedObjectTable, billed_logger: _RecordingBilledLogger):
from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost
prisma_client: Final = MagicMock()
prisma_client.db.litellm_managedobjecttable = table
poller: Final = CheckResponsesCost(
proxy_logging_obj=MagicMock(),
prisma_client=prisma_client,
llm_router=MagicMock(),
log_billed_response=billed_logger,
)
poller._expire_stale_rows = AsyncMock(return_value=0)
return poller
def _terminal_response(status: str = "completed") -> ResponsesAPIResponse:
return ResponsesAPIResponse(
id="resp_upstream",
object="response",
status=status,
created_at=int(datetime.now().timestamp()),
output=[],
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
)
def _claim(job_id: str) -> tuple[Mapping[str, object], Mapping[str, object]]:
return ({"id": job_id, "status": {"in": ["queued", "in_progress"]}}, {"status": "completed"})
_READ_OF_A_BACKGROUND_RESPONSE: Final = MappingProxyType({"background": True})
def _is_billed(litellm_metadata: Mapping[str, object]) -> bool:
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call
return not is_unbilled_non_inference_call(
"aget_responses", dict(litellm_metadata), _READ_OF_A_BACKGROUND_RESPONSE
)
class TestCheckResponsesCost:
"""Test suite for CheckResponsesCost class"""
@ -50,12 +191,16 @@ class TestCheckResponsesCost:
"""Create a mock LLM Router"""
router = MagicMock()
router.aget_responses = AsyncMock()
router.get_deployment = MagicMock()
router.get_deployment = MagicMock(side_effect=lambda model_id: _deployment(model_id))
return router
@pytest.fixture
def billed_logger(self) -> _RecordingBilledLogger:
return _RecordingBilledLogger()
@pytest.fixture
def check_responses_cost_instance(
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router, billed_logger
):
"""Create a CheckResponsesCost instance with mocked dependencies"""
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
@ -66,6 +211,7 @@ class TestCheckResponsesCost:
proxy_logging_obj=mock_proxy_logging_obj,
prisma_client=mock_prisma_client,
llm_router=mock_llm_router,
log_billed_response=billed_logger,
)
# Mock _expire_stale_rows (raw SQL) so _cleanup_stale_managed_objects
# succeeds without a real DB. Individual tests can override this.
@ -121,7 +267,7 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
async def test_check_responses_cost_with_completed_response(
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, billed_logger
):
"""Test check_responses_cost with a completed response"""
# Mock job with response ID
@ -150,23 +296,22 @@ class TestCheckResponsesCost:
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
return_value=1
)
# Run the check with mocked litellm.aget_responses
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
# update_many should only contain the job completion call
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
completion_call = calls[0]
assert completion_call[1]["data"]["status"] == "completed"
assert completion_call[1]["where"]["id"]["in"] == ["job-123"]
assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-123")]
assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False]
assert [(billed.response.id, billed.response.usage) for billed in billed_logger.billed] == [
("resp_test_123", mock_response.usage)
]
@pytest.mark.asyncio
async def test_check_responses_cost_with_failed_response(
@ -399,7 +544,8 @@ class TestCheckResponsesCost:
mock_job.id = "job-404"
mock_job.file_object = {"model": "gpt-4o", "id": encoded_response_id}
mock_llm_router.get_deployment.return_value = {"model_id": "deployment-404"}
mock_llm_router.get_deployment.side_effect = None
mock_llm_router.get_deployment.return_value = _deployment("deployment-404")
mock_llm_router.aget_responses = AsyncMock(
side_effect=litellm.NotFoundError(
message="Response with id 'resp_upstream_404' not found.", model="gpt-5", llm_provider="openai"
@ -461,7 +607,8 @@ class TestCheckResponsesCost:
mock_job.id = "job-404-mapped"
mock_job.file_object = {"model": "gpt-5.5", "id": encoded_response_id}
mock_llm_router.get_deployment.return_value = {"model_id": "deployment-404-mapped"}
mock_llm_router.get_deployment.side_effect = None
mock_llm_router.get_deployment.return_value = _deployment("deployment-404-mapped")
mock_llm_router.aget_responses = AsyncMock(side_effect=mapped.value)
table = _RecordingManagedObjectTable(rows=(mock_job,))
@ -493,6 +640,7 @@ class TestCheckResponsesCost:
mock_job.id = "job-404-fallback"
mock_job.file_object = {"model": "gpt-4o", "id": encoded_response_id}
mock_llm_router.get_deployment.side_effect = None
mock_llm_router.get_deployment.return_value = None
table = _RecordingManagedObjectTable(rows=(mock_job,))
@ -531,7 +679,8 @@ class TestCheckResponsesCost:
mock_job.id = "job-404-other"
mock_job.file_object = {"model": "azure-gpt-5", "id": encoded_response_id}
mock_llm_router.get_deployment.return_value = {"model_id": "deployment-renamed"}
mock_llm_router.get_deployment.side_effect = None
mock_llm_router.get_deployment.return_value = _deployment("deployment-renamed")
mock_llm_router.aget_responses = AsyncMock(
side_effect=litellm.NotFoundError(
message="Resource not found", model="gpt-5", llm_provider="azure"
@ -550,7 +699,7 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
async def test_check_responses_cost_deployment_lookup_error_skips_only_that_job(
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, billed_logger
):
"""A deployment lookup that raises skips its own job and the cycle still records the next job."""
from litellm.responses.utils import ResponsesAPIRequestUtils
@ -576,7 +725,7 @@ class TestCheckResponsesCost:
mock_llm_router.get_deployment.side_effect = [
Exception("Model invalid format - <class 'str'>"),
{"model_id": "deployment-healthy"},
_deployment("deployment-healthy"),
]
mock_llm_router.aget_responses = AsyncMock(
return_value=ResponsesAPIResponse(
@ -596,9 +745,11 @@ class TestCheckResponsesCost:
await check_responses_cost_instance.check_responses_cost()
mock_sdk_aget.assert_not_called()
mock_llm_router.aget_responses.assert_awaited_once()
assert mock_llm_router.aget_responses.call_args.kwargs["response_id"] == healthy_response_id
assert table.updates == (({"id": {"in": ["job-healthy"]}}, {"status": "completed"}),)
assert [call.kwargs["response_id"] for call in mock_llm_router.aget_responses.call_args_list] == [
healthy_response_id,
]
assert table.updates == (_claim("job-healthy"),)
assert [billed.response.id for billed in billed_logger.billed] == [healthy_response_id]
@pytest.mark.asyncio
async def test_check_responses_cost_non_404_error_keeps_row_for_retry(
@ -620,7 +771,8 @@ class TestCheckResponsesCost:
mock_job.id = "job-500"
mock_job.file_object = {"model": "gpt-4o", "id": encoded_response_id}
mock_llm_router.get_deployment.return_value = {"model_id": "deployment-500"}
mock_llm_router.get_deployment.side_effect = None
mock_llm_router.get_deployment.return_value = _deployment("deployment-500")
mock_llm_router.aget_responses = AsyncMock(
side_effect=litellm.InternalServerError(
message="boom", model="gpt-5", llm_provider="openai"
@ -639,7 +791,7 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
async def test_check_responses_cost_multiple_jobs(
self, check_responses_cost_instance, mock_prisma_client
self, check_responses_cost_instance, mock_prisma_client, billed_logger
):
"""Test check_responses_cost with multiple jobs"""
# Mock multiple jobs
@ -702,25 +854,23 @@ class TestCheckResponsesCost:
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
return_value=1
)
# Run the check
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.side_effect = [mock_response1, mock_response2, mock_response3]
await check_responses_cost_instance.check_responses_cost()
# update_many should only contain the job completion call
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
completion_call = calls[0]
assert len(completion_call[1]["where"]["id"]["in"]) == 2
assert "job-1" in completion_call[1]["where"]["id"]["in"]
assert "job-3" in completion_call[1]["where"]["id"]["in"]
assert "job-2" not in completion_call[1]["where"]["id"]["in"]
assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-1"), _claim("job-3")]
assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False, False, False]
assert [(billed.response.id, billed.response.usage) for billed in billed_logger.billed] == [
("resp_test_1", mock_response1.usage),
("resp_test_3", mock_response3.usage),
]
@pytest.mark.asyncio
async def test_encoded_response_id_is_fetched_through_router(
@ -787,9 +937,7 @@ class TestCheckResponsesCost:
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
assert calls[0][1]["data"]["status"] == "completed"
assert calls[0][1]["where"]["id"]["in"] == ["job-router"]
assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-router")]
@pytest.mark.asyncio
async def test_encrypted_response_id_is_fetched_through_router(
@ -860,8 +1008,7 @@ class TestCheckResponsesCost:
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
assert calls[0][1]["where"]["id"]["in"] == ["job-encrypted"]
assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-encrypted")]
@pytest.mark.asyncio
async def test_response_id_without_model_id_uses_sdk(
@ -955,9 +1102,7 @@ class TestCheckResponsesCost:
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
assert calls[0][1]["data"]["status"] == "completed"
assert calls[0][1]["where"]["id"]["in"] == ["job-missing-deployment"]
assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-missing-deployment")]
@pytest.mark.asyncio
async def test_check_responses_cost_with_incomplete_response(
@ -993,9 +1138,7 @@ class TestCheckResponsesCost:
calls = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
)
assert len(calls) == 1
assert calls[0][1]["data"]["status"] == "completed"
assert calls[0][1]["where"]["id"]["in"] == ["job-incomplete"]
assert [(call[1]["where"], call[1]["data"]) for call in calls] == [_claim("job-incomplete")]
@pytest.mark.asyncio
async def test_check_responses_cost_no_model_in_file_object(
@ -1012,7 +1155,7 @@ class TestCheckResponsesCost:
return_value=[mock_job]
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
return_value=1
)
mock_response = MagicMock()
@ -1028,12 +1171,12 @@ class TestCheckResponsesCost:
assert "model_group" not in call_kwargs.get("litellm_metadata", {})
@pytest.mark.asyncio
async def test_poll_stamps_internal_call_origin_so_the_read_is_billed(
self, check_responses_cost_instance, mock_prisma_client
async def test_poll_stamps_the_billing_event_so_it_is_priced_and_its_status_read_is_free(
self, check_responses_cost_instance, mock_prisma_client, billed_logger
):
"""A background create returns queued with no usage, so this poll's retrieval is the only
place the job's spend is ever seen. Without the origin stamp it is priced at zero like a
user-facing read (LIT-5602) and the job is never billed."""
"""A background create returns queued with no usage, so the event the poll logs from the
finished response is the only place the job's spend is ever seen. Without the origin stamp
it is priced at zero like a user-facing read (LIT-5602) and the job is never billed."""
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.litellm_core_utils.internal_call_metadata import (
is_unbilled_non_inference_call,
@ -1049,7 +1192,7 @@ class TestCheckResponsesCost:
return_value=[mock_job]
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
return_value=1
)
mock_response = MagicMock()
@ -1059,8 +1202,549 @@ class TestCheckResponsesCost:
mock_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
metadata = mock_aget.call_args[1]["litellm_metadata"]
foreground_read = {"background": False}
assert metadata[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "background_response_cost_poll"
assert is_unbilled_non_inference_call("aget_responses", metadata, foreground_read) is False
assert is_unbilled_non_inference_call("aget_responses", None, foreground_read) is True
assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False]
(billed,) = billed_logger.billed
assert billed.metadata[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "background_response_cost_poll"
assert (
is_unbilled_non_inference_call("aget_responses", dict(billed.metadata), _READ_OF_A_BACKGROUND_RESPONSE)
is False
)
@pytest.mark.asyncio
async def test_poll_bills_the_response_to_the_creating_key_team_org_and_tags(
self, check_responses_cost_instance, mock_prisma_client, billed_logger
):
"""The poll's billing event is the one billed event of a background job, so its spend row must name
the key, team, org, user, and tags of the create it settles, read off the managed object row,
never the poller itself."""
from types import SimpleNamespace
mock_job = MagicMock()
mock_job.unified_object_id = "resp_test_attributed"
mock_job.id = "job-attributed"
mock_job.created_by = "alice"
mock_job.api_key = "hash-alice"
mock_job.team_id = "team-alpha"
mock_job.org_id = "org-1"
mock_job.request_tags = ["env:prod"]
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_attributed"}
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=SimpleNamespace(key_alias="prod-key", organization_id=None, spend=4.0, max_budget=5.0)
)
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=SimpleNamespace(team_alias="Team Alpha", organization_id=None, spend=40.0, max_budget=50.0)
)
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=SimpleNamespace(user_email="alice@example.com", user_alias=None, spend=1.0, max_budget=2.0)
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.return_value = MagicMock(status="completed")
await check_responses_cost_instance.check_responses_cost()
(billed,) = billed_logger.billed
metadata = billed.metadata
assert metadata["user_api_key"] == "hash-alice"
assert metadata["user_api_key_hash"] == "hash-alice"
assert metadata["user_api_key_user_id"] == "alice"
assert metadata["user_api_key_user_email"] == "alice@example.com"
assert metadata["user_api_key_team_id"] == "team-alpha"
assert metadata["user_api_key_org_id"] == "org-1"
assert metadata["user_api_key_alias"] == "prod-key"
assert metadata["user_api_key_team_alias"] == "Team Alpha"
assert (metadata["user_api_key_spend"], metadata["user_api_key_max_budget"]) == (4.0, 5.0)
assert (metadata["user_api_key_team_spend"], metadata["user_api_key_team_max_budget"]) == (40.0, 50.0)
assert (metadata["user_api_key_user_spend"], metadata["user_api_key_user_max_budget"]) == (1.0, 2.0)
assert metadata["tags"] == ["env:prod"]
assert metadata["model_group"] == "gpt-5"
@pytest.mark.asyncio
async def test_two_pollers_racing_over_one_completed_response_bill_it_once():
"""Every worker and pod runs its own poll, so two of them can pick up the same completed
response in one window; only the one that moves the row out of a pending status bills it."""
table: Final = _InMemoryManagedObjectTable(
rows=(_StoredResponse(id="job-race", unified_object_id="resp_race", status="in_progress"),)
)
reads: Final[list[Mapping[str, object]]] = []
billed_logger: Final = _RecordingBilledLogger()
async def provider_read(response_id: str, litellm_metadata: Mapping[str, object]) -> ResponsesAPIResponse:
reads.append(litellm_metadata)
await asyncio.sleep(0)
return _terminal_response()
with patch("litellm.aget_responses", side_effect=provider_read):
await asyncio.gather(
_poller_over(table, billed_logger).check_responses_cost(),
_poller_over(table, billed_logger).check_responses_cost(),
)
assert [_is_billed(metadata) for metadata in reads] == [False, False]
assert len(billed_logger.billed) == 1
assert table.statuses == {"job-race": "completed"}
@pytest.mark.asyncio
async def test_a_failed_billing_event_puts_the_response_back_for_the_next_poll():
table: Final = _InMemoryManagedObjectTable(
rows=(_StoredResponse(id="job-retry", unified_object_id="resp_retry", status="queued"),)
)
billed_logger: Final = _RecordingBilledLogger(failures=(RuntimeError("spend log unavailable"),))
poller: Final = _poller_over(table, billed_logger)
with patch("litellm.aget_responses", new_callable=AsyncMock, return_value=_terminal_response()):
await poller.check_responses_cost()
assert table.statuses == {"job-retry": "queued"}
assert billed_logger.billed == ()
await poller.check_responses_cost()
assert len(billed_logger.billed) == 1
assert table.statuses == {"job-retry": "completed"}
@pytest.mark.asyncio
async def test_only_the_first_read_that_sees_a_response_finished_bills_it():
"""A client polling its background response bills it on the first read that sees it finished,
under the creating key, so it is billed even if the provider forgets it before the poll runs;
reads before and after that one stay free."""
table: Final = _InMemoryManagedObjectTable(
rows=(
_StoredResponse(
id="job-read",
unified_object_id="resp_read",
status="queued",
created_by="alice",
api_key="hash-alice",
team_id="team-alpha",
),
)
)
billed_logger: Final = _RecordingBilledLogger()
poller: Final = _poller_over(table, billed_logger)
poller.prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
poller.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
poller.prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
finished: Final = _terminal_response("completed")
with patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")):
costs: Final = (
await poller.bill_if_finished(response_id="resp_read", response=_terminal_response("in_progress")),
await poller.bill_if_finished(response_id="resp_read", response=finished),
await poller.bill_if_finished(response_id="resp_read", response=finished),
)
await poller.check_responses_cost()
assert costs == (None, billed_logger.charge, None)
(billed,) = billed_logger.billed
assert billed.response.usage == finished.usage
assert billed.response.id == "resp_read"
assert billed.metadata["user_api_key_hash"] == "hash-alice"
assert billed.metadata["user_api_key_team_id"] == "team-alpha"
assert table.statuses == {"job-read": "completed"}
_ROTATION_SALT_KEY: Final = "check-responses-cost-unit-test-salt"
@pytest.mark.asyncio
async def test_a_read_through_a_reissued_response_id_bills_the_stored_response(monkeypatch: pytest.MonkeyPatch):
"""The proxy re-encrypts a response id on every read, so a client polling with the id its last
read returned still bills the row stored under the id the create returned."""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity
monkeypatch.setenv("LITELLM_SALT_KEY", _ROTATION_SALT_KEY)
issuer: Final = ResponsesIDSecurity(general_settings_reader=lambda: {}, signing_key_reader=lambda: _ROTATION_SALT_KEY)
def issued_id() -> str:
return issuer._encrypt_response_id(
ResponsesAPIResponse(id="resp_provider_job", created_at=1, output=[], status="queued"),
UserAPIKeyAuth(user_id="alice"),
).id
created_id: Final = issued_id()
reissued_id: Final = issued_id()
assert created_id != reissued_id
table: Final = _InMemoryManagedObjectTable(
rows=(
_StoredResponse(
id="job-rotated",
unified_object_id=created_id,
model_object_id="resp_provider_job",
status="queued",
),
)
)
billed_logger: Final = _RecordingBilledLogger()
with patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")):
await _poller_over(table, billed_logger).bill_if_finished(response_id=reissued_id, response=_terminal_response())
assert len(billed_logger.billed) == 1
assert table.statuses == {"job-rotated": "completed"}
def _issued_ids(monkeypatch: pytest.MonkeyPatch, count: int) -> tuple[str, ...]:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity
monkeypatch.setenv("LITELLM_SALT_KEY", _ROTATION_SALT_KEY)
issuer: Final = ResponsesIDSecurity(general_settings_reader=lambda: {}, signing_key_reader=lambda: _ROTATION_SALT_KEY)
return tuple(
issuer._encrypt_response_id(
ResponsesAPIResponse(id="resp_provider_job", created_at=1, output=[], status="queued"),
UserAPIKeyAuth(user_id="alice"),
).id
for _ in range(count)
)
def _stored_before_the_upgrade(created_id: str) -> _StoredResponse:
return _StoredResponse(
id="job-legacy", unified_object_id=created_id, model_object_id=created_id, status="queued"
)
@pytest.mark.asyncio
async def test_a_response_stored_before_the_upgrade_bills_once_whichever_issued_id_reads_it(
monkeypatch: pytest.MonkeyPatch,
):
"""A row the old proxy stored keys the response by the encrypted id the create returned, so the
read that sends that id rekeys it by the provider id, and a read through any later id finds it
already billed."""
created_id, reissued_id = _issued_ids(monkeypatch, 2)
table: Final = _InMemoryManagedObjectTable(rows=(_stored_before_the_upgrade(created_id),))
billed_logger: Final = _RecordingBilledLogger()
poller: Final = _poller_over(table, billed_logger)
track_untracked: Final = AsyncMock(return_value=False)
with patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")):
await poller.bill_if_finished(response_id=created_id, response=_terminal_response())
await poller.bill_if_finished(
response_id=reissued_id, response=_terminal_response(), track_untracked=track_untracked
)
assert len(billed_logger.billed) == 1
assert table.model_object_ids == {"job-legacy": "resp_provider_job"}
track_untracked.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"read_rows, billed_by_the_poll",
[((), 1), ((_StoredResponse(id="job-read", unified_object_id="read-id", model_object_id="resp_provider_job", status="completed"),), 0)],
ids=["no_read_yet", "a_read_already_billed_it"],
)
async def test_the_poll_bills_a_response_stored_before_the_upgrade_only_when_no_read_billed_it_first(
monkeypatch: pytest.MonkeyPatch, read_rows: tuple[_StoredResponse, ...], billed_by_the_poll: int
):
"""A read through a later id misses the old row and tracks the response under its provider id,
so the poll retires the old row unbilled when that row exists and bills it once otherwise."""
(created_id,) = _issued_ids(monkeypatch, 1)
table: Final = _InMemoryManagedObjectTable(rows=(_stored_before_the_upgrade(created_id), *read_rows))
billed_logger: Final = _RecordingBilledLogger()
with patch("litellm.aget_responses", new_callable=AsyncMock, return_value=_terminal_response()):
await _poller_over(table, billed_logger).check_responses_cost()
await _poller_over(table, billed_logger).check_responses_cost()
assert len(billed_logger.billed) == billed_by_the_poll
assert table.statuses["job-legacy"] == "completed"
@pytest.mark.asyncio
async def test_a_cancelled_billing_event_puts_the_response_back_for_the_next_poll():
"""A client that disconnects mid-read cancels the billing event, so the claim is released and
the next poll still bills the response."""
table: Final = _InMemoryManagedObjectTable(
rows=(_StoredResponse(id="job-cancelled", unified_object_id="resp_cancelled", status="in_progress"),)
)
billed_logger: Final = _RecordingBilledLogger(failures=(asyncio.CancelledError(),))
with pytest.raises(asyncio.CancelledError):
await _poller_over(table, billed_logger).bill_if_finished(
response_id="resp_cancelled", response=_terminal_response()
)
assert billed_logger.billed == ()
assert table.statuses == {"job-cancelled": "in_progress"}
@pytest.mark.asyncio
async def test_the_billing_event_carries_the_creators_end_user():
"""The charge belongs to the end user the create named, whoever reads the response later."""
table: Final = _InMemoryManagedObjectTable(
rows=(
_StoredResponse(
id="job-end-user", unified_object_id="resp_end_user", status="queued", end_user_id="customer-42"
),
)
)
billed_logger: Final = _RecordingBilledLogger()
await _poller_over(table, billed_logger).bill_if_finished(
response_id="resp_end_user", response=_terminal_response()
)
assert table.statuses == {"job-end-user": "completed"}
(billed,) = billed_logger.billed
assert billed.metadata["user_api_key_end_user_id"] == "customer-42"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("created_by", "billed_user_id"),
[("alice", "alice"), ("key:hashed-token", None)],
ids=["key_with_a_user", "key_with_no_user_or_team"],
)
async def test_the_billing_event_bills_the_creating_user_and_never_a_keys_owner_scope(
created_by: str, billed_user_id: str | None
):
table: Final = _InMemoryManagedObjectTable(
rows=(
_StoredResponse(
id="job-owner",
unified_object_id="resp_owner",
status="queued",
created_by=created_by,
api_key="hashed-token",
),
)
)
billed_logger: Final = _RecordingBilledLogger()
await _poller_over(table, billed_logger).bill_if_finished(response_id="resp_owner", response=_terminal_response())
(billed,) = billed_logger.billed
assert billed.metadata["user_api_key_user_id"] == billed_user_id
@pytest.mark.asyncio
async def test_settling_from_the_deployment_bills_a_finished_response_once_from_a_free_fetch():
"""A read the caller retargets, or a delete, bills from a fetch through the configured
deployment, since the client controls the body it would otherwise hand in."""
table: Final = _InMemoryManagedObjectTable(
rows=(_StoredResponse(id="job-settle", unified_object_id="resp_settle", status="in_progress"),)
)
billed_logger: Final = _RecordingBilledLogger()
poller: Final = _poller_over(table, billed_logger)
fetched: Final = _terminal_response()
with patch("litellm.aget_responses", new_callable=AsyncMock, return_value=fetched) as fetch:
charge: Final = await poller.bill_from_deployment("resp_settle")
with patch("litellm.aget_responses", side_effect=AssertionError("a billed response is never fetched again")):
assert await poller.bill_from_deployment("resp_settle") is None
assert charge == billed_logger.charge
assert [_is_billed(call.kwargs["litellm_metadata"]) for call in fetch.call_args_list] == [False]
(billed,) = billed_logger.billed
assert billed.response.usage == fetched.usage
assert table.statuses == {"job-settle": "completed"}
@pytest.mark.asyncio
async def test_settling_from_the_deployment_leaves_a_running_response_for_the_poll():
table: Final = _InMemoryManagedObjectTable(
rows=(_StoredResponse(id="job-running", unified_object_id="resp_running", status="queued"),)
)
billed_logger: Final = _RecordingBilledLogger()
with patch("litellm.aget_responses", new_callable=AsyncMock, return_value=_terminal_response("in_progress")):
assert await _poller_over(table, billed_logger).bill_from_deployment("resp_running") is None
assert billed_logger.billed == ()
assert table.statuses == {"job-running": "queued"}
class _SpendCapture:
def __init__(self) -> None:
self.payloads: tuple[Mapping[str, object], ...] = ()
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
payload = kwargs["standard_logging_object"]
assert isinstance(payload, Mapping)
self.payloads = (*self.payloads, payload)
def _spend_capture_logger():
from litellm.integrations.custom_logger import CustomLogger
class _CapturingLogger(_SpendCapture, CustomLogger):
pass
return _CapturingLogger()
_PRICED_COST: Final = 100 * 0.001 + 50 * 0.002
@dataclass(frozen=True)
class _PricedBilling:
table: _InMemoryManagedObjectTable
response_id: str
cost: float | None
_PRICED_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType(
{"model": "openai/gpt-5.6", "api_key": "sk-unused", "input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
)
async def _bill_through_a_priced_deployment(
read: ResponsesAPIResponse,
litellm_params: Mapping[str, object] = _PRICED_LITELLM_PARAMS,
model_info: Mapping[str, object] = MappingProxyType({}),
) -> _PricedBilling:
import litellm
from litellm_enterprise.proxy.common_utils.check_responses_cost import CheckResponsesCost
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.responses.utils import ResponsesAPIRequestUtils
response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="deployment-custom", response_id="resp_upstream"
)
table: Final = _InMemoryManagedObjectTable(
rows=(_StoredResponse(id="job-priced", unified_object_id=response_id, status="queued", api_key="hash-alice"),)
)
prisma_client: Final = MagicMock()
prisma_client.db.litellm_managedobjecttable = table
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
router: Final = litellm.Router(
model_list=[
{
"model_name": "priced-gpt",
"litellm_params": dict(litellm_params),
"model_info": {"id": "deployment-custom", **model_info},
}
]
)
charge: Final = await CheckResponsesCost(
proxy_logging_obj=MagicMock(), prisma_client=prisma_client, llm_router=router
).bill_if_finished(response_id=response_id, response=read)
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
return _PricedBilling(table=table, response_id=response_id, cost=charge.cost if charge is not None else None)
@pytest.mark.asyncio
async def test_the_billing_event_prices_the_fetched_usage_at_the_deployment_rates(monkeypatch: pytest.MonkeyPatch):
"""Billing logs the body the read already fetched instead of reading it again, so the event must
carry the deployment's own pricing (100 input tokens at 0.001 plus 50 output tokens at 0.002 is
0.2), and neither take the read's zero price nor change the response the reader gets back."""
import litellm
capture: Final = _spend_capture_logger()
monkeypatch.setattr(litellm, "_async_success_callback", [capture])
read: Final = _terminal_response()
read.hidden_params = {"response_cost": 0.0}
billing: Final = await _bill_through_a_priced_deployment(read)
(payload,) = capture.payloads
assert payload["response_cost"] == pytest.approx(_PRICED_COST)
assert billing.cost == pytest.approx(_PRICED_COST)
assert payload["model_id"] == "deployment-custom"
assert payload["model_group"] == "priced-gpt"
assert payload["metadata"]["user_api_key_hash"] == "hash-alice"
assert read.hidden_params == {"response_cost": 0.0}
assert billing.table.statuses == {"job-priced": "completed"}
@pytest.mark.asyncio
@pytest.mark.parametrize("base_model_in", ["litellm_params", "model_info"])
async def test_the_billing_event_prices_the_usage_at_the_deployments_base_model(
monkeypatch: pytest.MonkeyPatch, base_model_in: str
):
"""A deployment that names a base_model for cost tracking bills its background responses at that
model's price, whether the base_model sits in litellm_params or model_info."""
import litellm
capture: Final = _spend_capture_logger()
monkeypatch.setattr(litellm, "_async_success_callback", [capture])
monkeypatch.setitem(
litellm.model_cost,
"qa-priced-base-model",
{
"litellm_provider": "openai",
"mode": "responses",
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
},
)
base_model: Final = {"base_model": "qa-priced-base-model"}
billing: Final = await _bill_through_a_priced_deployment(
_terminal_response(),
litellm_params={
"model": "openai/gpt-5.6",
"api_key": "sk-unused",
**(base_model if base_model_in == "litellm_params" else {}),
},
model_info=base_model if base_model_in == "model_info" else {},
)
(payload,) = capture.payloads
assert payload["response_cost"] == pytest.approx(_PRICED_COST)
assert billing.cost == pytest.approx(_PRICED_COST)
@pytest.mark.asyncio
async def test_the_billing_event_reaches_sync_only_success_callbacks(monkeypatch: pytest.MonkeyPatch):
"""Sync-only success callbacks (legacy string integrations, plain functions) get every async
call's event from a worker thread, so they must get the billed background response too."""
import queue
import litellm
priced_costs: Final[queue.Queue[object]] = queue.Queue()
def record_cost(kwargs: Mapping[str, object], *_: object) -> None:
priced_costs.put(kwargs.get("response_cost"))
monkeypatch.setattr(litellm, "success_callback", [record_cost])
monkeypatch.setattr(litellm, "_async_success_callback", [])
await _bill_through_a_priced_deployment(_terminal_response())
assert priced_costs.get(timeout=10) == pytest.approx(_PRICED_COST)
@pytest.mark.asyncio
async def test_the_billing_event_keys_its_spend_row_apart_from_the_read_it_bills(monkeypatch: pytest.MonkeyPatch):
"""The billing event logs the body the read returned, under the same response id, and spend logs
keep one row per request id. Its row must take its own request id, or the read's free row takes
the slot and the charge reaches the key with no spend row for an admin to see."""
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
class _RequestIdCapture(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.request_ids: tuple[str, ...] = ()
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.request_ids = (
*self.request_ids,
get_logging_payload(kwargs, response_obj, start_time, end_time)["request_id"],
)
capture: Final = _RequestIdCapture()
monkeypatch.setattr(litellm, "_async_success_callback", [capture])
read: Final = _terminal_response()
billing: Final = await _bill_through_a_priced_deployment(read)
(billed_request_id,) = capture.request_ids
assert billed_request_id != billing.response_id
assert billed_request_id.startswith(billing.response_id)

24
tests/unit/proxy/db/test_model_usage_rollup.py
View file

@ -1,4 +1,5 @@
import asyncio
import json
from datetime import datetime, timezone
from typing import Any
from unittest.mock import MagicMock
@ -6,6 +7,7 @@ from unittest.mock import MagicMock
import httpx
import pytest
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.proxy.db import rollup_lock_timeout
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.db.model_usage_rollup import (
@ -16,6 +18,7 @@ from litellm.proxy.db.model_usage_rollup import (
model_usage_task_type,
)
from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
class _FakeBatcher:
@ -155,6 +158,27 @@ async def test_flush_aggregates_each_rollup_row_into_one_upsert() -> None:
assert upserts["claude"]["create"]["request_count"] == 1
@pytest.mark.asyncio
async def test_a_background_response_billing_row_adds_spend_but_no_request() -> None:
"""A background response is billed by a later event, not by the create the client sent, so
model insights counts its spend and tokens once and its request once, from the create."""
batcher = _FakeBatcher()
prisma = _prisma(MagicMock(return_value=batcher))
create = build_model_usage_transaction(_payload(spend=0.0))
billing = build_model_usage_transaction(
_payload(spend=0.5, metadata=json.dumps({INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}))
)
assert create is not None and billing is not None
await flush_model_usage_transactions(prisma_client=prisma, transactions=[create, billing])
(upsert,) = batcher.litellm_dailymodelusage.upsert.call_args_list
created = upsert.kwargs["data"]["create"]
assert created["spend"] == 0.5
assert created["prompt_tokens"] == 20
assert (created["request_count"], created["successful_requests"], created["failed_requests"]) == (1, 1, 0)
@pytest.mark.asyncio
async def test_flush_with_no_transactions_touches_nothing() -> None:
prisma = _prisma(MagicMock())

1065
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

File diff suppressed because it is too large Load diff

45
tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -15,6 +15,7 @@ import litellm
import litellm.constants as litellm_constants
import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -48,6 +49,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
from litellm.proxy.utils import hash_token
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
StandardLoggingHiddenParams,
StandardLoggingMetadata,
StandardLoggingModelInformation,
@ -4400,7 +4402,10 @@ def test_user_traffic_carries_no_internal_call_origin():
def _spend_log_for_call_type(
call_type: str, internal_call_origin: str | None = None, background: bool | None = None
call_type: str,
internal_call_origin: str | None = None,
background: bool | None = None,
managed_read: bool = False,
) -> dict:
from litellm.types.llms.openai import ResponsesAPIResponse
@ -4415,6 +4420,7 @@ def _spend_log_for_call_type(
"metadata": {
"user_api_key": "test-key",
"internal_call_origin": internal_call_origin,
**({"_managed_background_response": True} if managed_read else {}),
}
},
},
@ -4460,6 +4466,15 @@ def test_spend_log_for_background_response_retrieval_counts_tokens():
assert payload["total_tokens"] == 6000
def test_spend_log_for_managed_background_response_read_counts_nothing():
"""A proxy that tracks a background response bills it once through the cost poll, so a client
read the managed files hook marked must not count the job's tokens a second time."""
payload = _spend_log_for_call_type("aget_responses", background=True, managed_read=True)
assert payload["total_tokens"] == 0
assert payload["spend"] == 0.0
def test_spend_log_for_foreground_response_retrieval_still_counts_nothing():
"""Guards the test above against a blanket exemption: an explicit background=false read was
already billed by its create and must stay at zero."""
@ -4573,6 +4588,34 @@ def test_get_spend_logs_id_separates_distinct_batches_whose_bodies_were_both_red
assert ids == ["batch_first_batch_cost", "batch_second_batch_cost"]
def test_a_background_response_billing_event_writes_a_row_apart_from_the_read_that_fetched_it():
"""The one billing event of a background response logs the body a read just fetched, so both
carry the same response id. Spend logs keep one row per request id, so the billing event must
key its row apart from the read's, or the charge reaches the key and no row ever shows it."""
response: Final = {"id": "resp_shared", "object": "response", "status": "completed", "background": True}
read_metadata: Final = {"user_api_key": "test-key", MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True}
billing_metadata: Final = {
"user_api_key": "test-key",
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
}
request_ids: Final = tuple(
get_logging_payload(
kwargs={
"call_type": "aget_responses",
"model": "gpt-5.6",
"litellm_call_id": call_id,
"litellm_params": {"metadata": metadata},
},
response_obj=dict(response),
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)["request_id"]
for call_id, metadata in (("call-read", read_metadata), ("call-bill", billing_metadata))
)
assert request_ids == ("resp_shared", "resp_shared_background_cost")
def test_get_spend_logs_id_prefers_the_response_id_over_the_standard_logging_id():
"""An unredacted response keeps deciding its own row key, so cache-hit ids and every
other call type behave exactly as they did before."""

2
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -1189,6 +1189,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"litellm_gateway_injected_cache": "forged-deployment-id",
"_session_deployment_affinity_ttl": 999999,
"internal_call_origin": "autorouter_classifier",
"_managed_background_response": True,
"_guardrail_pipelines": [{"name": "spoofed"}],
"_pipeline_managed_guardrails": ["evaded"],
"safe_user_metadata": "kept",
@ -1246,6 +1247,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"litellm_gateway_injected_cache",
"_session_deployment_affinity_ttl",
"internal_call_origin",
"_managed_background_response",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
}

4
tests/unit/test_responses_streaming_container_ownership.py
View file

@ -7,7 +7,7 @@ which returns ``FallbackResponsesStreamWrapper``. That wrapper set
``self.completed_response = None`` in __init__ and never updated it,
so the proxy's container-ownership hook (which reads
``getattr(stream_response, "completed_response", None)`` via
``ProxyBaseLLMRequestProcessing._extract_completed_responses_response``)
``ProxyBaseLLMRequestProcessing.extract_completed_responses_response``)
saw None on every streaming call and silently recorded nothing —
follow-up ``GET /v1/containers/<id>/files`` then 403'd for the very
key that created the container.
@ -194,7 +194,7 @@ class TestProxyOwnershipHookReadsCompletedResponse:
wrapper = wrapper_cls(gen())
asyncio.run(_drain(wrapper))
extracted = ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
extracted = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(
wrapper
)
assert extracted is not None

10
tests/unit/test_router/test_router.py
View file

@ -10410,14 +10410,16 @@ class TestBackgroundResponseCostPollCooldown:
)
assert "dep-1" in self._cooled_down_ids(router)
def test_cost_poll_not_found_does_not_cool_down_deployment(self):
@pytest.mark.parametrize("poll_origin", ["background_response_cost_poll", "background_response_status_poll"])
def test_background_response_poll_not_found_does_not_cool_down_deployment(self, poll_origin: str):
"""A client can delete its background response before the poll reads it, so a 404 on either
poll read says nothing about the deployment's health."""
from datetime import datetime
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
get_deployment_failures_for_current_minute,
)
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
router = self._router()
now = datetime.now()
@ -10427,9 +10429,7 @@ class TestBackgroundResponseCostPollCooldown:
"exception": self._not_found(),
"litellm_params": {
"model_info": {"id": "dep-1"},
"litellm_metadata": {
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
},
"litellm_metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: poll_origin},
},
},
None,