Merge pull request #47 from pattern-tech/feat/wallet-accounting

fix[api]: query usage route edited

Author: yasinfakhar

api/requirements.txt

@@ -32,3 +32,4 @@ langchain-groq==0.2.4
 langchain-huggingface==0.1.2
 sentry-sdk==2.22.0
 siwe==4.4.0
+bcrypt==4.0.1

api/src/agentflow/providers/goldrush_tools.py

@@ -68,6 +68,7 @@ def get_wallet_activity(wallet_address: str, output_include: list[str]) -> List[
 def get_balance_for_address(wallet_address: str, output_include: list[str]) -> str:
     """
      fetch the native, fungible (ERC20), and non-fungible (ERC721 & ERC1155) tokens held by an address
+     apply decimal conversion for balance
 
     Args:
         wallet_address (str): The wallet address to retrieve balance for.

api/src/agentflow/providers/moralis_tools.py

@@ -54,6 +54,7 @@ def get_wallet_active_chains(wallet_address: str, output_include: list[str]) ->
 def get_wallet_token_balances(wallet_address: str, output_include: list[str], cursor: str = None) -> dict:
     """
     Get token balances for a specific wallet address and their token prices in USD. (paginated)
+    apply decimal conversion for balance
 
     Args:
         wallet_address (str): Ethereum wallet address

api/src/auth/services/auth_service.py

@@ -115,7 +115,7 @@ def authenticate_user(self, email: str, password: str, db: Session):
         """
         # Fetch the user from the database using the provided email
         user = db.query(UserModel).filter_by(email=email).first()
-        print(password, user.password)
+
         if not user:
             raise HTTPException(
                 status_code=401, detail="Incorrect email or password")

api/src/conversation/routers/playground_conversation_router.py

@@ -9,6 +9,7 @@
 from src.db.sql_alchemy import Database
 from src.util.response import global_response
 from src.auth.utils.get_token import authenticate_user
+from src.query_usage.services.query_usage_service import QueryUsageService
 from src.conversation.services.conversation_service import ConversationService
 
 router = APIRouter(prefix="/playground/conversation")
@@ -33,6 +34,13 @@ def get_conversation_service() -> ConversationService:
     return ConversationService()
 
 
+def get_query_usage_service() -> QueryUsageService:
+    """
+    Dependency to instantiate the QueryUsageService.
+    """
+    return QueryUsageService()
+
+
 class CreateConversationInput(BaseModel):
     """
     Schema for creating a conversation.
@@ -41,7 +49,7 @@ class CreateConversationInput(BaseModel):
     project_id: UUID
 
     class Config:
-        orm_mode = True
+        from_attributes = True
 
 
 class ConversationOutput(BaseModel):
@@ -53,7 +61,7 @@ class ConversationOutput(BaseModel):
     project_id: UUID
 
     class Config:
-        orm_mode = True
+        from_attributes = True
 
 
 class MessageType(str, Enum):
@@ -254,7 +262,9 @@ async def send_message(
     conversation_id: UUID,
     project_id: UUID,
     db: Session = Depends(get_db),
-    service: ConversationService = Depends(get_conversation_service),
+    query_usage_service: QueryUsageService = Depends(get_query_usage_service),
+    conversation_service: ConversationService = Depends(
+        get_conversation_service),
     user_id: UUID = Depends(authenticate_user),
 ):
     """
@@ -272,21 +282,29 @@ async def send_message(
         dict: A JSON response containing the complete message data if `stream` is false.
     """
     try:
-        service.check_user_eligibility(db, user_id)
+
+        max_query_allowance = query_usage_service.get_user_max_query_allowance(
+            db, user_id)
+        is_eligible = query_usage_service.check_user_eligibility(
+            db, user_id, max_query_allowance)
+        if not is_eligible:
+            raise Exception(
+                "You have reached your daily query limit. Please try again tomorrow or stake more to get more queries."
+            )
 
         if input.stream:
             return StreamingResponse(
-                service.send_message(db,
-                                     input.message,
-                                     user_id,
-                                     conversation_id,
-                                     input.message_type,
-                                     input.stream),
+                conversation_service.send_message(db,
+                                                  input.message,
+                                                  user_id,
+                                                  conversation_id,
+                                                  input.message_type,
+                                                  input.stream),
                 media_type="text/plain"
             )
         else:
             response = None
-            async for item in service.send_message(
+            async for item in conversation_service.send_message(
                 db,
                 input.message,
                 user_id,

api/src/conversation/services/conversation_service.py

@@ -9,7 +9,6 @@
 from src.agentflow.agents.hub import AgentHub
 from src.db.models import Conversation, QueryUsage
 from src.user.services.user_service import UserService
-from src.share.staked_tokens import get_user_staked_tokens
 from src.agent.services.memory_service import MemoryService
 from src.project.services.project_service import ProjectService
 from src.agent.services.agent_service import RouterAgentService
@@ -138,56 +137,6 @@ def get_project_associated_with_conversation(self, db_session: Session, conversa
         """
         return self.repository.get_project_associated_with_conversation(db_session, conversation_id)
 
-    def check_user_eligibility(self, db_session: Session, user_id: UUID) -> bool:
-        """
-        Checks if a user is eligible to use the service by checking their payment status.
-
-        Args:
-            db_session (Session): The database session.
-            user_id (UUID): The ID of the user to check.
-
-        Returns:
-            bool: True if the user is eligible, False otherwise.
-
-        Raises:
-            Exception: If the user is not eligible, an exception is raised with a message explaining why.
-        """
-        user = self.user_service.get_user(db_session, user_id)
-
-        whitelist = self.user_service.get_whitelist(db_session)
-
-        # check user payment
-        for wl in whitelist:
-            if str(user_id) == str(wl.user_id):
-                max_allowed_query = wl.max_query
-                break
-        else:
-            staked_morpheus = get_user_staked_tokens(
-                wallet_address=user.wallet_address, provider="morpheus")
-
-            if staked_morpheus == 0:
-                raise Exception(
-                    "You need to stake Morpheus tokens to use this service")
-
-            usage_setting = self.query_usage_service.get_usage_setting(
-                db_session)
-            max_allowed_query = 0
-            for setting in usage_setting:
-                if setting.provider == "morpheus":
-                    max_allowed_query = setting.max_query * \
-                        (int(staked_morpheus) / 1e18)
-
-        user_query_usage_until_previous_24h = self.query_usage_service.get_all_query_usages(
-            db_session, user_id, "morpheus",
-            timedelta(hours=24))
-
-        if len(user_query_usage_until_previous_24h) >= max_allowed_query:
-            raise Exception(
-                "You have reached your daily query limit. Please try again tomorrow or stake more to get more queries."
-            )
-
-        return True
-
     async def send_message(
         self,
         db_session: Session,

api/src/project/routers/project_router.py

@@ -39,7 +39,7 @@ class CreateProjectInput(BaseModel):
     workspace_id: UUID
 
     class Config:
-        orm_mode = True
+        from_attributes = True
 
 
 class ProjectOutput(BaseModel):
@@ -51,7 +51,7 @@ class ProjectOutput(BaseModel):
     workspace_id: UUID
 
     class Config:
-        orm_mode = True
+        from_attributes = True
 
 
 @router.post(

api/src/query_usage/routers/query_usage_router.py

@@ -1,8 +1,8 @@
 from uuid import UUID
+from datetime import timedelta
 from typing import List, Optional
 from sqlalchemy.orm import Session
 from pydantic import BaseModel, Field
-from datetime import datetime, timedelta
 from fastapi import APIRouter, Depends, HTTPException, status
 
 from src.db.sql_alchemy import Database
@@ -40,7 +40,7 @@ class CreateQueryUsageInput(BaseModel):
     provider: str = Field(..., example="morpheus")
 
     class Config:
-        orm_mode = True
+        from_attributes = True
 
 
 class QueryUsageOutput(BaseModel):
@@ -51,7 +51,7 @@ class QueryUsageOutput(BaseModel):
     provider: str = Field(..., example="morpheus")
 
     class Config:
-        orm_mode = True
+        from_attributes = True
 
 
 @router.get(
@@ -89,29 +89,35 @@ def get_query_usage(
 @router.get(
     "",
     response_model=List[QueryUsageOutput],
-    summary="List All Query Usages",
-    description="Lists all query usage records for the authenticated user.",
-    response_description="A list of all query usage records."
+    summary="Get user query usage",
+    description="Get number of used and total number of allowed query for a user",
+    response_description="Number of used and total number of allowed query for a user"
 )
-def get_all_query_usages(
+def get_user_query_usages(
     provider: Optional[str] = None,
     duration: Optional[timedelta] = timedelta(hours=24),
     user_id: UUID = Depends(authenticate_user),
     db: Session = Depends(get_db),
     service: QueryUsageService = Depends(get_query_usage_service),
 ):
     """
-    List all query usage records for the authenticated user.
+    Get user query usage
 
     - **provider**: Optional filter by provider.
     - **duration**: Optional filter by a specific datetime.
     - **user_id**: The authenticated user's ID.
     - **db**: Database session.
     - **service**: QueryUsage service handling business logic.
+    - **conversation_service**: Conversation service handling business logic.
 
     Returns:
-        List[QueryUsageOutput]: A list of all the user's query usage records.
+        dict: A dictionary containing the number of used and total number of allowed query.
     """
     query_usages = service.get_all_query_usages(
         db, user_id, provider, duration)
-    return global_response(query_usages)
+
+    data = {
+        "query_usages": len(query_usages),
+        "max_query_allowance_per_day": service.get_user_max_query_allowance(db, user_id)
+    }
+    return global_response(data)

api/src/query_usage/services/query_usage_service.py

@@ -1,16 +1,19 @@
 from uuid import UUID
-from datetime import datetime, timedelta
 from typing import List, Optional
 from sqlalchemy.orm import Session
+from datetime import datetime, timedelta
 
 from src.db.models import QueryUsage
 from src.share.base_service import BaseService
+from src.user.services.user_service import UserService
+from src.share.staked_tokens import get_user_staked_tokens
 from src.query_usage.repositories.query_usage_repository import QueryUsageRepository
 
 
 class QueryUsageService(BaseService):
     def __init__(self):
         self.repository = QueryUsageRepository()
+        self.user_service = UserService()
 
     def create_query_usage(self, db_session: Session, query_usage: QueryUsage) -> QueryUsage:
         """
@@ -65,3 +68,67 @@ def get_usage_setting(self, db_session: Session) -> List[QueryUsage]:
             List[UsageSetting]: A list of all usage settings.
         """
         return self.repository.get_usage_setting(db_session)
+
+    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.
+
+        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)
+
+        # 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
+
+        staked_morpheus = get_user_staked_tokens(
+            wallet_address=user.wallet_address, provider="morpheus")
+
+        if staked_morpheus == 0:
+            raise Exception(
+                "You need to stake Morpheus tokens to use this service")
+
+        usage_setting = self.get_usage_setting(
+            db_session)
+
+        max_allowed_query = 0
+        for setting in usage_setting:
+            if setting.provider == "morpheus":
+                max_allowed_query = setting.max_query * \
+                    (int(staked_morpheus) / 1e18)
+                break
+
+        return max_allowed_query
+
+    def check_user_eligibility(self, db_session: Session, user_id: UUID, max_query_allowance: int) -> bool:
+        """
+        Checks if a user is eligible to make a query based on their daily query limit.
+
+        Args:
+            db_session (Session): The database session.
+            user_id (UUID): The ID of the user.
+            max_query_allowance (int): The number of queries the user is allowed to make.
+
+        Returns:
+            bool: True if the user is eligible, False otherwise.
+        """
+        user_query_usage_until_previous_24h = self.get_all_query_usages(
+            db_session, user_id, "morpheus", timedelta(hours=24)
+        )
+
+        if len(user_query_usage_until_previous_24h) >= max_query_allowance:
+            return False
+
+        return True