From b18d0cd83f05f2964bb3de25f668297e5603b8ee Mon Sep 17 00:00:00 2001 From: tibG Date: Mon, 23 Mar 2026 13:48:25 +0100 Subject: [PATCH] style: apply black formatting to db_spend_update_writer and test --- litellm/proxy/db/db_spend_update_writer.py | 39 +++++++++++++------ .../proxy/db/test_db_spend_update_writer.py | 4 +- 2 files changed, 29 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 85c4c71c604..9c725e2fda4 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1578,7 +1578,8 @@ class DBSpendUpdateWriter: "mcp_namespaced_tool_name" ) or "", - "endpoint": transaction.get("endpoint") or "", + "endpoint": transaction.get("endpoint") + or "", } } @@ -1609,22 +1610,29 @@ class DBSpendUpdateWriter: "successful_requests": transaction[ "successful_requests" ], - "failed_requests": transaction["failed_requests"], + "failed_requests": transaction[ + "failed_requests" + ], } # Add cache-related fields if they exist if "cache_read_input_tokens" in transaction: - common_data["cache_read_input_tokens"] = ( - transaction.get("cache_read_input_tokens", 0) + common_data[ + "cache_read_input_tokens" + ] = transaction.get( + "cache_read_input_tokens", 0 ) if "cache_creation_input_tokens" in transaction: - common_data["cache_creation_input_tokens"] = ( - transaction.get( - "cache_creation_input_tokens", 0 - ) + common_data[ + "cache_creation_input_tokens" + ] = transaction.get( + "cache_creation_input_tokens", 0 ) - if entity_type == "tag" and "request_id" in transaction: + if ( + entity_type == "tag" + and "request_id" in transaction + ): common_data["request_id"] = transaction.get( "request_id" ) @@ -1635,14 +1643,18 @@ class DBSpendUpdateWriter: "increment": transaction["prompt_tokens"] }, "completion_tokens": { - "increment": transaction["completion_tokens"] + "increment": transaction[ + "completion_tokens" + ] }, "spend": {"increment": transaction["spend"]}, "api_requests": { "increment": transaction["api_requests"] }, "successful_requests": { - "increment": transaction["successful_requests"] + "increment": transaction[ + "successful_requests" + ] }, "failed_requests": { "increment": transaction["failed_requests"] @@ -1663,7 +1675,10 @@ class DBSpendUpdateWriter: ) } - if entity_type == "tag" and "request_id" in transaction: + if ( + entity_type == "tag" + and "request_id" in transaction + ): update_data["request_id"] = transaction.get( "request_id" ) 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 81fd5a1d2a9..5e716fc5055 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 @@ -1444,8 +1444,8 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() # Return all-None tuple (no data to commit) - mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( - return_value=(None, None, None, None, None, None, None) + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = ( + AsyncMock(return_value=(None, None, None, None, None, None, None)) ) db_writer.redis_update_buffer = mock_redis_update_buffer