scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

style: apply black formatting to db_spend_update_writer and test

Browse source
This commit is contained in:
tibG 2026-03-23 13:48:25 +01:00
parent b0304ca589
commit b18d0cd83f
2 changed files with 29 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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"
)

4
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -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