diff --git a/prometheus/app/api/routes/issue.py b/prometheus/app/api/routes/issue.py index e58582c6..05253252 100644 --- a/prometheus/app/api/routes/issue.py +++ b/prometheus/app/api/routes/issue.py @@ -7,6 +7,7 @@ from prometheus.app.services.issue_service import IssueService from prometheus.app.services.knowledge_graph_service import KnowledgeGraphService from prometheus.app.services.repository_service import RepositoryService +from prometheus.app.services.user_service import UserService from prometheus.configuration.config import settings from prometheus.exceptions.server_exception import ServerException @@ -24,19 +25,35 @@ @requireLogin async def answer_issue(issue: IssueRequest, request: Request) -> Response[IssueResponse]: repository_service: RepositoryService = request.app.state.service["repository_service"] + # Fetch the repository by ID repository = repository_service.get_repository_by_id(issue.repository_id) + # Ensure the repository exists if not repository: raise ServerException(code=404, message="Repository not found") + # Ensure the user has access to the repository 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"] + 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: if issue.workdir is None: raise ServerException( code=400, message="workdir must be provided for user defined environment", ) - + # Ensure the repository is not currently being used if repository.is_working: raise ServerException( code=400, diff --git a/prometheus/app/api/routes/repository.py b/prometheus/app/api/routes/repository.py index 7aff519f..f78d44f4 100644 --- a/prometheus/app/api/routes/repository.py +++ b/prometheus/app/api/routes/repository.py @@ -23,7 +23,7 @@ def get_github_token(request: Request, github_token: str) -> str: if not settings.ENABLE_AUTHENTICATION: # If the user is not authenticated, raise an exception raise ServerException( - code=401, message="GitHub token is required, please provide it or log in" + code=400, message="GitHub token is required, please provide it or log in" ) # If the user is authenticated, get the user service and fetch the token user_service: UserService = request.app.state.service["user_service"] @@ -33,7 +33,7 @@ def get_github_token(request: Request, github_token: str) -> str: # If the token is still not available, raise an exception if not github_token: raise ServerException( - code=423, message="Either provide a GitHub token or set it in your user profile" + code=400, message="Either provide a GitHub token or set it in your user profile" ) return github_token diff --git a/prometheus/app/services/user_service.py b/prometheus/app/services/user_service.py index 49fa6030..92d5c628 100644 --- a/prometheus/app/services/user_service.py +++ b/prometheus/app/services/user_service.py @@ -120,3 +120,33 @@ def get_user_by_id(self, user_id: int) -> Optional[User]: with Session(self.engine) as session: statement = select(User).where(User.id == user_id) return session.exec(statement).first() + + def get_issue_credit(self, user_id: int) -> int: + """ + Retrieve the issue credit of a user by their ID. + + Args: + user_id (int): The ID of the user. + + Returns: + int: The issue credit of the user. + """ + with Session(self.engine) as session: + statement = select(User.issue_credit).where(User.id == user_id) + result = session.exec(statement).first() + return result + + def update_issue_credit(self, user_id: int, new_issue_credit) -> None: + """ + Update the issue credit of a user by their ID. + + Args: + user_id (int): The ID of the user. + new_issue_credit (int): The new issue credit. + """ + with Session(self.engine) as session: + statement = select(User).where(User.id == user_id) + user = session.exec(statement).first() + user.issue_credit = new_issue_credit + session.add(user) + session.commit()