From 3f42295d93ce322e45907415b53a660d06792337 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Wed, 22 Apr 2026 14:31:25 -0700
Subject: [PATCH] [Fix] Satisfy mypy on spend buffer restore helper
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 .../redis_update_buffer.py                    | 66 ++++++++++++++-----
 1 file changed, 48 insertions(+), 18 deletions(-)

diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
index 3a008d265fc..1e3014dbf3c 100644
--- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
+++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
@@ -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),