diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 522e05ebbd5..d272a235758 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -10,6 +10,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils import batch_utils as core_batch_utils from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo @@ -39,6 +40,7 @@ class BatchCostUsageResult: _DECODED_BATCH_LINE: Final = TypeAdapter(object) +batch_cost_is_final: Final = core_batch_utils.batch_cost_is_final def _uses_native_vertex_output( diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 83508a08de6..35e7ae1c579 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -160,6 +160,22 @@ def _org_member_transaction_key(org_id: str, user_id: str) -> str: return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}" +async def _track_guardrail_usage_of_cost_row( + payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool +) -> None: + """A job's cost row skips the spend log queue, so its guardrail runs are counted here instead of at the flush.""" + if disable_spend_logs: + return + from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage + + try: + await process_spend_logs_guardrail_usage(prisma_client=prisma_client, logs_to_process=(payload,)) + except Exception as e: # noqa: BLE001 # guardrail usage is best effort, as at the spend log flush + verbose_proxy_logger.warning( + "Guardrail usage tracking failed for spend row %s (non-fatal): %s", payload["request_id"], e + ) + + def _is_settle_once_cost_row(payload: SpendLogsPayload) -> bool: """A poll of a job whose cost is charged once, by the first poll that sees it, however often it is polled. @@ -807,15 +823,16 @@ class DBSpendUpdateWriter: data=[prisma_client.jsonify_object(row)], skip_duplicates=True, ) - if claimed == 1: - return True - existing: Final = await spend_logs.find_unique(where={"request_id": request_id}) + existing: Final = None if claimed == 1 else await spend_logs.find_unique(where={"request_id": request_id}) 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 job'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 claimed == 1: + await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs) + return True if existing is None or existing.call_type != payload["call_type"] or existing.status != "success": verbose_proxy_logger.warning( "Spend row %s belongs to a %s request, so this job's cost is charged without a row of its own", @@ -826,7 +843,12 @@ class DBSpendUpdateWriter: if existing.spend > 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this job", request_id) return False - return await self._take_over_uncharged_cost_row(payload=payload, prisma_client=prisma_client, row=row) + taken_over: Final = await self._take_over_uncharged_cost_row( + payload=payload, prisma_client=prisma_client, row=row + ) + if taken_over: + await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs) + return taken_over async def _take_over_uncharged_cost_row( self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 00939cbffe7..686527842dd 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -28,6 +28,7 @@ from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import ( ChargeOnceJob, + charged_poll_cost, defers_batch_cost, job_charged_once_marker_key, job_to_charge_once, @@ -177,7 +178,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return callback_fields: Final = _CALLBACK_FIELDS.validate_python(kwargs) - response_cost: Final = finite_cost(callback_fields.get("response_cost")) or 0.0 + response_cost: Final = finite_cost(charged_poll_cost(callback_fields)) or 0.0 if response_cost <= 0: return @@ -196,7 +197,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): verbose_proxy_logger.debug( "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f", session_id, - response_cost, + charge, ) except Exception as e: verbose_proxy_logger.warning( diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 999f538cf85..3a79df0370b 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1,6 +1,7 @@ import asyncio import copy import functools +import itertools import json import logging import re @@ -4787,6 +4788,52 @@ async def test_update_database_charges_a_video_once_from_the_poll_that_sees_it_c } +def _indexed_guardrail_request_ids(prisma: MagicMock, index_table: str) -> tuple[str, ...]: + calls: Final = getattr(prisma.db, index_table).create_many.await_args_list + batches: Final = tuple(call.kwargs["data"] for call in calls) + return tuple(row["request_id"] for row in itertools.chain.from_iterable(batches)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("disable_spend_logs", [False, True], ids=["spend_logs_on", "spend_logs_off"]) +@pytest.mark.parametrize("left_uncharged", [False, True], ids=["claimed_row", "taken_over_row"]) +async def test_a_job_cost_row_counts_its_polls_guardrail_run_in_guardrail_usage_once( + disable_spend_logs: bool, left_uncharged: bool +): + """The job's cost row is written by its claim, not the spend log queue whose flush counts guardrail runs, so a + finished video's guardrail run must still reach guardrail usage, once, and not when spend logs are off.""" + from litellm.repositories.table_repositories import ( + DailyGuardrailMetricsRepository, + SpendLogGuardrailIndexRepository, + ) + + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + getattr(prisma.db, SpendLogGuardrailIndexRepository.table_name).create_many = AsyncMock(return_value=1) + getattr(prisma.db, DailyGuardrailMetricsRepository.table_name).upsert = AsyncMock() + if left_uncharged: + prisma.db.litellm_spendlogs.rows["video_abc_video_cost"] = {**_video_status_payload(), "spend": 0.0} + guarded_poll: Final = { + **_video_status_payload(), + "metadata": json.dumps( + {"guardrail_information": [{"guardrail_name": "priced", "guardrail_status": "success"}]} + ), + } + + charged = [ + await _update_database_with( + db_writer, prisma, guarded_poll, disable_spend_logs=disable_spend_logs, response_cost=0.43 + ) + for _ in range(2) + ] + + assert charged == [True, False] + assert _indexed_guardrail_request_ids(prisma, SpendLogGuardrailIndexRepository.table_name) == ( + () if disable_spend_logs else ("video_abc_video_cost",) + ) + + def _job_poll( litellm_call_id: str, response_cost: float, diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index 1f7af3122eb..3d1ce77aae5 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -222,14 +222,14 @@ async def test_a_batch_retrieve_before_its_cost_is_final_charges_the_session_not def _video_poll( - cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs", status: str = "completed" + video_cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs", status: str = "completed" ) -> dict[str, object]: return { "call_type": "avideo_status", - "response_cost": cost, + "response_cost": video_cost, POLLED_VIDEO_STATUS_KWARG: status, "standard_logging_object": { - "response_cost": cost, + "response_cost": video_cost + guardrail_cost, "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}], }, "litellm_params": {"metadata": {"session_id": session_id, "agent_id": "agent-budget-123"}}, @@ -237,28 +237,39 @@ def _video_poll( @pytest.mark.asyncio -async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_session_for_the_video(): +@pytest.mark.parametrize( + ("video_cost", "guardrail_cost"), + [(0.40, 0.01), (0.05, 0.10)], + ids=["video_costs_more_than_its_guardrail", "video_costs_no_more_than_its_guardrail"], +) +async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_session_for_the_video( + video_cost: float, guardrail_cost: float +): from litellm.types.videos.main import VideoObject handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + completed_polls: Final = 3 with patch( "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" ) as mock_registry: mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) await handler.async_log_success_event( - _video_poll(0.01, guardrail_cost=0.01, status="in_progress"), + _video_poll(0.0, guardrail_cost=guardrail_cost, status="in_progress"), VideoObject(id="job-1", object="video", status="in_progress"), None, None, ) - for _ in range(3): + for _ in range(completed_polls): await handler.async_log_success_event( - _video_poll(0.41, guardrail_cost=0.01), VideoObject(id="job-1", object="video", status="completed"), None, None + _video_poll(video_cost, guardrail_cost=guardrail_cost), + VideoObject(id="job-1", object="video", status="completed"), + None, + None, ) spend = await handler._get_current_spend(handler._make_cache_key("session-jobs")) - assert spend == pytest.approx(0.01 + 0.41 + 2 * 0.01) + assert spend == pytest.approx(video_cost + (1 + completed_polls) * guardrail_cost) class _ClockedRedis: