diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ae3bc5663eb..ee7802a45d3 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -412,9 +412,11 @@ class DBSpendUpdateWriter: "spend": 0.0, }, ) - except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; a row it cannot take over charges the batch - verbose_proxy_logger.warning("Could not take over spend row %s for a batch's cost: %s", request_id, e) - return True + 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 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 33b7af06e1a..4efb94b60aa 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -3093,17 +3093,21 @@ async def test_update_database_charges_a_batch_whose_row_a_pre_upgrade_poll_left @pytest.mark.asyncio -async def test_update_database_charges_a_batch_whose_zero_row_it_could_not_take_over(): - """A DB that refuses the takeover must not swallow the batch's cost.""" +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 True + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False - assert db_writer._batch_database_updates.await_count == 1 + assert db_writer._batch_database_updates.await_count == 0 @pytest.mark.asyncio