diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 59e27496c27..fc04882244a 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -28,6 +28,7 @@ from litellm.proxy.db.db_spend_update_writer import ( ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.spend_event import ( + ObjectMapping, SpendEventBuildError, SpendEventDecodeError, build_spend_event, @@ -85,7 +86,9 @@ class _ProxyDBLogger(CustomLogger): super().__init__() self.spend_event_producer = spend_event_producer - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + 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): await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) return diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index d96cf1d0228..cddb69b03f7 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1881,9 +1881,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit(): } with ( - patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as mock_increment, # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), # test-quality-ok: same function-body import, no injection seam - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, # test-quality-ok: same function-body import, no injection seam + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() @@ -2210,8 +2216,12 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski logger = _ProxyDBLogger(producer) with ( - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, - patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as counters, + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as counters, ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() await logger.async_log_success_event(_offload_kwargs(), _offload_response(), datetime.now(), datetime.now()) @@ -2242,7 +2252,11 @@ async def test_async_log_success_event_keeps_batch_retrieves_in_process(): logger = _ProxyDBLogger(producer) kwargs = _batch_retrieve_kwargs(CallTypes.aretrieve_batch.value) - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging + ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() await logger.async_log_success_event( kwargs, _retrieved_batch("in_progress", output_file_id=None), datetime.now(), datetime.now() @@ -2253,9 +2267,15 @@ async def test_async_log_success_event_keeps_batch_retrieves_in_process(): async def _spend_row_written_by(run) -> tuple[SpendLogsPayload, dict, tuple[str, ...]]: with ( - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, - patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as counters, - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as counters, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True) mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() @@ -2310,8 +2330,12 @@ async def test_sidecar_writes_the_same_spend_row_and_counters_as_the_in_process_ @pytest.mark.asyncio -async def test_sidecar_ignores_an_undecodable_event(): - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: +async def test_sidecar_ignores_an_undecodable_event(): # test-quality-ok: a discarded event has no observable output other than the DB writer never being reached + with ( + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging + ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() await run_spend_event(b"garbage\n") mock_proxy_logging.db_spend_update_writer.update_database.assert_not_awaited() diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 82953b82cfb..22ab3a2d64c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -180,9 +180,7 @@ def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") - cost_session: Final = _get_batch_trace_session_id( - call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost" - ) + cost_session: Final = _get_batch_trace_session_id(call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost") assert create_session == cost_session == "batch-uid-1" @@ -4293,7 +4291,7 @@ ANTHROPIC_MESSAGES_SSE_CHUNKS: Final = ( 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' '"usage":{"output_tokens":4}}\n\n', - "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', ) @@ -4331,9 +4329,7 @@ def test_spend_log_request_id_is_the_message_id_a_non_streaming_messages_caller_ """ logging_obj = _anthropic_messages_logging_obj(stream=False) - logged_response = logging_obj._handle_anthropic_messages_response_logging( - result=ANTHROPIC_MESSAGES_RESPONSE - ) + logged_response = logging_obj._handle_anthropic_messages_response_logging(result=ANTHROPIC_MESSAGES_RESPONSE) assert logged_response.id == "msg_01Lit6806NonStreaming" assert ( @@ -4409,9 +4405,7 @@ def test_spend_log_request_id_still_falls_back_to_litellm_call_id_without_a_prov end_time=datetime.datetime.now(timezone.utc), logging_obj=logging_obj, ) - assert logging_obj.model_call_details["complete_streaming_response"].id == ( - "6806cafe-0000-4000-8000-000000000001" - ) + assert logging_obj.model_call_details["complete_streaming_response"].id == ("6806cafe-0000-4000-8000-000000000001") def test_spend_log_request_id_for_chat_completions_is_untouched():