Implement Issue Question Answering Function and Ready for Production

prometheus/app/api/main.py

@@ -1,6 +1,6 @@
 from fastapi import APIRouter
 
-from prometheus.app.api.routes import auth, issue, repository
+from prometheus.app.api.routes import auth, invitation_code, issue, repository, user
 from prometheus.configuration.config import settings
 
 api_router = APIRouter()
@@ -9,3 +9,7 @@
 
 if settings.ENABLE_AUTHENTICATION:
     api_router.include_router(auth.router, prefix="/auth", tags=["auth"])
+    api_router.include_router(
+        invitation_code.router, prefix="/invitation-code", tags=["invitation_code"]
+    )
+    api_router.include_router(user.router, prefix="/user", tags=["user"])

prometheus/app/api/routes/auth.py

@@ -1,9 +1,12 @@
 from fastapi import APIRouter, Request
 
-from prometheus.app.models.requests.auth import LoginRequest
+from prometheus.app.models.requests.auth import CreateUserRequest, LoginRequest
 from prometheus.app.models.response.auth import LoginResponse
 from prometheus.app.models.response.response import Response
+from prometheus.app.services.invitation_code_service import InvitationCodeService
 from prometheus.app.services.user_service import UserService
+from prometheus.configuration.config import settings
+from prometheus.exceptions.server_exception import ServerException
 
 router = APIRouter()
 
@@ -28,3 +31,38 @@ def login(login_request: LoginRequest, request: Request) -> Response[LoginRespon
         password=login_request.password,
     )
     return Response(data=LoginResponse(access_token=access_token))
+
+
+@router.post(
+    "/register/",
+    summary="Register a new user",
+    description="Register a new user with username, email, password and invitation code.",
+    response_description="Returns a success message upon successful registration",
+    response_model=Response,
+)
+def register(request: Request, create_user_request: CreateUserRequest) -> Response:
+    """
+    Register a new user with username, email, password and invitation code.
+    Returns a success message upon successful registration.
+    """
+    invitation_code_service: InvitationCodeService = request.app.state.service[
+        "invitation_code_service"
+    ]
+    user_service: UserService = request.app.state.service["user_service"]
+
+    # Check if the invitation code is valid
+    if not invitation_code_service.check_invitation_code(create_user_request.invitation_code):
+        raise ServerException(code=400, message="Invalid or expired invitation code")
+
+    # Create the user
+    user_service.create_user(
+        username=create_user_request.username,
+        email=create_user_request.email,
+        password=create_user_request.password,
+        issue_credit=settings.DEFAULT_USER_ISSUE_CREDIT,
+    )
+
+    # Mark the invitation code as used
+    invitation_code_service.mark_code_as_used(create_user_request.invitation_code)
+
+    return Response(message="User registered successfully")

prometheus/app/api/routes/invitation_code.py

@@ -0,0 +1,57 @@
+from typing import Sequence
+
+from fastapi import APIRouter, Request
+
+from prometheus.app.decorators.require_login import requireLogin
+from prometheus.app.entity.invitation_code import InvitationCode
+from prometheus.app.models.response.response import Response
+from prometheus.app.services.user_service import UserService
+from prometheus.exceptions.server_exception import ServerException
+
+router = APIRouter()
+
+
+@router.post(
+    "/create/",
+    summary="Create a new invitation code",
+    description="Generates a new invitation code for user registration.",
+    response_description="Returns the newly created invitation code",
+    response_model=Response[InvitationCode],
+)
+@requireLogin
+def create_invitation_code(request: Request) -> Response[InvitationCode]:
+    """
+    Create a new invitation code.
+    """
+    # Check if the user is an admin
+    user_service: UserService = request.app.state.service["user_service"]
+    if not user_service.is_admin(request.state.user_id):
+        raise ServerException(code=403, message="Only admins can create invitation codes")
+
+    # Create a new invitation code
+    invitation_code_service = request.app.state.service["invitation_code_service"]
+    invitation_code = invitation_code_service.create_invitation_code()
+    return Response(data=invitation_code)
+
+
+@router.get(
+    "/list/",
+    summary="List all invitation codes",
+    description="Retrieves a list of all invitation codes.",
+    response_description="Returns a list of invitation codes",
+    response_model=Response[Sequence[InvitationCode]],
+)
+@requireLogin
+def list_invitation_codes(request: Request) -> Response[Sequence[InvitationCode]]:
+    """
+    List all invitation codes.
+    """
+    # Check if the user is an admin
+    user_service: UserService = request.app.state.service["user_service"]
+    if not user_service.is_admin(request.state.user_id):
+        raise ServerException(code=403, message="Only admins can list invitation codes")
+
+    # List all invitation codes
+    invitation_code_service = request.app.state.service["invitation_code_service"]
+    invitation_codes = invitation_code_service.list_invitation_codes()
+    return Response(data=invitation_codes)

prometheus/app/api/routes/issue.py

@@ -26,7 +26,14 @@
 )
 @requireLogin
 async def answer_issue(issue: IssueRequest, request: Request) -> Response[IssueResponse]:
+    # Retrieve necessary services from the application state
     repository_service: RepositoryService = request.app.state.service["repository_service"]
+    user_service: UserService = request.app.state.service["user_service"]
+    issue_service: IssueService = request.app.state.service["issue_service"]
+    knowledge_graph_service: KnowledgeGraphService = request.app.state.service[
+        "knowledge_graph_service"
+    ]
+
     # Fetch the repository by ID
     repository = repository_service.get_repository_by_id(issue.repository_id)
     # Ensure the repository exists
@@ -36,17 +43,15 @@ async def answer_issue(issue: IssueRequest, request: Request) -> Response[IssueR
     if settings.ENABLE_AUTHENTICATION and repository.user_id != request.state.user_id:
         raise ServerException(code=403, message="You do not have access to this repository")
 
-    # Deduct issue credit if authentication is enabled
-    user_service: UserService = request.app.state.service["user_service"]
+    # Check issue credit
+    user_issue_credit = None
     if settings.ENABLE_AUTHENTICATION:
-        # Check and deduct issue credit
         user_issue_credit = user_service.get_issue_credit(request.state.user_id)
         if user_issue_credit <= 0:
             raise ServerException(
                 code=403,
                 message="Insufficient issue credits. Please purchase more to continue.",
             )
-        user_service.update_issue_credit(request.state.user_id, user_issue_credit - 1)
 
     # Validate Dockerfile and workdir inputs
     if issue.dockerfile_content or issue.image_name:
@@ -62,10 +67,7 @@ async def answer_issue(issue: IssueRequest, request: Request) -> Response[IssueR
             message="The repository is currently being used. Please try again later.",
         )
 
-    knowledge_graph_service: KnowledgeGraphService = request.app.state.service[
-        "knowledge_graph_service"
-    ]
-
+    # Load the git repository and knowledge graph
     git_repository = repository_service.get_repository(repository.playground_path)
     knowledge_graph = knowledge_graph_service.get_knowledge_graph(
         repository.kg_root_node_id,
@@ -74,8 +76,7 @@ async def answer_issue(issue: IssueRequest, request: Request) -> Response[IssueR
         repository.kg_chunk_overlap,
     )
 
-    issue_service: IssueService = request.app.state.service["issue_service"]
-
+    # Process the issue in a separate thread to avoid blocking the event loop
     (
         patch,
         passed_reproducing_test,
@@ -104,6 +105,8 @@ async def answer_issue(issue: IssueRequest, request: Request) -> Response[IssueR
         build_commands=issue.build_commands,
         test_commands=issue.test_commands,
     )
+
+    # Check if all outputs are in their initial state, indicating a failure
     if (
         patch,
         passed_reproducing_test,
@@ -117,6 +120,12 @@ async def answer_issue(issue: IssueRequest, request: Request) -> Response[IssueR
             code=500,
             message="Failed to process the issue. Please try again later.",
         )
+
+    # Deduct issue credit after successful processing
+    if settings.ENABLE_AUTHENTICATION:
+        user_service.update_issue_credit(request.state.user_id, user_issue_credit - 1)
+
+    # Return the response
     return Response(
         data=IssueResponse(
             patch=patch,

prometheus/app/api/routes/repository.py

@@ -57,35 +57,48 @@ async def upload_github_repository(
 ):
     # Get the repository and knowledge graph services
     repository_service: RepositoryService = request.app.state.service["repository_service"]
-    repository = repository_service.get_repository_by_url_and_commit_id(
-        upload_repository_request.https_url, commit_id=upload_repository_request.commit_id
-    )
+    knowledge_graph_service: KnowledgeGraphService = request.app.state.service[
+        "knowledge_graph_service"
+    ]
+
+    # Check if the repository already exists
     if settings.ENABLE_AUTHENTICATION:
-        if repository and request.state.user_id == repository.user_id:
-            return Response(
-                message="Repository already exists", data={"repository_id": repository.id}
-            )
+        repository = repository_service.get_repository_by_url_commit_id_and_user_id(
+            upload_repository_request.https_url,
+            upload_repository_request.commit_id,
+            request.state.user_id,
+        )
     else:
-        if repository:
-            # If the repository already exists, return its ID
-            return Response(
-                message="Repository already exists", data={"repository_id": repository.id}
+        repository = repository_service.get_repository_by_url_and_commit_id(
+            upload_repository_request.https_url, commit_id=upload_repository_request.commit_id
+        )
+
+    # If the repository already exists, return its ID
+    if repository:
+        return Response(message="Repository already exists", data={"repository_id": repository.id})
+
+    # Check if the number of repositories exceeds the limit
+    if settings.ENABLE_AUTHENTICATION:
+        user_repositories = repository_service.get_repositories_by_user_id(request.state.user_id)
+        if len(user_repositories) >= settings.DEFAULT_USER_REPOSITORY_LIMIT:
+            raise ServerException(
+                code=400,
+                message=f"You have reached the maximum number of repositories ({settings.DEFAULT_USER_REPOSITORY_LIMIT}). Please delete some repositories before uploading new ones.",
             )
 
-    knowledge_graph_service: KnowledgeGraphService = request.app.state.service[
-        "knowledge_graph_service"
-    ]
+    # Get the GitHub token
     github_token = get_github_token(request, upload_repository_request.github_token)
 
+    # Clone the repository
     try:
-        # Clone the repository
         saved_path = await repository_service.clone_github_repo(
             github_token, upload_repository_request.https_url, upload_repository_request.commit_id
         )
     except git.exc.GitCommandError:
         raise ServerException(
             code=400, message=f"Unable to clone {upload_repository_request.https_url}."
         )
+
     # Build and save the knowledge graph from the cloned repository
     root_node_id = await knowledge_graph_service.build_and_save_knowledge_graph(saved_path)
     repository_id = repository_service.create_new_repository(

prometheus/app/api/routes/user.py

@@ -0,0 +1,54 @@
+from typing import Sequence
+
+from fastapi import APIRouter, Request
+
+from prometheus.app.decorators.require_login import requireLogin
+from prometheus.app.entity.user import User
+from prometheus.app.models.requests.user import SetGithubTokenRequest
+from prometheus.app.models.response.response import Response
+from prometheus.app.models.response.user import UserResponse
+from prometheus.app.services.user_service import UserService
+from prometheus.exceptions.server_exception import ServerException
+
+router = APIRouter()
+
+
+@router.get(
+    "/list/",
+    summary="List all users in the database",
+    description="Retrieves a list of all users.",
+    response_description="Returns a list of users",
+    response_model=Response[Sequence[UserResponse]],
+)
+@requireLogin
+def list_users(request: Request) -> Response[Sequence[User]]:
+    """
+    List all users in the database.
+    """
+    # Check if the user is an admin
+    user_service: UserService = request.app.state.service["user_service"]
+    if not user_service.is_admin(request.state.user_id):
+        raise ServerException(code=403, message="Only admins can list users")
+
+    # List all users
+    users = user_service.list_users()
+    return Response(data=[UserResponse.model_validate(user) for user in users])
+
+
+@router.put(
+    "/set-github-token/",
+    summary="Set GitHub token for the user",
+    description="Sets the GitHub token for the authenticated user.",
+    response_description="Returns the updated user information",
+    response_model=Response,
+)
+@requireLogin
+def set_github_token(request: Request, set_github_token_request: SetGithubTokenRequest) -> Response:
+    """
+    Set GitHub token for the user.
+    """
+    user_service: UserService = request.app.state.service["user_service"]
+
+    # Update the user's GitHub token
+    user_service.set_github_token(request.state.user_id, set_github_token_request.github_token)
+    return Response()

prometheus/app/dependencies.py

@@ -2,6 +2,7 @@
 
 from prometheus.app.services.base_service import BaseService
 from prometheus.app.services.database_service import DatabaseService
+from prometheus.app.services.invitation_code_service import InvitationCodeService
 from prometheus.app.services.issue_service import IssueService
 from prometheus.app.services.knowledge_graph_service import KnowledgeGraphService
 from prometheus.app.services.llm_service import LLMService
@@ -69,6 +70,7 @@ def initialize_services() -> dict[str, BaseService]:
     )
 
     user_service = UserService(database_service)
+    invitation_code_service = InvitationCodeService(database_service)
 
     return {
         "neo4j_service": neo4j_service,
@@ -78,4 +80,5 @@ def initialize_services() -> dict[str, BaseService]:
         "issue_service": issue_service,
         "database_service": database_service,
         "user_service": user_service,
+        "invitation_code_service": invitation_code_service,
     }

prometheus/app/entity/invitation_code.py

@@ -0,0 +1,19 @@
+from datetime import datetime, timedelta, timezone
+
+from sqlmodel import Field, SQLModel
+
+from prometheus.configuration.config import settings
+
+
+class InvitationCode(SQLModel, table=True):
+    """
+    InvitationCode model for managing invitation codes.
+    """
+
+    id: int = Field(primary_key=True, description="ID")
+    code: str = Field(index=True, unique=True, max_length=36, description="Invitation code")
+    is_used: bool = Field(default=False, description="Whether the invitation code has been used")
+    expiration_time: datetime = Field(
+        default=datetime.now(timezone.utc) + timedelta(days=settings.INVITATION_CODE_EXPIRE_TIME),
+        description="Expiration time of the invitation code",
+    )