mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-15 17:43:03 +03:00
9aaedcff51
* 🗃️ Rename users table into user_daily_usage * 💥 replace User model with UserIdentity model * 🗃️ New UserDailyUsage class for database interaction * 🐛 fix daily requests rate limiting per user * 🐛 fix user stats and properties update * ✏️ add typing and linting * 🚚 rename user_dialy_usage Class into user_usage & requests_count into daily_requests_count * 🚑 fix some rebase errors
55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
from logger import get_logger
|
|
from models.databases.supabase.supabase import SupabaseDB
|
|
from models.settings import get_supabase_db
|
|
from models.user_identity import UserIdentity
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class UserUsage(UserIdentity):
|
|
daily_requests_count: int = 0
|
|
|
|
def __init__(self, **data):
|
|
super().__init__(**data)
|
|
|
|
@property
|
|
def supabase_db(self) -> SupabaseDB:
|
|
return get_supabase_db()
|
|
|
|
def get_user_usage(self):
|
|
"""
|
|
Fetch the user request stats from the database
|
|
"""
|
|
request = self.supabase_db.get_user_usage(self.id)
|
|
|
|
return request
|
|
|
|
def handle_increment_user_request_count(self, date):
|
|
"""
|
|
Increment the user request count in the database
|
|
"""
|
|
current_requests_count = self.supabase_db.get_user_requests_count_for_day(
|
|
self.id, date
|
|
)
|
|
|
|
if current_requests_count is None:
|
|
if self.email is None:
|
|
raise ValueError("User Email should be defined for daily usage table")
|
|
self.supabase_db.create_user_daily_usage(
|
|
user_id=self.id, date=date, user_email=self.email
|
|
)
|
|
self.daily_requests_count = 1
|
|
return
|
|
|
|
self.supabase_db.increment_user_request_count(
|
|
user_id=self.id,
|
|
date=date,
|
|
current_requests_count=current_requests_count,
|
|
)
|
|
|
|
self.daily_requests_count = current_requests_count
|
|
|
|
logger.info(
|
|
f"User {self.email} request count updated to {current_requests_count}"
|
|
)
|