mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge pull request #39980 from BerriAI/litellm_lit_7048_batch_cost_row_once
fix(batches): account a batch's cost once, from the first retrieve that sees it final
This commit is contained in:
Mateo Wang
GitHub
commit
b09b7d3eb8
9 changed files with 732 additions and 87 deletions
|
|
@ -25,6 +25,27 @@ class BatchCostUsageResult:
|
|||
failed_requests: int
|
||||
|
||||
|
||||
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
|
||||
_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"})
|
||||
|
||||
|
||||
def batch_cost_is_final(batch: Batch) -> bool:
|
||||
"""Whether this retrieve of the batch is the one to account its cost from.
|
||||
|
||||
A batch still in flight has nothing to price, and a "completed" batch can report
|
||||
no output_file_id for a moment before the output populates; pricing either records
|
||||
$0 under the batch's single spend row and pins it there. Final means a completed
|
||||
batch whose output file has arrived or whose counts prove no line succeeded, or
|
||||
any other terminal status (failed, cancelled, expired).
|
||||
"""
|
||||
if batch.status not in _TERMINAL_BATCH_STATUSES:
|
||||
return False
|
||||
if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None:
|
||||
return True
|
||||
request_counts: Final = batch.request_counts
|
||||
return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0
|
||||
|
||||
|
||||
async def calculate_batch_cost_and_usage(
|
||||
file_content_dictionary: list[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ from litellm._logging import (
|
|||
verbose_logger,
|
||||
)
|
||||
from litellm._uuid import uuid
|
||||
from litellm.batches.batch_utils import _handle_completed_batch
|
||||
from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final
|
||||
from litellm.caching.caching import DualCache, InMemoryCache
|
||||
from litellm.caching.caching_handler import LLMCachingHandler
|
||||
from litellm.constants import (
|
||||
|
|
@ -2899,13 +2899,6 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
): # polling job will query these frequently, don't spam db logs
|
||||
return
|
||||
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
)
|
||||
|
||||
# check if file id is a unified file id
|
||||
is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(result.id)
|
||||
|
||||
batch_cost: Final = kwargs.get("batch_cost", None)
|
||||
batch_usage = kwargs.get("batch_usage", None)
|
||||
batch_models = kwargs.get("batch_models", None)
|
||||
|
|
@ -2913,9 +2906,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
batch_failed_requests: Final = kwargs.get("batch_failed_requests", None)
|
||||
has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models))
|
||||
|
||||
should_compute_batch_data: Final = (
|
||||
not is_base64_unified_file_id or not has_explicit_batch_data and result.status == "completed"
|
||||
)
|
||||
should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result)
|
||||
if has_explicit_batch_data:
|
||||
result._hidden_params["response_cost"] = batch_cost
|
||||
result._hidden_params["batch_models"] = batch_models
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import time
|
|||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
|
||||
|
||||
import litellm
|
||||
|
|
@ -84,6 +85,25 @@ else:
|
|||
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
|
||||
|
||||
|
||||
def _is_batch_cost_row(payload: SpendLogsPayload) -> bool:
|
||||
return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success"
|
||||
|
||||
|
||||
_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"})
|
||||
|
||||
|
||||
def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]:
|
||||
"""Reduce a batch's cost row to what tells the retrieves apart when logging is off.
|
||||
|
||||
A proxy run with spend logs disabled still needs one row per batch to charge it once,
|
||||
so the row is written either way, but it carries no request of its own: no metadata,
|
||||
no requester IP, no key, model, or token counts (LIT-7048).
|
||||
"""
|
||||
if disable_spend_logs is False:
|
||||
return payload
|
||||
return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS})
|
||||
|
||||
|
||||
class _SpendBatch(Protocol):
|
||||
litellm_usertable: BatchTable
|
||||
litellm_verificationtoken: BatchTable
|
||||
|
|
@ -215,7 +235,12 @@ class DBSpendUpdateWriter:
|
|||
start_time: datetime | None,
|
||||
end_time: datetime | None,
|
||||
response_cost: float | None,
|
||||
) -> None:
|
||||
) -> bool:
|
||||
"""Record the request's spend, answering whether its cost still needs charging.
|
||||
|
||||
False only for a batch retrieve whose cost row another retrieve already wrote,
|
||||
so the caller leaves the key, team, and user counters alone (LIT-7048).
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
disable_spend_logs,
|
||||
litellm_proxy_budget_name,
|
||||
|
|
@ -232,7 +257,7 @@ class DBSpendUpdateWriter:
|
|||
team_id,
|
||||
)
|
||||
if ProxyUpdateSpend.disable_spend_updates() is True:
|
||||
return
|
||||
return True
|
||||
if token is not None and isinstance(token, str) and token.startswith("sk-"):
|
||||
hashed_token = hash_token(token=token)
|
||||
else:
|
||||
|
|
@ -262,11 +287,12 @@ class DBSpendUpdateWriter:
|
|||
if team_id is not None and team_id != "":
|
||||
payload["team_id"] = team_id
|
||||
|
||||
if not await self._record_spend_log(
|
||||
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
|
||||
):
|
||||
return False
|
||||
|
||||
if disable_spend_logs is False:
|
||||
await self._insert_spend_log_to_db(
|
||||
payload=payload,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
await self._enqueue_tool_usage_transaction(
|
||||
payload=payload,
|
||||
completion_response=completion_response,
|
||||
|
|
@ -306,6 +332,7 @@ class DBSpendUpdateWriter:
|
|||
)
|
||||
|
||||
verbose_proxy_logger.debug("Runs spend update on all tables")
|
||||
return True
|
||||
except Exception:
|
||||
spend_log_error(
|
||||
"Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue "
|
||||
|
|
@ -318,7 +345,102 @@ class DBSpendUpdateWriter:
|
|||
org_id,
|
||||
end_user_id,
|
||||
)
|
||||
return
|
||||
return True
|
||||
|
||||
async def _record_spend_log(
|
||||
self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool
|
||||
) -> bool:
|
||||
if prisma_client is not None and _is_batch_cost_row(payload):
|
||||
return await self._claim_batch_cost_spend_log(
|
||||
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
|
||||
)
|
||||
if disable_spend_logs is False:
|
||||
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client)
|
||||
return True
|
||||
|
||||
async def _claim_batch_cost_spend_log(
|
||||
self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool
|
||||
) -> bool:
|
||||
"""Write the batch's cost row now, or learn that another retrieve already did.
|
||||
|
||||
Every retrieve of one batch shares this row, so the insert that lands first owns
|
||||
the charge and every later one finds the row and charges nothing (LIT-7048). Only
|
||||
a row that recorded a charge counts: a failed retrieve, a request whose client
|
||||
picked the batch id as its call id, and the $0 row an older proxy left behind
|
||||
while the batch was still running all leave the charge to be made.
|
||||
"""
|
||||
from litellm.repositories.table_repositories import SpendLogsRepository
|
||||
|
||||
request_id: Final = payload["request_id"]
|
||||
row: Final = _batch_cost_row_to_write(payload, disable_spend_logs)
|
||||
spend_logs: Final = SpendLogsRepository(prisma_client).table
|
||||
try:
|
||||
claimed: Final = await spend_logs.create_many(
|
||||
data=[prisma_client.jsonify_object(row)], # mutable-ok: prisma create_many takes a list
|
||||
skip_duplicates=True,
|
||||
)
|
||||
if claimed == 1:
|
||||
return True
|
||||
existing: Final = await spend_logs.find_unique(
|
||||
where={"request_id": request_id} # mutable-ok: prisma where clause
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log
|
||||
verbose_proxy_logger.warning(
|
||||
"Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e
|
||||
)
|
||||
await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client)
|
||||
return True
|
||||
if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success":
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own",
|
||||
request_id,
|
||||
getattr(existing, "call_type", None),
|
||||
)
|
||||
return True
|
||||
if existing.spend > 0:
|
||||
verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id)
|
||||
return False
|
||||
return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row)
|
||||
|
||||
async def _take_over_uncharged_batch_cost_row(
|
||||
self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object]
|
||||
) -> bool:
|
||||
"""Take the batch's cost row over from the poll that left it charging nothing.
|
||||
|
||||
A pre-upgrade proxy wrote that row every time it polled the batch while it was still
|
||||
running, so the charge is still to be made and the row still has to end up carrying
|
||||
it. The row stops matching the moment it carries a charge, so it is one retrieve that
|
||||
takes it over and charges, and every later one reads the charge and charges nothing.
|
||||
"""
|
||||
from litellm.repositories.table_repositories import SpendLogsRepository
|
||||
|
||||
request_id: Final = payload["request_id"]
|
||||
if payload["spend"] <= 0:
|
||||
verbose_proxy_logger.debug(
|
||||
"Cost tracking skipped: this batch costs nothing and spend row %s says so", request_id
|
||||
)
|
||||
return False
|
||||
try:
|
||||
taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many(
|
||||
data=prisma_client.jsonify_object(
|
||||
MappingProxyType({field: value for field, value in row.items() if field != "request_id"})
|
||||
),
|
||||
where={ # mutable-ok: prisma where clause
|
||||
"request_id": request_id,
|
||||
"call_type": CallTypes.aretrieve_batch.value,
|
||||
"status": "success",
|
||||
"spend": 0.0,
|
||||
},
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over
|
||||
verbose_proxy_logger.warning(
|
||||
"Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e
|
||||
)
|
||||
return False
|
||||
if taken_over == 0:
|
||||
verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id)
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _enqueue_tool_usage_transaction(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.batches.batch_utils import batch_cost_is_final
|
||||
from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
|
|
@ -37,6 +38,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
|||
from litellm.proxy.utils import ProxyUpdateSpend
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
LiteLLMBatch,
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
)
|
||||
|
|
@ -248,6 +250,18 @@ class _ProxyDBLogger(CustomLogger):
|
|||
)
|
||||
_write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
|
||||
budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
|
||||
if (
|
||||
isinstance(completion_response, LiteLLMBatch)
|
||||
and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
|
||||
and not batch_cost_is_final(completion_response)
|
||||
):
|
||||
verbose_proxy_logger.debug(
|
||||
"Cost tracking deferred for batch %s still in status %s",
|
||||
completion_response.id,
|
||||
completion_response.status,
|
||||
)
|
||||
await _release_budget_reservation(budget_reservation=budget_reservation)
|
||||
return
|
||||
user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
|
||||
team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
|
||||
org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
|
||||
|
|
@ -285,7 +299,7 @@ class _ProxyDBLogger(CustomLogger):
|
|||
call_type=call_type,
|
||||
):
|
||||
## UPDATE DATABASE
|
||||
await _update_database_and_spend_counters(
|
||||
charged: Final = await _update_database_and_spend_counters(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
increment_spend_counters=increment_spend_counters,
|
||||
user_api_key=user_api_key,
|
||||
|
|
@ -302,6 +316,8 @@ class _ProxyDBLogger(CustomLogger):
|
|||
request_tags=tags,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
if not charged:
|
||||
return
|
||||
|
||||
# update cache (fire-and-forget for backward compat:
|
||||
# cached object fields, soft budget alerts, etc.)
|
||||
|
|
@ -578,9 +594,9 @@ async def _update_database_and_spend_counters(
|
|||
budget_reservation: dict | None,
|
||||
request_tags: list[str] | None = None,
|
||||
model_access_groups: Sequence[str] | None = None,
|
||||
) -> None:
|
||||
) -> bool:
|
||||
try:
|
||||
await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
token=user_api_key,
|
||||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
|
|
@ -605,6 +621,9 @@ async def _update_database_and_spend_counters(
|
|||
"Failed to invalidate budget reservation counters after release failed"
|
||||
)
|
||||
raise
|
||||
if not charged:
|
||||
await _release_budget_reservation(budget_reservation=budget_reservation)
|
||||
return False
|
||||
|
||||
try:
|
||||
await increment_spend_counters(
|
||||
|
|
@ -630,6 +649,7 @@ async def _update_database_and_spend_counters(
|
|||
finally:
|
||||
budget_reservation["finalized"] = True
|
||||
raise
|
||||
return True
|
||||
|
||||
|
||||
async def _release_budget_reservation(budget_reservation: dict | None) -> None:
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from typing import (
|
|||
runtime_checkable,
|
||||
)
|
||||
|
||||
from litellm.batches.batch_utils import batch_cost_is_final
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.repositories.table_repositories import (
|
||||
ManagedFileRepository,
|
||||
|
|
@ -1357,12 +1358,7 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:
|
|||
enumerated the batch and none succeeded. A zero or unknown total means counts
|
||||
are unreported, so stay eligible and let the next poller pass revisit it. (#37713)
|
||||
"""
|
||||
if response.output_file_id is not None:
|
||||
return True
|
||||
request_counts = response.request_counts
|
||||
if request_counts is None:
|
||||
return False
|
||||
return request_counts.total > 0 and request_counts.completed == 0
|
||||
return batch_cost_is_final(response)
|
||||
|
||||
|
||||
async def update_batch_in_database(
|
||||
|
|
|
|||
|
|
@ -21,11 +21,12 @@ from types import MappingProxyType
|
|||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from openai.types.batch import BatchRequestCounts
|
||||
|
||||
|
||||
import litellm
|
||||
import litellm.batches.batch_utils as bu
|
||||
from litellm.types.utils import Usage
|
||||
from litellm.types.utils import LiteLLMBatch, Usage
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Builders for batch OUTPUT file rows.
|
||||
|
|
@ -1718,3 +1719,57 @@ def test_unparsable_bedrock_batch_usage_warns(caplog):
|
|||
assert usage.total_tokens == 0
|
||||
assert "does not understand" in caplog.text
|
||||
assert "inputTextTokenCount" in caplog.text
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# batch_cost_is_final
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _retrieved_batch(
|
||||
status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None
|
||||
) -> LiteLLMBatch:
|
||||
return LiteLLMBatch(
|
||||
id="batch_abc",
|
||||
completion_window="24h",
|
||||
created_at=1,
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="file-in",
|
||||
object="batch",
|
||||
status="validating",
|
||||
output_file_id=output_file_id,
|
||||
request_counts=counts,
|
||||
).model_copy(update={"status": status})
|
||||
|
||||
|
||||
class TestBatchCostIsFinal:
|
||||
"""Every retrieve of one batch writes the same spend row, so the first retrieve
|
||||
that prices it decides the row for good. A poll before the output exists must
|
||||
therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048)."""
|
||||
|
||||
@pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"])
|
||||
def test_in_flight_batch_is_not_final(self, status):
|
||||
assert bu.batch_cost_is_final(_retrieved_batch(status)) is False
|
||||
|
||||
@pytest.mark.parametrize("status", ["completed", "complete"])
|
||||
def test_completed_with_output_is_final(self, status):
|
||||
assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True
|
||||
|
||||
def test_completed_without_output_and_unknown_counts_is_not_final(self):
|
||||
assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False
|
||||
|
||||
def test_completed_without_output_and_zero_counts_is_not_final(self):
|
||||
counts = BatchRequestCounts(total=0, completed=0, failed=0)
|
||||
assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False
|
||||
|
||||
def test_completed_without_output_but_successful_lines_is_not_final(self):
|
||||
counts = BatchRequestCounts(total=2, completed=2, failed=0)
|
||||
assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False
|
||||
|
||||
@pytest.mark.parametrize("status", ["completed", "complete"])
|
||||
def test_completed_without_output_and_every_line_failed_is_final(self, status):
|
||||
counts = BatchRequestCounts(total=2, completed=0, failed=2)
|
||||
assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True
|
||||
|
||||
@pytest.mark.parametrize("status", ["failed", "expired", "cancelled"])
|
||||
def test_other_terminal_statuses_are_final(self, status):
|
||||
assert bu.batch_cost_is_final(_retrieved_batch(status)) is True
|
||||
|
|
|
|||
|
|
@ -632,6 +632,86 @@ class TestRetrieveBatchCostPassesModelIdentity:
|
|||
assert captured["model_info"]["input_cost_per_token"] == 0.0
|
||||
|
||||
|
||||
class TestRetrieveBatchPricesOnlyFinalBatches:
|
||||
"""Regression (LIT-7048): retrieving a provider-id batch priced it on every poll.
|
||||
|
||||
Every retrieve of one batch logs under the same spend row, so pricing a poll
|
||||
that landed before the output existed wrote that row at $0 and pinned it there.
|
||||
Only a final batch gets priced; an in-flight poll carries no cost at all.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _logging_obj() -> LitellmLogging:
|
||||
obj = LitellmLogging(
|
||||
model="gpt-5.6-luna",
|
||||
messages=[{"role": "user", "content": "Hey"}],
|
||||
stream=False,
|
||||
call_type="aretrieve_batch",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="batch-call-2",
|
||||
function_id="f",
|
||||
)
|
||||
obj.custom_llm_provider = "openai"
|
||||
return obj
|
||||
|
||||
@staticmethod
|
||||
def _batch(status: str, output_file_id: str | None):
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
return LiteLLMBatch(
|
||||
id="batch_6a9c99e185588190877d391f8b9d7f8a",
|
||||
completion_window="24h",
|
||||
created_at=1,
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="file-in",
|
||||
object="batch",
|
||||
status="validating",
|
||||
output_file_id=output_file_id,
|
||||
).model_copy(update={"status": status})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("status", "output_file_id"),
|
||||
[("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None), ("complete", None)],
|
||||
)
|
||||
async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None:
|
||||
from litellm.litellm_core_utils import litellm_logging as logging_module
|
||||
|
||||
handle_completed_batch = AsyncMock()
|
||||
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
|
||||
batch = self._batch(status, output_file_id)
|
||||
|
||||
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
|
||||
|
||||
handle_completed_batch.assert_not_awaited()
|
||||
assert "response_cost" not in batch._hidden_params
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None:
|
||||
from litellm.batches.batch_utils import BatchCostUsageResult
|
||||
from litellm.litellm_core_utils import litellm_logging as logging_module
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
handle_completed_batch = AsyncMock(
|
||||
return_value=BatchCostUsageResult(
|
||||
cost=8e-06,
|
||||
usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35),
|
||||
models=["gpt-5.6-luna"],
|
||||
successful_requests=2,
|
||||
failed_requests=0,
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
|
||||
batch = self._batch("completed", "file-out")
|
||||
|
||||
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
|
||||
|
||||
handle_completed_batch.assert_awaited_once()
|
||||
assert batch._hidden_params["response_cost"] == 8e-06
|
||||
assert batch.usage is not None
|
||||
assert batch.usage.total_tokens == 35
|
||||
|
||||
|
||||
class TestAnthropicPassthroughCustomPricing:
|
||||
"""Verify the Anthropic pass-through handler forwards custom pricing."""
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import re
|
|||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -2936,7 +2937,9 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
|
|||
"call_type, expects_flush",
|
||||
[("aresponses", True), ("responses", True), ("acompletion", False)],
|
||||
)
|
||||
async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool):
|
||||
async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back(
|
||||
call_type: str, expects_flush: bool
|
||||
):
|
||||
"""
|
||||
A `previous_response_id` chained straight off the previous turn reads the DB, so a
|
||||
Responses row cannot sit in this worker's queue until the monitor's next poll.
|
||||
|
|
@ -2957,6 +2960,303 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(c
|
|||
PrismaClient.spend_log_flush_requested.clear()
|
||||
|
||||
|
||||
def _batch_cost_payload() -> dict:
|
||||
return {
|
||||
**_minimal_spend_payload(),
|
||||
"request_id": "batch_abc_batch_cost",
|
||||
"call_type": "aretrieve_batch",
|
||||
"status": "success",
|
||||
}
|
||||
|
||||
|
||||
def _spend_logs_prisma(inserted: int, existing: object, taken_over: int = 1) -> MagicMock:
|
||||
prisma = _tool_usage_prisma()
|
||||
prisma.jsonify_object = lambda data: dict(data)
|
||||
prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted)
|
||||
prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing)
|
||||
prisma.db.litellm_spendlogs.update_many = AsyncMock(return_value=taken_over)
|
||||
return prisma
|
||||
|
||||
|
||||
async def _update_database_with(
|
||||
db_writer: DBSpendUpdateWriter,
|
||||
prisma: MagicMock,
|
||||
payload: dict,
|
||||
disable_spend_logs: bool = False,
|
||||
response_cost: float = 0.25,
|
||||
) -> bool:
|
||||
with (
|
||||
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
|
||||
"litellm.proxy.proxy_server.disable_spend_logs", disable_spend_logs
|
||||
),
|
||||
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
|
||||
"litellm.proxy.proxy_server.prisma_client", prisma
|
||||
),
|
||||
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
|
||||
"litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"
|
||||
),
|
||||
patch( # test-quality-ok: update_database imports the payload builder inside its body, no seam
|
||||
"litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload",
|
||||
return_value=payload,
|
||||
),
|
||||
):
|
||||
charged = await db_writer.update_database(
|
||||
token="test-token",
|
||||
user_id="test-user",
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
org_id=None,
|
||||
kwargs={"model": "gpt-5.6-luna", "call_type": "aretrieve_batch"},
|
||||
completion_response=None,
|
||||
start_time=datetime.now(timezone.utc),
|
||||
end_time=datetime.now(timezone.utc),
|
||||
response_cost=response_cost,
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
return charged
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("inserted", "existing", "charged"),
|
||||
[
|
||||
(1, None, True),
|
||||
(0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25), False),
|
||||
(0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0), True),
|
||||
(0, SimpleNamespace(call_type="aretrieve_batch", status="failure", spend=0.0), True),
|
||||
(0, SimpleNamespace(call_type="aembedding", status="success", spend=0.25), True),
|
||||
(0, None, True),
|
||||
],
|
||||
ids=[
|
||||
"first_retrieve_owns_the_row",
|
||||
"another_retrieve_already_charged",
|
||||
"an_older_proxy_left_a_zero_row_while_the_batch_ran",
|
||||
"failed_retrieve_holds_the_row",
|
||||
"client_chosen_call_id_holds_the_row",
|
||||
"row_gone_between_insert_and_lookup",
|
||||
],
|
||||
)
|
||||
async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote_its_row(
|
||||
inserted: int, existing: object, charged: bool
|
||||
):
|
||||
"""
|
||||
Every retrieve of one batch shares one spend row, so the insert that lands first is
|
||||
the charge and every later retrieve must leave the counters alone (LIT-7048). A row
|
||||
that recorded no charge must not be able to take the charge away: neither one a
|
||||
client planted under the batch id, nor the $0 row a pre-upgrade proxy wrote every
|
||||
time it polled the batch while it was still running.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
prisma = _spend_logs_prisma(inserted, existing)
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged
|
||||
|
||||
claimed_rows = prisma.db.litellm_spendlogs.create_many.await_args.kwargs
|
||||
assert claimed_rows["skip_duplicates"] is True
|
||||
assert [(row["request_id"], row["spend"]) for row in claimed_rows["data"]] == [("batch_abc_batch_cost", 0.25)]
|
||||
assert prisma.spend_log_transactions == []
|
||||
assert db_writer._batch_database_updates.await_count == (1 if charged else 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("taken_over", "charged"),
|
||||
[(1, True), (0, False)],
|
||||
ids=["this_retrieve_takes_it_over", "another_one_got_there_first"],
|
||||
)
|
||||
async def test_update_database_charges_a_batch_whose_row_a_pre_upgrade_poll_left_at_zero(
|
||||
taken_over: int, charged: bool
|
||||
):
|
||||
"""
|
||||
A proxy without this fix wrote the batch's row at $0 on every poll of a running batch,
|
||||
and the row outlives the upgrade, so the charge has to land on the row itself. Charging
|
||||
without writing it there would charge again on every later retrieve (LIT-7048).
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0)
|
||||
prisma = _spend_logs_prisma(0, existing, taken_over)
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged
|
||||
|
||||
taken = prisma.db.litellm_spendlogs.update_many.await_args.kwargs
|
||||
assert taken["where"] == {
|
||||
"request_id": "batch_abc_batch_cost",
|
||||
"call_type": "aretrieve_batch",
|
||||
"status": "success",
|
||||
"spend": 0.0,
|
||||
}
|
||||
assert taken["data"]["spend"] == 0.25
|
||||
assert "request_id" not in taken["data"]
|
||||
assert db_writer._batch_database_updates.await_count == (1 if charged else 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_leaves_a_batch_whose_zero_row_it_could_not_take_over_to_the_next_retrieve():
|
||||
"""
|
||||
A DB that refuses the takeover leaves the row reading $0, so charging here would charge
|
||||
the batch again on every later retrieve. The retrieve that does take the row over is the
|
||||
one that charges.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0)
|
||||
prisma = _spend_logs_prisma(0, existing)
|
||||
prisma.db.litellm_spendlogs.update_many = AsyncMock(side_effect=RuntimeError("db unreachable"))
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False
|
||||
|
||||
assert db_writer._batch_database_updates.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_leaves_a_batch_that_cost_nothing_to_the_retrieve_that_wrote_its_row():
|
||||
"""
|
||||
A batch every line of which failed costs $0, so its row reads $0 for the honest reason
|
||||
and the retrieve that wrote it is still the one that accounted it. Taking that row over
|
||||
on every later retrieve would count one batch as many requests.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0)
|
||||
prisma = _spend_logs_prisma(0, existing)
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, _batch_cost_payload(), response_cost=0.0) is False
|
||||
|
||||
prisma.db.litellm_spendlogs.update_many.assert_not_called()
|
||||
assert db_writer._batch_database_updates.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("inserted", "existing", "charged"),
|
||||
[
|
||||
(1, None, True),
|
||||
(0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25), False),
|
||||
],
|
||||
ids=["first_retrieve_owns_the_row", "another_retrieve_already_charged"],
|
||||
)
|
||||
async def test_update_database_charges_a_batch_once_even_with_spend_logs_disabled(
|
||||
inserted: int, existing: object, charged: bool
|
||||
):
|
||||
"""
|
||||
disable_spend_logs drops the per-request logs, not the batch's charge, so the one row
|
||||
that makes a batch chargeable exactly once is still written and still read back.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
prisma = _spend_logs_prisma(inserted, existing)
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, _batch_cost_payload(), True) is charged
|
||||
|
||||
assert prisma.db.litellm_spendlogs.create_many.await_count == 1
|
||||
assert db_writer._batch_database_updates.await_count == (1 if charged else 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_writes_no_ordinary_spend_row_with_spend_logs_disabled():
|
||||
"""The batch carve-out above stays a carve-out: every other row still goes unwritten."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
prisma = _spend_logs_prisma(1, None)
|
||||
payload = {**_batch_cost_payload(), "call_type": "acompletion"}
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, payload, True) is True
|
||||
|
||||
prisma.db.litellm_spendlogs.create_many.assert_not_called()
|
||||
assert prisma.spend_log_transactions == []
|
||||
assert db_writer._batch_database_updates.await_count == 1
|
||||
|
||||
|
||||
_BATCH_CLAIM_FIELDS = {"request_id", "call_type", "status", "spend", "startTime", "endTime"}
|
||||
|
||||
|
||||
def _logged_batch_cost_payload() -> dict:
|
||||
return {
|
||||
**_batch_cost_payload(),
|
||||
"api_key": "0e5b0e9e5f",
|
||||
"model": "gpt-5.6-luna",
|
||||
"user": "test-user",
|
||||
"metadata": '{"batch_models": ["gpt-5.6-luna"]}',
|
||||
"requester_ip_address": "127.0.0.1",
|
||||
"proxy_server_request": '{"headers": {"user-agent": "litellm-batch-cost-check"}}',
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("disable_spend_logs", "logs_the_request"),
|
||||
[(False, True), (True, False)],
|
||||
ids=["spend_logs_on", "spend_logs_off"],
|
||||
)
|
||||
async def test_update_database_claims_a_batch_without_logging_the_request_that_polled_it(
|
||||
disable_spend_logs: bool, logs_the_request: bool
|
||||
):
|
||||
"""
|
||||
disable_spend_logs has to keep meaning that no request gets logged, and the batch's cost
|
||||
row is the one row it cannot drop, so with logging off that row carries only what tells
|
||||
the retrieves apart: no metadata, no requester IP, no key, model, or token counts.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
prisma = _spend_logs_prisma(1, None)
|
||||
payload = _logged_batch_cost_payload()
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, payload, disable_spend_logs) is True
|
||||
|
||||
claimed = prisma.db.litellm_spendlogs.create_many.await_args.kwargs["data"][0]
|
||||
assert set(claimed) == (set(payload) if logs_the_request else _BATCH_CLAIM_FIELDS)
|
||||
assert claimed["spend"] == 0.25
|
||||
assert db_writer._batch_database_updates.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_queues_only_the_claim_for_a_batch_it_could_not_write_with_logs_disabled():
|
||||
"""A refused claim is retried through the queue, so what it queues has to stay unlogged too."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
prisma = _spend_logs_prisma(0, None)
|
||||
prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable"))
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, _logged_batch_cost_payload(), True) is True
|
||||
|
||||
assert [set(row) for row in prisma.spend_log_transactions] == [_BATCH_CLAIM_FIELDS]
|
||||
assert db_writer._batch_database_updates.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_queues_a_batch_cost_row_it_could_not_claim():
|
||||
"""An unreachable DB must not drop the batch's only spend row, nor its charge."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
prisma = _spend_logs_prisma(0, None)
|
||||
prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable"))
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True
|
||||
|
||||
assert [row["request_id"] for row in prisma.spend_log_transactions] == ["batch_abc_batch_cost"]
|
||||
assert db_writer._batch_database_updates.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"payload",
|
||||
[{**_batch_cost_payload(), "call_type": "acompletion"}, {**_batch_cost_payload(), "status": "failure"}],
|
||||
ids=["not_a_batch_retrieve", "failed_batch_retrieve"],
|
||||
)
|
||||
async def test_update_database_queues_every_other_spend_row_for_the_next_flush(payload: dict):
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
prisma = _spend_logs_prisma(1, None)
|
||||
|
||||
assert await _update_database_with(db_writer, prisma, payload) is True
|
||||
|
||||
prisma.db.litellm_spendlogs.create_many.assert_not_called()
|
||||
assert prisma.spend_log_transactions == [payload]
|
||||
assert db_writer._batch_database_updates.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"injected_deployment, attributed",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -70,9 +70,7 @@ async def test_async_post_call_failure_hook():
|
|||
|
||||
# Check that metadata was properly updated
|
||||
assert "litellm_params" in call_args["kwargs"]
|
||||
assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {
|
||||
"request_id": "test_request_id"
|
||||
}
|
||||
assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"}
|
||||
metadata = call_args["kwargs"]["litellm_params"]["metadata"]
|
||||
assert metadata["user_api_key"] == "test_api_key"
|
||||
assert metadata["status"] == "failure"
|
||||
|
|
@ -336,9 +334,7 @@ async def test_should_continue_failure_tracking_when_budget_release_fails():
|
|||
)
|
||||
assert mock_invalidate_budget_reservation_counters.await_count == 1
|
||||
assert (
|
||||
mock_invalidate_budget_reservation_counters.await_args.kwargs[
|
||||
"budget_reservation"
|
||||
]
|
||||
mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"]
|
||||
is user_api_key_dict.budget_reservation
|
||||
)
|
||||
assert user_api_key_dict.budget_reservation["finalized"] is True
|
||||
|
|
@ -433,36 +429,21 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object():
|
|||
"entries": [{"counter_key": "spend:key:test_api_key"}],
|
||||
}
|
||||
|
||||
assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None
|
||||
assert (
|
||||
_get_budget_reservation_from_metadata(
|
||||
metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
_get_budget_reservation_from_metadata(
|
||||
metadata={
|
||||
"user_api_key_auth": UserAPIKeyAuth(
|
||||
budget_reservation=budget_reservation
|
||||
)
|
||||
}
|
||||
metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)}
|
||||
)
|
||||
== budget_reservation
|
||||
)
|
||||
assert (
|
||||
_get_budget_reservation_from_metadata(
|
||||
metadata={
|
||||
"user_api_key_auth": dict(
|
||||
UserAPIKeyAuth(budget_reservation=budget_reservation)
|
||||
)
|
||||
}
|
||||
metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))}
|
||||
)
|
||||
== budget_reservation
|
||||
)
|
||||
assert (
|
||||
_get_budget_reservation_from_metadata(
|
||||
metadata={"user_api_key_budget_reservation": budget_reservation}
|
||||
)
|
||||
_get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation})
|
||||
is budget_reservation
|
||||
)
|
||||
|
||||
|
|
@ -470,9 +451,7 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object():
|
|||
@pytest.mark.asyncio
|
||||
async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails():
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
|
||||
side_effect=Exception("db unavailable")
|
||||
)
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable"))
|
||||
increment_spend_counters = AsyncMock()
|
||||
budget_reservation = {"reserved_cost": 0.5, "entries": []}
|
||||
|
||||
|
|
@ -508,9 +487,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u
|
|||
async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails():
|
||||
proxy_logging_obj = MagicMock()
|
||||
db_exception = RuntimeError("db unavailable")
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
|
||||
side_effect=db_exception
|
||||
)
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception)
|
||||
increment_spend_counters = AsyncMock()
|
||||
budget_reservation = {"reserved_cost": 0.5, "entries": []}
|
||||
|
||||
|
|
@ -554,12 +531,8 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re
|
|||
budget_reservation=budget_reservation,
|
||||
)
|
||||
assert mock_log_exception.call_count == 2
|
||||
mock_log_exception.assert_any_call(
|
||||
"Failed to release budget reservation after database update failed"
|
||||
)
|
||||
mock_log_exception.assert_any_call(
|
||||
"Failed to invalidate budget reservation counters after release failed"
|
||||
)
|
||||
mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed")
|
||||
mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed")
|
||||
|
||||
increment_spend_counters.assert_not_awaited()
|
||||
|
||||
|
|
@ -778,6 +751,107 @@ async def test_track_cost_callback_defers_in_progress_background_interaction():
|
|||
mock_proxy_logging.failed_tracking_alert.assert_not_called()
|
||||
|
||||
|
||||
def _batch_retrieve_kwargs(call_type: str, reservation: dict | None = None) -> dict:
|
||||
metadata = {
|
||||
"user_api_key": "hashed_key",
|
||||
"user_api_key_user_id": "user-1",
|
||||
"user_api_key_team_id": "team-1",
|
||||
**({"user_api_key_budget_reservation": reservation} if reservation is not None else {}),
|
||||
}
|
||||
return {
|
||||
"call_type": call_type,
|
||||
"model": "gpt-5.6-luna",
|
||||
"litellm_call_id": "test-call-id",
|
||||
"litellm_params": {"metadata": metadata},
|
||||
"standard_logging_object": {"response_cost": 0.0, "request_tags": None},
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
|
||||
def _retrieved_batch(status: str, output_file_id: str | None):
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
return LiteLLMBatch(
|
||||
id="batch_abc",
|
||||
completion_window="24h",
|
||||
created_at=1,
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="file-in",
|
||||
object="batch",
|
||||
status=status,
|
||||
output_file_id=output_file_id,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"),
|
||||
[
|
||||
("aretrieve_batch", "in_progress", None, True, False, False),
|
||||
("aretrieve_batch", "completed", None, True, False, False),
|
||||
("aretrieve_batch", "completed", "file-out", False, True, False),
|
||||
("aretrieve_batch", "completed", "file-out", True, True, True),
|
||||
("aretrieve_batch", "failed", None, True, True, True),
|
||||
("acreate_batch", "validating", None, True, True, True),
|
||||
],
|
||||
ids=[
|
||||
"retrieve_before_final",
|
||||
"retrieve_completed_without_output_yet",
|
||||
"retrieve_after_another_retrieve_charged",
|
||||
"retrieve_first_final",
|
||||
"retrieve_failed_batch",
|
||||
"create_before_final",
|
||||
],
|
||||
)
|
||||
async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs, whether the counters move, and whether the poll's reservation is handed back is the whole observable contract of the gate
|
||||
call_type, status, output_file_id, row_claimed, spend_written, charged
|
||||
):
|
||||
"""
|
||||
A poll before the batch is final used to pin its shared spend row at $0, and every
|
||||
completed retrieve after the first charged the key again (LIT-7048). Only retrieves
|
||||
are gated, since creating a batch is its own billable request, and a retrieve that
|
||||
charges nothing hands its budget reservation back instead.
|
||||
"""
|
||||
logger = _ProxyDBLogger()
|
||||
budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []}
|
||||
kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam
|
||||
"litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock
|
||||
) as mock_increment_spend_counters,
|
||||
patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam
|
||||
"litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock
|
||||
) as mock_update_cache,
|
||||
patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj"
|
||||
) as mock_proxy_logging,
|
||||
patch( # test-quality-ok: the release is imported inside the callback's helper, no seam
|
||||
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", new_callable=AsyncMock
|
||||
) as mock_release_budget_reservation,
|
||||
):
|
||||
mock_proxy_logging.failed_tracking_alert = AsyncMock()
|
||||
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=row_claimed)
|
||||
mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
|
||||
|
||||
await logger._PROXY_track_cost_callback(
|
||||
kwargs=kwargs,
|
||||
completion_response=_retrieved_batch(status, output_file_id),
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
mock_proxy_logging.failed_tracking_alert.assert_not_called()
|
||||
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if spend_written else 0)
|
||||
assert mock_increment_spend_counters.await_count == (1 if charged else 0)
|
||||
assert mock_update_cache.await_count == (1 if charged else 0)
|
||||
if charged:
|
||||
mock_release_budget_reservation.assert_not_awaited()
|
||||
else:
|
||||
mock_release_budget_reservation.assert_awaited_once_with(budget_reservation=budget_reservation)
|
||||
|
||||
|
||||
def _in_progress_interaction_kwargs(reservation: dict) -> dict:
|
||||
return {
|
||||
"call_type": "acreate_interaction",
|
||||
|
|
@ -1101,10 +1175,7 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj
|
|||
|
||||
# standard_logging_object should have been propagated from logging obj
|
||||
assert call_kwargs.get("standard_logging_object") is not None
|
||||
assert (
|
||||
call_kwargs["standard_logging_object"]["trace_id"]
|
||||
== "trace-id-from-logging-obj"
|
||||
)
|
||||
assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj"
|
||||
# litellm_trace_id should also be propagated as a fallback
|
||||
assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj"
|
||||
|
||||
|
|
@ -1691,9 +1762,7 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend():
|
|||
"metadata": {},
|
||||
"proxy_server_request": {"request_id": "rid"},
|
||||
"response_cost": 3.5e-05,
|
||||
"combined_usage_object": Usage(
|
||||
prompt_tokens=30, completion_tokens=1, total_tokens=31
|
||||
),
|
||||
"combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31),
|
||||
}
|
||||
|
||||
with patch(
|
||||
|
|
@ -1772,15 +1841,10 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata():
|
|||
assert mock_increment.call_args.kwargs["team_id"] == "team-123"
|
||||
assert mock_increment.call_args.kwargs["org_id"] == "org-456"
|
||||
|
||||
update_kwargs = (
|
||||
mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs
|
||||
)
|
||||
update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs
|
||||
assert update_kwargs["user_id"] == "mcp-user@example.com"
|
||||
assert update_kwargs["team_id"] == "team-123"
|
||||
assert (
|
||||
kwargs["litellm_params"]["metadata"]["user_api_key_user_id"]
|
||||
== "mcp-user@example.com"
|
||||
)
|
||||
assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1875,9 +1939,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect
|
|||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
|
||||
call_type, expect_spend_log
|
||||
):
|
||||
async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log):
|
||||
"""Regression for LIT-3782: a pass-through request with auth=false reaches the
|
||||
cost callback with no key/user/team/end-user. Before the fix the spend-log
|
||||
write was skipped and the request never appeared in request/usage logs. It
|
||||
|
|
@ -1923,9 +1985,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
|
|||
end_time=datetime.now(),
|
||||
)
|
||||
|
||||
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (
|
||||
1 if expect_spend_log else 0
|
||||
)
|
||||
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0)
|
||||
|
||||
|
||||
class _FakeDeploymentLookup:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue