Merge pull request #82 from pattern-tech/refactor/ai-v2

feat: enhance user whitelist management and automatic entry creation

Author: yasinfakhar

src/auth/services/auth_service.py

@@ -81,7 +81,7 @@ def register(self, input: RegisterInput, db: Session) -> str:
         existing_user = db.query(UserModel).filter_by(
             wallet_address=input.wallet_address).first()
         if existing_user:
-            raise AlreadyExistsError("User already exists")
+            raise AlreadyExistsError("User already exists with this wallet address")
 
         # # Create a new user record
         if input.email and input.password:

src/conversation/services/conversation_service.py

@@ -30,10 +30,11 @@ class ConversationService:
     def __init__(self):
         self.repository = ConversationRepository()
         self.project_repository = ProjectRepository()
+
+        self.user_service = UserService()
         self.memory_service = MemoryService()
         self.project_service = ProjectService()
         self.query_usage_service = QueryUsageService()
-        self.user_service = UserService()
 
     def create_conversation(
         self, db_session: Session, name: str, project_id: UUID, user_id: UUID, conversation_id: UUID = None

src/query_usage/repositories/query_usage_repository.py

@@ -140,6 +140,25 @@ def get_user_query_count_for_today(self, db_session: Session, user_id: UUID, pro
 
         return result_count, result_oldest, next_reset
 
+    def get_user_whitelist_allowance(self, db_session: Session, user_id: UUID) -> Optional[int]:
+        """
+        Retrieves the user's whitelist max_query allowance directly from the database.
+
+        Args:
+            db_session (Session): The database session to use.
+            user_id (UUID): The unique identifier of the user.
+
+        Returns:
+            Optional[int]: The user's max query allowance from the whitelist, or None if not whitelisted.
+        """
+        from src.db.models import WhiteList
+
+        whitelist_entry = db_session.query(WhiteList.max_query)\
+            .filter(WhiteList.user_id == user_id)\
+            .first()
+
+        return whitelist_entry.max_query if whitelist_entry else None
+
     def update(self, db_session, id: UUID) -> None:
         pass
 

src/query_usage/routers/query_usage_router.py

@@ -59,6 +59,10 @@ class TodayQueryUsage(BaseModel):
     today_query_count: int = Field(..., example=2)
     remaining_query_allowance: int = Field(..., example=3)
     max_query_allowance_per_day: int = Field(..., example=5)
+    # Base allowance from whitelist
+    whitelist_allowance: int = Field(..., example=20)
+    # Additional allowance from staked tokens
+    token_based_allowance: int = Field(..., example=5)
     next_reset_time: datetime = Field(..., example="2025-03-15T15:30:20+03:30")
 
 
@@ -142,14 +146,23 @@ def get_user_daily_query_usages(
         dict: A dictionary containing:
             - today_query_count: Number of queries used today
             - max_query_allowance_per_day: Total number of allowed queries per day
+            - whitelist_allowance: Base query allowance from whitelist (default 20)
+            - token_based_allowance: Additional allowance from staked tokens
             - remaining_queries_today: Number of remaining queries for today
             - next_reset_time: The datetime when the query usage will reset to zero
     """
     try:
-        # Get the max query allowance for this user
-        max_query_allowance = service.get_user_max_query_allowance(db, user_id)
+        # Get the user's base whitelist allowance
+        whitelist_allowance = service.get_user_whitelist_allowance(db, user_id)
 
-        # Get today's query count using the new repository method
+        # Get the token-based allowance using our new dedicated method
+        token_based_allowance = service.get_user_token_based_allowance(
+            db, user_id)
+
+        # Calculate total max query allowance (whitelist + token-based)
+        max_query_allowance = whitelist_allowance + token_based_allowance
+
+        # Get today's query count
         today_query_count, _, next_reset_time = service.get_user_query_count_for_today(
             db, user_id, provider)
 

src/query_usage/services/query_usage_service.py

@@ -77,45 +77,76 @@ def get_usage_setting(self, db_session: Session) -> List[QueryUsage]:
     def get_user_max_query_allowance(self, db_session: Session, user_id: UUID) -> int:
         """
         Retrieves the maximum number of queries a user is allowed to make.
+        This is calculated as the sum of:
+        1. User's whitelist allowance
+        2. Allowance based on staked Morpheus tokens
 
         Args:
             db_session (Session): The database session.
             user_id (UUID): The ID of the user.
 
         Returns:
             int: The maximum number of queries the user is allowed to make.
-
-        Raises:
-            Exception: If the user has not staked any Morpheus tokens.
         """
         user = self.user_service.get_user(db_session, user_id)
 
-        whitelist = self.user_service.get_whitelist(db_session)
+        # Get the user's base whitelist allowance using the dedicated method
+        whitelist_allowance = self.get_user_whitelist_allowance(
+            db_session, user_id)
+
+        # Calculate additional allowance based on staked tokens
+        token_based_allowance = self.get_user_token_based_allowance(
+            db_session, user_id)
+
+        # Total allowance is the sum of whitelist allowance and token-based allowance
+        max_allowed_query = whitelist_allowance + token_based_allowance
+
+        return max_allowed_query
+
+    def get_user_token_based_allowance(self, db_session: Session, user_id: UUID) -> int:
+        """
+        Calculates a user's token-based query allowance based on their staked MOR tokens.
 
-        # check user payment
-        for wl in whitelist:
-            if str(user_id) == str(wl.user_id):
-                max_allowed_query = wl.max_query
-                return max_allowed_query
+        Args:
+            db_session (Session): The database session.
+            user_id (UUID): The ID of the user.
 
+        Returns:
+            int: The token-based query allowance (0 if no tokens are staked)
+        """
+        user = self.user_service.get_user(db_session, user_id)
+
+        token_based_allowance = 0
         staked_morpheus = get_user_staked_tokens(
             wallet_address=user.wallet_address.lower(), provider="morpheus")
 
-        if staked_morpheus == 0:
-            raise RateLimitError(
-                "You need to stake Morpheus tokens to use this service")
+        if staked_morpheus > 0:
+            usage_setting = self.get_usage_setting(db_session)
+            for setting in usage_setting:
+                if setting.provider == "morpheus":
+                    token_based_allowance = setting.max_query * \
+                        int(int(staked_morpheus) / 1e18)
+                    break
 
-        usage_setting = self.get_usage_setting(
-            db_session)
+        return token_based_allowance
 
-        max_allowed_query = 0
-        for setting in usage_setting:
-            if setting.provider == "morpheus":
-                max_allowed_query = setting.max_query * \
-                    int(int(staked_morpheus) / 1e18)
-                break
+    def get_user_whitelist_allowance(self, db_session: Session, user_id: UUID) -> int:
+        """
+        Retrieves the base query allowance for a user from the whitelist using a direct database query.
 
-        return max_allowed_query
+        Args:
+            db_session (Session): The database session.
+            user_id (UUID): The ID of the user.
+
+        Returns:
+            int: The base query allowance from the whitelist (default: 0 if not whitelisted)
+        """
+        # Use the repository to directly query the database for whitelist allowance
+        allowance = self.repository.get_user_whitelist_allowance(
+            db_session, user_id)
+
+        # Return the allowance if found, otherwise return 0
+        return allowance if allowance is not None else 0
 
     def get_user_query_count_for_today(self, db_session: Session, user_id: UUID, provider: Optional[str] = None) -> Tuple[int, Optional[datetime]]:
         """

src/user/repositories/user_repository.py

@@ -122,3 +122,21 @@ def get_whitelist(self, db_session: Session) -> list[WhiteList]:
             List[WhiteList]: A list of all whitelisted users.
         """
         return db_session.query(WhiteList).all()
+
+    def create_whitelist_entry(self, db_session: Session, user_id: UUID, max_query: int = 20) -> WhiteList:
+        """
+        Creates a whitelist entry for a user with a default max query limit of 20.
+
+        Args:
+            db_session (Session): The database session to use.
+            user_id (UUID): The unique identifier of the user to whitelist.
+            max_query (int): The maximum number of queries allowed for this user.
+
+        Returns:
+            WhiteList: The created whitelist entry.
+        """
+        whitelist_entry = WhiteList(user_id=user_id, max_query=max_query)
+        db_session.add(whitelist_entry)
+        db_session.commit()
+        db_session.refresh(whitelist_entry)
+        return whitelist_entry

src/user/services/user_service.py

@@ -2,7 +2,7 @@
 from typing import List
 from sqlalchemy.orm import Session
 
-from src.db.models import UserModel, WhiteList
+from src.db.models import UserModel, WhiteList, UsageSetting
 from src.user.repositories.user_repository import UserRepository
 
 
@@ -13,7 +13,7 @@ def __init__(self):
     def create_user(self, db_session: Session, wallet_address: str, chain_id: int, email: str = None, password: str = None,
                     ) -> UserModel:
         """
-        Creates a new user.
+        Creates a new user and automatically whitelists them with the default max query value from database.
 
         Args:
             db_session (Session): The database session.
@@ -31,6 +31,16 @@ def create_user(self, db_session: Session, wallet_address: str, chain_id: int, e
                           chain_id=chain_id, email=email, password=password)
         user = self.repository.create(db_session, _user)
 
+        default_max_query = 0
+        usage_settings = db_session.query(UsageSetting).filter(
+            UsageSetting.provider == "pattern").first()
+        if usage_settings:
+            default_max_query = usage_settings.max_query
+
+        # Automatically whitelist the user with the default query credits from database
+        self.repository.create_whitelist_entry(
+            db_session, user.id, max_query=default_max_query)
+
         return user
 
     def get_user(self, db_session: Session, user_id: UUID) -> UserModel: