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

Feat/wallet accounting

Author: yasinfakhar

.github/workflows/dev.yml

@@ -73,7 +73,7 @@ jobs:
               -e MORALIS_API_KEY=${{ secrets.MORALIS_API_KEY }} \
               -e ETHER_SCAN_API_KEY=${{ secrets.ETHER_SCAN_API_KEY }} \
               -e GOLDRUSH_API_KEY=${{ secrets.GOLDRUSH_API_KEY }} \
-              -e ETH_RPC_URL=${{ secrets.ETH_RPC_URL }} \
+              -e ETH_RPC=${{ secrets.ETH_RPC }} \
               -e LLM_PROVIDER=${{ secrets.LLM_PROVIDER }} \
               -e LLM_MODEL=${{ secrets.LLM_MODEL }} \
               -e LLM_API_KEY=${{ secrets.LLM_API_KEY }} \

api/.env.sample

@@ -55,7 +55,7 @@ EXA_API_KEY=
 PERPLEXITY_API_KEY=
 TAVILY_API_KEY=
 
-ETH_RPC_URL=
+ETH_RPC=
 
 
 #------------------------

api/src/__init__.py

@@ -5,6 +5,7 @@
 from src.agent.routers import agent_router
 from src.project.routers import project_router
 from src.workspace.routers import workspace_router
+from src.query_usage.routers import query_usage_router
 from src.conversation.routers import playground_conversation_router
 
 api_router = APIRouter()
@@ -14,4 +15,5 @@
 api_router.include_router(project_router.router, tags=["Project"])
 api_router.include_router(agent_router.router, tags=["Agent"])
 api_router.include_router(
-    playground_conversation_router.router, tags=["Conversation"])
+    playground_conversation_router.router, tags=["Conversation"])
+api_router.include_router(query_usage_router.router, tags=["Query Usage"])

api/src/auth/routers/auth_router.py

@@ -1,14 +1,15 @@
+from sqlalchemy.orm import Session
+from fastapi import APIRouter, Depends
 from src.auth.utils.bcrypt_helper import generate_access_token
 from src.auth.services.auth_service import (
     AuthService,
     LoginInput,
     RegisterInput,
     VerifyInput
 )
+
 from src.db.sql_alchemy import Database
 from src.util.response import global_response
-from fastapi import APIRouter, Depends, HTTPException
-from sqlalchemy.orm import Session
 
 router = APIRouter(prefix="/auth")
 database = Database()
@@ -69,19 +70,18 @@ def login(input: LoginInput, db: Session = Depends(database.get_db)):
     summary="Verify a Signature",
     description="Verify a signature according to SIWE spec",
 )
-def verify(input: VerifyInput):
+def verify(input: VerifyInput, db: Session = Depends(database.get_db)):
     """
     Verify a signature according to SIWE spec and return an access token
 
     - **message**: A SIWE message
     - **signature**: User's signature for the message
     """
-    verification_result = auth.verify_signature(
-        input.message, input.signature)
 
-    payload = {"id": "{}:{}".format(
-        verification_result["chain_id"], verification_result["address"])}
+    user = auth.verify_signature(
+        input.message, input.signature, db)
 
+    payload = {"id": str(user.id)}
     token = generate_access_token(data=payload)
     return global_response(
         {

api/src/auth/services/auth_service.py

@@ -1,23 +1,35 @@
+from typing import Optional
+from siwe import SiweMessage
 from fastapi import HTTPException
-from pydantic import BaseModel, Field, EmailStr
 from sqlalchemy.orm import Session
-from src.workspace.services.workspace_service import WorkspaceService
-from src.auth.utils.bcrypt_helper import hash_password, verify_password
+from pydantic import BaseModel, Field, EmailStr
+
 from src.db.models import UserModel
 from src.db.sql_alchemy import Database
-from siwe import SiweMessage
+from src.share.base_types import WalletAddress
+from src.user.services.user_service import UserService
+from src.workspace.services.workspace_service import WorkspaceService
+from src.auth.utils.bcrypt_helper import hash_password, verify_password
 
 database = Database()
 
 
 class RegisterInput(BaseModel):
-    email: EmailStr = Field(
-        ..., example="user@example.com", description="The email address of the user"
+    email: Optional[EmailStr] = Field(
+        None, example="user@example.com", description="The email address of the user"
     )
-    password: str = Field(
+    password: Optional[str] = Field(
+        None, example="securepassword123", description="The password for the user account"
+    )
+    wallet_address: WalletAddress = Field(
         ...,
-        example="securepassword123",
-        description="The password for the user account",
+        example="0x0...",
+        description="The wallet address of the user"
+    )
+    chain_id: int = Field(
+        ...,
+        example="1",
+        description="The chain id of the user"
     )
 
 
@@ -47,34 +59,44 @@ class AuthService:
     """
 
     def __init__(self):
-        self.workspace = WorkspaceService()
+        self.workspace_service = WorkspaceService()
+        self.user_service = UserService()
 
     def register(self, input: RegisterInput, db: Session) -> str:
         """
         Registers a new user by saving their details into the database.
 
         Args:
-            input (RegisterInput): The registration input containing email and password.
+            input (RegisterInput): The registration input .
             db (Session): The database session for executing queries.
 
         Returns:
             str: Success message indicating the user was registered.
 
         Raises:
-            HTTPException: If a user with the same email already exists.
+            HTTPException: If a user with the same wallet_address already exists.
         """
         # Check if the user already exists
-        existing_user = db.query(UserModel).filter_by(email=input.email.lower()).first()
+        existing_user = db.query(UserModel).filter_by(
+            wallet_address=input.wallet_address).first()
         if existing_user:
             raise HTTPException(status_code=400, detail="User already exists")
 
-        # Hash the user's password and create a new user record
-        hashed_password = hash_password(input.password)
-        new_user = UserModel(email=input.email.lower(), password=hashed_password)
-        db.add(new_user)
-        db.commit()
+        # # Create a new user record
+        if input.email and input.password:
+            existing_user = db.query(UserModel).filter_by(
+                email=input.email.lower()).first()
+            if existing_user:
+                raise HTTPException(
+                    status_code=400, detail="This email is already exists")
+
+        if input.password:
+            input.password = hash_password(input.password)
+
+        new_user = self.user_service.create_user(
+            db, input.wallet_address, input.chain_id, input.email, input.password, )
 
-        self.workspace.create_workspace(db, "Default", new_user.id)
+        self.workspace_service.create_workspace(db, "Default", new_user.id)
 
         return new_user
 
@@ -93,16 +115,19 @@ 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")
+            raise HTTPException(
+                status_code=401, detail="Incorrect email or password")
 
         # Verify the provided password matches the stored hash
         if not verify_password(password, user.password):
-            raise HTTPException(status_code=401, detail="Incorrect email or password")
+            raise HTTPException(
+                status_code=401, detail="Incorrect email or password")
 
         return user
 
-    def verify_signature(self, message: str, signature: str):
+    def verify_signature(self, message: str, signature: str, db: Session):
         """
         Verify a signature according to SIWE spec
 
@@ -116,7 +141,6 @@ def verify_signature(self, message: str, signature: str):
         Raises:
             HTTPException: If message cannot be parsed or verified
         """
-
         try:
             siwe_message = SiweMessage.from_message(message)
             siwe_message.verify(signature)
@@ -128,7 +152,15 @@ def verify_signature(self, message: str, signature: str):
             raise HTTPException(
                 status_code=401, detail="Signature is not valid")
 
-        return {
-            "chain_id": siwe_message.chain_id,
-            "address": siwe_message.address,
-        }
+        # check if not exist create new user
+        user = self.user_service.get_user_by_wallet_address(
+            siwe_message.address, db)
+        if not user:
+            user = self.register(RegisterInput(
+                email=None,
+                password=None,
+                wallet_address=siwe_message.address,
+                chain_id=siwe_message.chain_id
+            ), db)
+
+        return user

api/src/conversation/routers/playground_conversation_router.py

@@ -272,6 +272,8 @@ 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)
+
         if input.stream:
             return StreamingResponse(
                 service.send_message(db,

api/src/conversation/services/conversation_service.py

@@ -1,15 +1,19 @@
 from uuid import UUID
 from typing import List
+from datetime import timedelta
 from sqlalchemy.orm import Session
+from fastapi import HTTPException, status
 from langchain_core.messages.human import HumanMessage
 
-from src.db.models import Conversation
 from src.util.configuration import Config
 from src.agentflow.agents.hub import AgentHub
-from src.agentflow.utils.tools_index import get_all_tools
+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
+from src.query_usage.services.query_usage_service import QueryUsageService
 from src.conversation.repositories.conversation_repository import ConversationRepository
 
 
@@ -22,6 +26,8 @@ def __init__(self):
         self.repository = ConversationRepository()
         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
@@ -132,6 +138,56 @@ 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,
@@ -163,6 +219,7 @@ async def send_message(
         Raises:
             Exception: If associated project is not found
         """
+
         config = Config.get_config()
 
         sub_agents = AgentHub().get_agents(config["agents"])
@@ -173,9 +230,14 @@ async def send_message(
             sub_agents=sub_agents, memory=memory, streaming=stream)
 
         if stream:
-            # Stream tokens as they become available.
-            async for token in agent.stream(message):
-                yield token
+            try:
+                # Stream tokens as they become available.
+                async for token in agent.stream(message):
+                    yield token
+            except Exception as e:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)
+                )
         else:
             result = agent.ask(message)
 
@@ -192,6 +254,13 @@ async def send_message(
                 "intermediate_steps": intermediate_steps
             }
 
+        query_usage = QueryUsage(
+            user_id=user_id,
+            provider="morpheus",
+        )
+        self.query_usage_service.create_query_usage(
+            db_session, query_usage)
+
     def get_history(
         self,
         db_session: Session,