quivr/backend/models/user_usage.py
Zineb El Bachiri 9aaedcff51
Fix/requests limiting (#980)
* 🗃️ 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
2023-08-21 14:05:13 +02:00

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