diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index dd7a24ae879..554fc7ed1d5 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -425,7 +425,7 @@ class _ProxyDBLogger(CustomLogger): async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime ) -> None: - if self.spend_event_producer is None or not is_offloadable_success(response_obj, kwargs.get("call_type")): + if self.spend_event_producer is None or not is_offloadable_success(response_obj, kwargs): await self._PROXY_track_cost_callback(dict(kwargs), response_obj, start_time, end_time) return event: Final = build_spend_event( diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py index d054519cb9c..52b9acf44c0 100644 --- a/litellm/proxy/spend_tracking/spend_event.py +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -23,10 +23,15 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.utils import LiteLLMBatch, Usage from litellm.types.videos.utils import is_video_status_call_type @@ -206,10 +211,19 @@ class SpendEventDecodeError: reason: str -def is_offloadable_success(response_obj: object, call_type: object) -> bool: - """Batch retrieves and interaction polls branch on the concrete response class, and a video status poll on - the status and video id the event leaves out, so they stay in-process.""" - return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) and not is_video_status_call_type(call_type) +def is_offloadable_success(response_obj: object, kwargs: ObjectMapping) -> bool: + """A success stays in-process when its cost logic reads something the spend event leaves out. + + Batches and interaction polls need the concrete response class. A batch retrieve needs the polled batch id, + which message redaction can hide. A video status poll needs the status and the video id. A background + response's settling read needs the settlement flag.""" + call_type: Final = kwargs.get("call_type") + return ( + not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) + and not is_batch_retrieve_call_type(call_type) + and not is_video_status_call_type(call_type) + and kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is not BackgroundResponseSettlement.SETTLES + ) def _json_fallback(value: object) -> str: diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 43f529faa6b..dcec934188f 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -18,7 +18,9 @@ from fastapi import FastAPI, HTTPException import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.litellm_core_utils.batch_utils import POLLED_BATCH_COST_IS_FINAL_KWARG, POLLED_BATCH_ID_KWARG from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.litellm_core_utils.redact_messages import perform_redaction from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded @@ -47,6 +49,10 @@ from litellm.proxy.utils import ( ) from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.types.router import Deployment, LiteLLM_Params +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject @@ -3692,17 +3698,41 @@ def _completed_batch() -> LiteLLMBatch: ) +def _redacted_completed_batch() -> object: + return perform_redaction({}, _completed_batch()) + + def _completed_video() -> VideoObject: return VideoObject(id="video_abc", object="video", status="completed") @pytest.mark.asyncio @pytest.mark.parametrize( - ("call_type", "polled_job"), - [(CallTypes.aretrieve_batch.value, _completed_batch), ("avideo_status", _completed_video)], + ("job_kwargs", "polled_job"), + [ + ({"call_type": CallTypes.aretrieve_batch.value}, _completed_batch), + ( + { + "call_type": CallTypes.aretrieve_batch.value, + POLLED_BATCH_ID_KWARG: "batch_abc", + POLLED_BATCH_COST_IS_FINAL_KWARG: True, + }, + _redacted_completed_batch, + ), + ({"call_type": "avideo_status"}, _completed_video), + ( + { + "call_type": "pass_through_endpoint", + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, + }, + _offload_response, + ), + ], + ids=["batch_retrieve", "redacted_batch_retrieve", "video_status_poll", "background_response_settling_read"], ) -async def test_async_log_success_event_keeps_job_polls_in_process(call_type, polled_job): - """The spend event leaves out the polled video's status and id, which decide whether the poll charges the job.""" +async def test_async_log_success_event_keeps_job_polls_in_process(job_kwargs, polled_job): + """The spend event leaves out the polled video's status and id and a background response's settlement flag, + which decide whether the poll charges the job and which spend row it writes.""" producer = SpendEventProducer( address=UnixAddress(path="/nonexistent/spend.sock"), on_unavailable="drop", @@ -3711,7 +3741,7 @@ async def test_async_log_success_event_keeps_job_polls_in_process(call_type, pol fallback=_no_fallback, ) logger = ProxyDBLogger(producer) - kwargs = {**_offload_kwargs(), "call_type": call_type} + kwargs = {**_offload_kwargs(), **job_kwargs} completed_job = polled_job() with ( diff --git a/tests/unit/proxy/spend_tracking/test_spend_event.py b/tests/unit/proxy/spend_tracking/test_spend_event.py index e1912d36d91..52422533213 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_event.py +++ b/tests/unit/proxy/spend_tracking/test_spend_event.py @@ -16,6 +16,10 @@ from litellm.proxy.spend_tracking.spend_event import ( is_offloadable_success, spend_event_callback_args, ) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.utils import LiteLLMBatch, ModelResponse, Usage _BIG_PROMPT: Final = "x" * 20_000 @@ -194,11 +198,23 @@ def test_undecodable_line_is_an_error_value(): assert isinstance(decode_spend_event(b"not json\n"), SpendEventDecodeError) +@pytest.mark.parametrize( + ("settlement", "offloadable"), + [(BackgroundResponseSettlement.SETTLES, False), (None, True), ("settles", True)], + ids=["settling_read_stays_in_process", "other_read_is_offloaded", "lookalike_string_is_offloaded"], +) +def test_only_a_background_response_settling_read_stays_in_process( + settlement: BackgroundResponseSettlement | str | None, offloadable: bool +) -> None: + kwargs: Final = {"call_type": "pass_through_endpoint", SETTLES_BACKGROUND_RESPONSE_KWARG: settlement} + assert is_offloadable_success(None, kwargs) is offloadable + + def test_batch_retrieves_and_video_status_polls_stay_in_process(): - assert is_offloadable_success(_response(), "acompletion") is True - assert is_offloadable_success(None, None) is True - assert is_offloadable_success(None, "avideo_generation") is True - assert is_offloadable_success(None, "avideo_status") is False + assert is_offloadable_success(_response(), {"call_type": "acompletion"}) is True + assert is_offloadable_success(None, {}) is True + assert is_offloadable_success(None, {"call_type": "avideo_generation"}) is True + assert is_offloadable_success(None, {"call_type": "avideo_status"}) is False assert ( is_offloadable_success( LiteLLMBatch( @@ -210,7 +226,7 @@ def test_batch_retrieves_and_video_status_polls_stay_in_process(): object="batch", status="completed", ), - "aretrieve_batch", + {"call_type": "aretrieve_batch"}, ) is False )