scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

[Fix] Satisfy mypy on spend buffer restore helper

Browse source 
The daily queue parameter types on _restore_spend_updates_to_in_memory_queues
were narrowed to specific subtypes (DailyUserSpendTransaction, etc), but
the caller passes Dict[str, BaseDailySpendTransaction] — the return type
of flush_and_get_aggregated_daily_spend_update_transactions. Widen the
parameters to the base type.

Also replace dynamic TypedDict key lookup (which returned object) with
explicit literal-keyed get() calls so mypy can type-narrow each field.
This commit is contained in:
Yuneng Jiang 2026-04-22 14:31:25 -07:00
parent 288d403529
commit 3f42295d93
No known key found for this signature in database
1 changed files with 48 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

66
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -22,6 +22,7 @@ from litellm.constants import (
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
BaseDailySpendTransaction,
DailyAgentSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
@ -303,18 +304,18 @@ class RedisUpdateBuffer:
@staticmethod
async def _restore_spend_updates_to_in_memory_queues(
db_spend_update_transactions: Optional[DBSpendUpdateTransactions],
daily_spend_update_transactions: Optional[Dict[str, DailyUserSpendTransaction]],
daily_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]],
daily_team_spend_update_transactions: Optional[
Dict[str, DailyTeamSpendTransaction]
Dict[str, BaseDailySpendTransaction]
],
daily_org_spend_update_transactions: Optional[
Dict[str, DailyOrganizationSpendTransaction]
Dict[str, BaseDailySpendTransaction]
],
daily_end_user_spend_update_transactions: Optional[
Dict[str, DailyEndUserSpendTransaction]
Dict[str, BaseDailySpendTransaction]
],
daily_agent_spend_update_transactions: Optional[
Dict[str, DailyAgentSpendTransaction]
Dict[str, BaseDailySpendTransaction]
],
spend_update_queue: SpendUpdateQueue,
daily_spend_update_queue: DailySpendUpdateQueue,
@ -330,19 +331,46 @@ class RedisUpdateBuffer:
data aggregated during the current scheduler tick is permanently lost
because the source queues were already drained before the rpush.
"""
entity_type_field_pairs = [
(Litellm_EntityType.USER, "user_list_transactions"),
(Litellm_EntityType.END_USER, "end_user_list_transactions"),
(Litellm_EntityType.KEY, "key_list_transactions"),
(Litellm_EntityType.TEAM, "team_list_transactions"),
(Litellm_EntityType.TEAM_MEMBER, "team_member_list_transactions"),
(Litellm_EntityType.ORGANIZATION, "org_list_transactions"),
(Litellm_EntityType.TAG, "tag_list_transactions"),
(Litellm_EntityType.AGENT, "agent_list_transactions"),
]
if db_spend_update_transactions is not None:
for entity_type, field in entity_type_field_pairs:
entities = db_spend_update_transactions.get(field) or {} # type: ignore[call-overload]
entity_entries: List[
Tuple[Litellm_EntityType, Optional[Dict[str, float]]]
] = [
(
Litellm_EntityType.USER,
db_spend_update_transactions.get("user_list_transactions"),
),
(
Litellm_EntityType.END_USER,
db_spend_update_transactions.get("end_user_list_transactions"),
),
(
Litellm_EntityType.KEY,
db_spend_update_transactions.get("key_list_transactions"),
),
(
Litellm_EntityType.TEAM,
db_spend_update_transactions.get("team_list_transactions"),
),
(
Litellm_EntityType.TEAM_MEMBER,
db_spend_update_transactions.get("team_member_list_transactions"),
),
(
Litellm_EntityType.ORGANIZATION,
db_spend_update_transactions.get("org_list_transactions"),
),
(
Litellm_EntityType.TAG,
db_spend_update_transactions.get("tag_list_transactions"),
),
(
Litellm_EntityType.AGENT,
db_spend_update_transactions.get("agent_list_transactions"),
),
]
for entity_type, entities in entity_entries:
if not entities:
continue
for entity_id, cost in entities.items():
await spend_update_queue.add_update(
SpendUpdateQueueItem(
@ -352,7 +380,9 @@ class RedisUpdateBuffer:
)
)
daily_pairs = [
daily_pairs: List[
Tuple[Optional[Dict[str, BaseDailySpendTransaction]], DailySpendUpdateQueue]
] = [
(daily_spend_update_transactions, daily_spend_update_queue),
(daily_team_spend_update_transactions, daily_team_spend_update_queue),
(daily_org_spend_update_transactions, daily_org_spend_update_queue),