From cc9f47b1d29f1572e56026dad0c6d077be2cebbb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 15:17:01 -0700 Subject: [PATCH] fix loc of DBSpendUpdateTransactions --- litellm/proxy/_types.py | 9 +++++++++ litellm/proxy/db/db_spend_update_writer.py | 6 ++---- litellm/proxy/db/redis_update_buffer.py | 12 ++---------- 3 files changed, 13 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 17b4acd138..1c3251c12a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2733,3 +2733,12 @@ class DailyUserSpendTransaction(TypedDict): prompt_tokens: int completion_tokens: int spend: float + + +class DBSpendUpdateTransactions(TypedDict): + user_list_transactions: Optional[Dict[str, float]] + end_user_list_transactions: Optional[Dict[str, float]] + key_list_transactions: Optional[Dict[str, float]] + team_list_transactions: Optional[Dict[str, float]] + team_member_list_transactions: Optional[Dict[str, float]] + org_list_transactions: Optional[Dict[str, float]] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 17247aff74..8f175fc905 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -18,15 +18,13 @@ from litellm.caching import DualCache, RedisCache from litellm.constants import DB_SPEND_UPDATE_JOB_NAME from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DBSpendUpdateTransactions, Litellm_EntityType, LiteLLM_UserTable, SpendLogsPayload, ) from litellm.proxy.db.pod_lock_manager import PodLockManager -from litellm.proxy.db.redis_update_buffer import ( - DBSpendUpdateTransactions, - RedisUpdateBuffer, -) +from litellm.proxy.db.redis_update_buffer import RedisUpdateBuffer if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 2207911024..6f334e90ed 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -4,10 +4,11 @@ Handles buffering database `UPDATE` transactions in Redis before committing them This is to prevent deadlocks and improve reliability """ -from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Dict, Optional, Union, cast from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache +from litellm.proxy._types import DBSpendUpdateTransactions from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: @@ -16,15 +17,6 @@ else: PrismaClient = Any -class DBSpendUpdateTransactions(TypedDict): - user_list_transactions: Optional[Dict[str, float]] - end_user_list_transactions: Optional[Dict[str, float]] - key_list_transactions: Optional[Dict[str, float]] - team_list_transactions: Optional[Dict[str, float]] - team_member_list_transactions: Optional[Dict[str, float]] - org_list_transactions: Optional[Dict[str, float]] - - class RedisUpdateBuffer: """ Handles buffering database `UPDATE` transactions in Redis before committing them to the database