diff --git a/.github/workflows/pytest_and_coverage.yml b/.github/workflows/pytest_and_coverage.yml index 15aa9547..c80effdc 100644 --- a/.github/workflows/pytest_and_coverage.yml +++ b/.github/workflows/pytest_and_coverage.yml @@ -14,6 +14,10 @@ jobs: # Logging PROMETHEUS_LOGGING_LEVEL: DEBUG + # General settings + PROMETHEUS_ENVIRONMENT: local + PROMETHEUS_BACKEND_CORS_ORIGINS: "[\"*\"]" + # Neo4j settings PROMETHEUS_NEO4J_URI: bolt://localhost:7687 PROMETHEUS_NEO4J_USERNAME: neo4j @@ -48,6 +52,9 @@ jobs: # DATABASE settings PROMETHEUS_DATABASE_URL: postgresql://postgres:password@localhost:5432/postgres?sslmode=disable + # JWT settings + PROMETHEUS_JWT_SECRET_KEY: your_jwt_secret_key + steps: - name: Check out code uses: actions/checkout@v4 diff --git a/README.md b/README.md index 0044bed3..780984ac 100644 --- a/README.md +++ b/README.md @@ -73,10 +73,20 @@ governed by a state machine to ensure code quality through automated reviews, bu ``` 2. #### Copy the `example.env` file to `.env` and update it with your API keys and other required configurations: + ```bash mv example.env .env ``` + > You need to provide a secure `JWT_SECRET_KEY` in the `.env` file. + > You can generate a strong key by running the following command: + + ```bash + python -m prometheus.script.generate_jwt_token + ``` + + This will print a secure token you can copy and paste into your `.env` file + 3. #### Create the working directory to store logs and cloned repositories: ```bash @@ -124,7 +134,7 @@ governed by a state machine to ensure code quality through automated reviews, bu You can ask Prometheus to analyze and answer a specific issue in your codebase using the `/issue/answer/` API endpoint. - **Endpoint:** `POST /issue/answer/` - - **Request Body:** JSON object matching the `IssueRequest` schema (see [API Documents](http://localhost:9002/docs#/issue/answer_issue_issue_answer__post)) + - **Request Body:** JSON object matching the `IssueRequest` schema (see [API Documents](http://127.0.0.1:9002/docs#/issue/issue-answer_issue)) - **Response:** Returns the generated patch, test/build results, and a summary response. --- diff --git a/docker-compose.win_mac.yml b/docker-compose.win_mac.yml index 8eeec133..bdeb6cd2 100644 --- a/docker-compose.win_mac.yml +++ b/docker-compose.win_mac.yml @@ -36,6 +36,10 @@ services: # Logging - PROMETHEUS_LOGGING_LEVEL=${PROMETHEUS_LOGGING_LEVEL} + # General settings + - PROMETHEUS_ENVIRONMENT=${PROMETHEUS_ENVIRONMENT} + - PROMETHEUS_BACKEND_CORS_ORIGINS=${PROMETHEUS_BACKEND_CORS_ORIGINS} + # Neo4j settings - PROMETHEUS_NEO4J_URI=${PROMETHEUS_NEO4J_URI} - PROMETHEUS_NEO4J_USERNAME=${PROMETHEUS_NEO4J_USERNAME} @@ -64,9 +68,6 @@ services: - PROMETHEUS_MAX_OUTPUT_TOKENS=${PROMETHEUS_MAX_OUTPUT_TOKENS} - PROMETHEUS_TEMPERATURE=${PROMETHEUS_TEMPERATURE} - # GitHub settings - - PROMETHEUS_GITHUB_ACCESS_TOKEN=${PROMETHEUS_GITHUB_ACCESS_TOKEN} - # Database settings - PROMETHEUS_DATABASE_URL=${PROMETHEUS_DATABASE_URL} networks: diff --git a/docker-compose.yml b/docker-compose.yml index 03fb66f3..72f3a56b 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -57,6 +57,10 @@ services: # Logging - PROMETHEUS_LOGGING_LEVEL=${PROMETHEUS_LOGGING_LEVEL} + # General settings + - PROMETHEUS_ENVIRONMENT=${PROMETHEUS_ENVIRONMENT} + - PROMETHEUS_BACKEND_CORS_ORIGINS=${PROMETHEUS_BACKEND_CORS_ORIGINS} + # Neo4j settings - PROMETHEUS_NEO4J_URI=${PROMETHEUS_NEO4J_URI} - PROMETHEUS_NEO4J_USERNAME=${PROMETHEUS_NEO4J_USERNAME} @@ -85,9 +89,6 @@ services: - PROMETHEUS_MAX_OUTPUT_TOKENS=${PROMETHEUS_MAX_OUTPUT_TOKENS} - PROMETHEUS_TEMPERATURE=${PROMETHEUS_TEMPERATURE} - # GitHub settings - - PROMETHEUS_GITHUB_ACCESS_TOKEN=${PROMETHEUS_GITHUB_ACCESS_TOKEN} - # Database settings - PROMETHEUS_DATABASE_URL=${PROMETHEUS_DATABASE_URL} volumes: diff --git a/example.env b/example.env index d2b2c063..8699b7d4 100644 --- a/example.env +++ b/example.env @@ -1,6 +1,10 @@ # Logging PROMETHEUS_LOGGING_LEVEL=DEBUG +# General settings +PROMETHEUS_ENVIRONMENT=local +PROMETHEUS_BACKEND_CORS_ORIGINS=["*] + # Neo4j settings PROMETHEUS_NEO4J_URI=bolt://neo4j:7687 PROMETHEUS_NEO4J_USERNAME=neo4j @@ -29,8 +33,8 @@ PROMETHEUS_MAX_INPUT_TOKENS=64000 PROMETHEUS_TEMPERATURE=0.3 PROMETHEUS_MAX_OUTPUT_TOKENS=15000 -# GitHub settings -PROMETHEUS_GITHUB_ACCESS_TOKEN=github_access_token - # Database settings PROMETHEUS_DATABASE_URL=postgresql://postgres:password@localhost:5432/postgres?sslmode=disable + +# JWT settings +PROMETHEUS_JWT_SECRET_KEY=your_jwt_secret_key diff --git a/prometheus/app/api/auth.py b/prometheus/app/api/auth.py new file mode 100644 index 00000000..cb55a352 --- /dev/null +++ b/prometheus/app/api/auth.py @@ -0,0 +1,30 @@ +from fastapi import APIRouter, Request + +from prometheus.app.models.requests.auth import LoginRequest +from prometheus.app.models.response.auth import LoginResponse +from prometheus.app.models.response.response import Response +from prometheus.app.services.user_service import UserService + +router = APIRouter() + + +@router.post( + "/login/", + summary="Login to the system", + description="Login to the system using username, email, and password. Returns an access token.", + response_description="Returns an access token for authenticated requests", + response_model=Response[LoginResponse], +) +def login(login_request: LoginRequest, request: Request) -> Response[LoginResponse]: + """ + Login to the system using username, email, and password. + Returns an access token for authenticated requests. + """ + + user_service: UserService = request.app.state.service["user_service"] + access_token = user_service.login( + username=login_request.username, + email=login_request.email, + password=login_request.password, + ) + return Response(data=LoginResponse(access_token=access_token)) diff --git a/prometheus/app/api/issue.py b/prometheus/app/api/issue.py index 19963fb9..803ebb90 100644 --- a/prometheus/app/api/issue.py +++ b/prometheus/app/api/issue.py @@ -1,105 +1,33 @@ -from typing import Mapping, Optional, Sequence +from fastapi import APIRouter, Request -from fastapi import APIRouter, HTTPException, Request -from litellm import Field -from pydantic import BaseModel - -from prometheus.lang_graph.graphs.issue_state import IssueType +from prometheus.app.models.requests.issue import IssueRequest +from prometheus.app.models.response.issue import IssueResponse +from prometheus.app.models.response.response import Response +from prometheus.exceptions.server_exception import ServerException router = APIRouter() -class IssueRequest(BaseModel): - issue_number: int = Field(description="The number of the issue", examples=[42]) - issue_title: str = Field( - description="The title of the issue", examples=["There is a memory leak"] - ) - issue_body: str = Field( - description="The description of the issue", examples=["foo/bar.c is causing a memory leak"] - ) - issue_comments: Optional[Sequence[Mapping[str, str]]] = Field( - default=None, - description="Comments on the issue", - examples=[ - [ - {"username": "user1", "comment": "I've experienced this issue as well."}, - { - "username": "user2", - "comment": "A potential fix is to adjust the memory settings.", - }, - ] - ], - ) - issue_type: IssueType = Field( - default=IssueType.AUTO, - description="The type of the issue, set to auto if you do not know", - examples=[IssueType.AUTO], - ) - run_build: Optional[bool] = Field( - default=False, - description="When editing the code, whenver we should run the build to verify the fix", - examples=[False], - ) - run_existing_test: Optional[bool] = Field( - default=False, - description="When editing the code, whenver we should run the existing test to verify the fix", - examples=[False], - ) - number_of_candidate_patch: Optional[int] = Field( - default=4, - description="When the patch is not verfied (through build or test), number of candidate patches we generate to select the best one", - examples=[4], - ) - dockerfile_content: Optional[str] = Field( - default=None, - description="Specify the containerized environment with dockerfile content", - examples=["FROM python:3.11\nWORKDIR /app\nCOPY . /app"], - ) - image_name: Optional[str] = Field( - default=None, - description="Specify the containerized environment with image name that should be pulled from dockerhub", - examples=["python:3.11-slim"], - ) - workdir: Optional[str] = Field( - default=None, - description="If you specified the container environment, you must also specify the workdir", - examples=["/app"], - ) - build_commands: Optional[Sequence[str]] = Field( - default=None, - description="If you specified dockerfile_content and run_build is True, you must also specify the build commands.", - examples=[["pip install -r requirements.txt", "python -m build"]], - ) - test_commands: Optional[Sequence[str]] = Field( - default=None, - description="If you specified dockerfile_content and run_test is True, you must also specify the test commands.", - examples=[["pytest ."]], - ) - push_to_remote: Optional[bool] = Field( - default=False, - description="When editing the code, whenver we should push the changes to a remote branch", - examples=[True], - ) - - @router.post( "/answer/", summary="Process and generate a response for an issue", - description="Analyzes an issue, generates patches if needed, runs optional builds and tests, and can push changes to a remote branch.", + description="Analyzes an issue, generates patches if needed, runs optional builds and tests, and can push changes " + "to a remote branch.", response_description="Returns the patch, test results, and issue response", + response_model=Response[IssueResponse], ) -def answer_issue(issue: IssueRequest, request: Request): - if not request.app.state.service_coordinator.exists_knowledge_graph(): - raise HTTPException( - status_code=404, - detail="A repository is not uploaded, use /repository/ endpoint to upload one", +def answer_issue(issue: IssueRequest, request: Request) -> Response[IssueResponse]: + if not request.app.state.service["knowledge_graph_service"].exists(): + raise ServerException( + code=404, + message="A repository is not uploaded, use /repository/ endpoint to upload one", ) if issue.dockerfile_content or issue.image_name: if issue.workdir is None: - raise HTTPException( - status_code=400, - detail="workdir must be provided for user defined environment", + raise ServerException( + code=400, + message="workdir must be provided for user defined environment", ) ( @@ -109,7 +37,7 @@ def answer_issue(issue: IssueRequest, request: Request): passed_build, passed_existing_test, issue_response, - ) = request.app.state.service_coordinator.answer_issue( + ) = request.app.state.service["issue_service"].answer_issue( issue_number=issue.issue_number, issue_title=issue.issue_title, issue_body=issue.issue_body, @@ -125,11 +53,13 @@ def answer_issue(issue: IssueRequest, request: Request): test_commands=issue.test_commands, push_to_remote=issue.push_to_remote, ) - return { - "patch": patch, - "passed_reproducing_test": passed_reproducing_test, - "passed_build": passed_build, - "passed_existing_test": passed_existing_test, - "issue_response": issue_response, - "remote_branch_name": remote_branch_name, - } + return Response( + data=IssueResponse( + patch=patch, + passed_reproducing_test=passed_reproducing_test, + passed_build=passed_build, + passed_existing_test=passed_existing_test, + issue_response=issue_response, + remote_branch_name=remote_branch_name, + ) + ) diff --git a/prometheus/app/api/repository.py b/prometheus/app/api/repository.py index 65655cf4..49784a32 100644 --- a/prometheus/app/api/repository.py +++ b/prometheus/app/api/repository.py @@ -1,33 +1,12 @@ -from pathlib import Path - import git -from fastapi import APIRouter, HTTPException, Request - -router = APIRouter() - - -@router.get( - "/local/", - description=""" - Upload a local codebase to Prometheus. - """, - responses={ - 404: {"description": "Local repository not found"}, - 200: {"description": "Repository uploaded successfully"}, - }, -) -def upload_local_repository(local_repository: str, request: Request): - local_path = Path(local_repository) +from fastapi import APIRouter, Request - if not local_path.exists(): - raise HTTPException( - status_code=404, - detail=f"Local repository not found at path: {local_repository}", - ) +from prometheus.app.models.response.response import Response +from prometheus.app.services.knowledge_graph_service import KnowledgeGraphService +from prometheus.app.services.repository_service import RepositoryService +from prometheus.exceptions.server_exception import ServerException - request.app.state.service_coordinator.upload_local_repository(local_path) - - return {"message": "Repository uploaded successfully"} +router = APIRouter() @router.get( @@ -35,12 +14,26 @@ def upload_local_repository(local_repository: str, request: Request): description=""" Upload a GitHub repository to Prometheus, default to the latest commit in the main branch. """, + response_model=Response, ) -def upload_github_repository(https_url: str, request: Request): +def upload_github_repository(github_token: str, https_url: str, request: Request): + # Get the repository and knowledge graph services + repository_service: RepositoryService = request.app.state.service["repository_service"] + knowledge_graph_service: KnowledgeGraphService = request.app.state.service[ + "knowledge_graph_service" + ] + + # Clean the services to ensure no previous data is present + repository_service.clean() + knowledge_graph_service.clear() try: - request.app.state.service_coordinator.upload_github_repository(https_url) + # Clone the repository + saved_path = repository_service.clone_github_repo(github_token, https_url) except git.exc.GitCommandError: - raise HTTPException(status_code=400, detail=f"Unable to clone {https_url}") + raise ServerException(code=400, message=f"Unable to clone {https_url}") + # Build and save the knowledge graph from the cloned repository + knowledge_graph_service.build_and_save_knowledge_graph(saved_path, https_url) + return Response() @router.get( @@ -48,14 +41,29 @@ def upload_github_repository(https_url: str, request: Request): description=""" Upload a GitHub repository at a specific commit to Prometheus. """, + response_model=Response, ) -def upload_github_repository_at_commit(https_url: str, commit_id: str, request: Request): +def upload_github_repository_at_commit( + github_token, https_url: str, commit_id: str, request: Request +): + # Get the repository and knowledge graph services + repository_service: RepositoryService = request.app.state.service["repository_service"] + knowledge_graph_service: KnowledgeGraphService = request.app.state.service[ + "knowledge_graph_service" + ] + + # Clean the services to ensure no previous data is present + repository_service.clean() + knowledge_graph_service.clear() + try: - request.app.state.service_coordinator.upload_github_repository(https_url, commit_id) + # Clone the repository + saved_path = repository_service.clone_github_repo(github_token, https_url, commit_id) except git.exc.GitCommandError: - raise HTTPException( - status_code=400, detail=f"Unable to clone {https_url} with commit {commit_id}" - ) + raise ServerException(code=400, message=f"Unable to clone {https_url}") + # Build and save the knowledge graph from the cloned repository + knowledge_graph_service.build_and_save_knowledge_graph(saved_path, https_url, commit_id) + return Response() @router.get( @@ -63,13 +71,21 @@ def upload_github_repository_at_commit(https_url: str, commit_id: str, request: description=""" Delete the repository uploaded to Prometheus, along with other information. """, + response_model=Response, ) def delete(request: Request): - if not request.app.state.service_coordinator.exists_knowledge_graph(): - return {"message": "No knowledge graph to delete"} + knowledge_graph_service: KnowledgeGraphService = request.app.state.service[ + "knowledge_graph_service" + ] + if not knowledge_graph_service.exists(): + return Response(message="No knowledge graph to delete") + # Get the repository service to clean up the repository data + repository_service: RepositoryService = request.app.state.service["repository_service"] - request.app.state.service_coordinator.clear() - return {"message": "Successfully deleted knowledge graph"} + # Clear the knowledge graph and repository data + knowledge_graph_service.clear() + repository_service.clean() + return Response() @router.get( @@ -77,7 +93,7 @@ def delete(request: Request): description=""" If there is a codebase uploaded to Promtheus. """, - response_model=bool, + response_model=Response[bool], ) -def knowledge_graph_exists(request: Request) -> bool: - return request.app.state.service_coordinator.exists_knowledge_graph() +def knowledge_graph_exists(request: Request) -> Response[bool]: + return Response(data=request.app.state.service["knowledge_graph_service"].exists()) diff --git a/prometheus/app/db.py b/prometheus/app/db.py deleted file mode 100644 index a95fa8f1..00000000 --- a/prometheus/app/db.py +++ /dev/null @@ -1,57 +0,0 @@ -import logging -from typing import Optional - -from passlib.hash import bcrypt -from sqlmodel import Session, SQLModel, create_engine - -from prometheus.app.entity.user import User -from prometheus.configuration.config import settings - -engine = create_engine(settings.DATABASE_URL, echo=True) -_logger = logging.getLogger("prometheus.app.db") - - -# Create the database and tables -def create_db_and_tables(): - SQLModel.metadata.create_all(engine) - - -# Create a superuser and commit it to the database -def create_superuser( - username: str, - email: str, - password: str, - github_token: Optional[str] = None, -) -> None: - """ - Create a new superuser and commit it to the database. - - Args: - username (str): Desired username. - email (str): Email address. - password (str): Plaintext password (will be hashed). - github_token (Optional[str]): Optional GitHub token. - - Returns: - User: The created superuser instance. - """ - with Session(engine) as session: - if session.query(User).filter(User.username == username).first(): - raise ValueError(f"Username '{username}' already exists") - if session.query(User).filter(User.email == email).first(): - raise ValueError(f"Email '{email}' already exists") - - hashed_password = bcrypt.hash(password) - - user = User( - username=username, - email=email, - password_hash=hashed_password, - github_token=github_token, - issue_credit=999999, - is_superuser=True, - ) - session.add(user) - session.commit() - session.refresh(user) - _logger.info(f"Superuser '{username}' created successfully.") diff --git a/prometheus/app/dependencies.py b/prometheus/app/dependencies.py index 5b3986d1..c66b7470 100644 --- a/prometheus/app/dependencies.py +++ b/prometheus/app/dependencies.py @@ -1,17 +1,17 @@ """Initializes and configures all prometheus services.""" -from pathlib import Path - +from prometheus.app.services.base_service import BaseService +from prometheus.app.services.database_service import DatabaseService 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 from prometheus.app.services.neo4j_service import Neo4jService from prometheus.app.services.repository_service import RepositoryService -from prometheus.app.services.service_coordinator import ServiceCoordinator +from prometheus.app.services.user_service import UserService from prometheus.configuration.config import settings -def initialize_services() -> ServiceCoordinator: +def initialize_services() -> dict[str, BaseService]: """Initializes and configures the complete prometheus service stack. This function creates and configures all required services for prometheus @@ -58,17 +58,17 @@ def initialize_services() -> ServiceCoordinator: neo4j_service, llm_service, settings.MAX_TOKEN_PER_NEO4J_RESULT, + settings.WORKING_DIRECTORY, ) + database_service = DatabaseService(settings.DATABASE_URL) + user_service = UserService(database_service) - service_coordinator = ServiceCoordinator( - issue_service, - knowledge_graph_service, - llm_service, - neo4j_service, - repository_service, - settings.MAX_TOKEN_PER_NEO4J_RESULT, - settings.GITHUB_ACCESS_TOKEN, - Path(settings.WORKING_DIRECTORY).absolute(), - ) - - return service_coordinator + return { + "neo4j_service": neo4j_service, + "llm_service": llm_service, + "knowledge_graph_service": knowledge_graph_service, + "repository_service": repository_service, + "issue_service": issue_service, + "database_service": database_service, + "user_service": user_service, + } diff --git a/prometheus/app/entity/user.py b/prometheus/app/entity/user.py index 6ed3776c..e6bdf13f 100644 --- a/prometheus/app/entity/user.py +++ b/prometheus/app/entity/user.py @@ -5,15 +5,18 @@ class User(SQLModel, table=True): id: int = Field(primary_key=True, description="User ID") username: str = Field( - index=True, unique=True, max_length=50, description="Username of the user" + index=True, unique=True, max_length=20, description="Username of the user" ) email: str = Field( - index=True, unique=True, max_length=100, description="Email address of the user" + index=True, unique=True, max_length=30, description="Email address of the user" ) password_hash: str = Field(max_length=128, description="Hashed password of the user") github_token: str = Field( - default=None, nullable=True, description="Optional GitHub token for integrations" + default=None, + nullable=True, + description="Optional GitHub token for integrations", + max_length=100, ) issue_credit: int = Field(default=0, ge=0, description="Number of issue credits the user has") diff --git a/prometheus/app/exception_handler.py b/prometheus/app/exception_handler.py new file mode 100644 index 00000000..2a9a56f0 --- /dev/null +++ b/prometheus/app/exception_handler.py @@ -0,0 +1,26 @@ +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +from prometheus.exceptions.server_exception import ServerException + + +def register_exception_handlers(app: FastAPI): + """Global exception handlers for the FastAPI application.""" + + @app.exception_handler(ServerException) + async def custom_exception_handler(_request: Request, exc: ServerException): + """ + Custom exception handler for ServerException. + """ + return JSONResponse( + status_code=exc.code, content={"code": exc.code, "message": exc.message, "data": None} + ) + + @app.exception_handler(Exception) + async def global_exception_handler(_request: Request, _exc: Exception): + """ + Global exception handler for all uncaught exceptions. + """ + return JSONResponse( + status_code=500, content={"code": 500, "message": "Internal Server Error", "data": None} + ) diff --git a/prometheus/app/main.py b/prometheus/app/main.py index 2eeca317..fe7bc600 100644 --- a/prometheus/app/main.py +++ b/prometheus/app/main.py @@ -3,9 +3,11 @@ from datetime import datetime, timezone from fastapi import FastAPI +from fastapi.routing import APIRoute from prometheus.app import dependencies from prometheus.app.api import issue, repository +from prometheus.app.exception_handler import register_exception_handlers from prometheus.configuration.config import settings # Create a logger for the application's namespace @@ -19,6 +21,8 @@ # Log the configuration settings logger.info(f"LOGGING_LEVEL={settings.LOGGING_LEVEL}") +logger.info(f"ENVIRONMENT={settings.ENVIRONMENT}") +logger.info(f"BACKEND_CORS_ORIGINS={settings.BACKEND_CORS_ORIGINS}") logger.info(f"ADVANCED_MODEL={settings.ADVANCED_MODEL}") logger.info(f"BASE_MODEL={settings.BASE_MODEL}") logger.info(f"NEO4J_BATCH_SIZE={settings.NEO4J_BATCH_SIZE}") @@ -35,15 +39,36 @@ @asynccontextmanager async def lifespan(app: FastAPI): # Initialization on startup - app.state.service_coordinator = dependencies.initialize_services() + app.state.service = dependencies.initialize_services() + logger.info("Starting services...") + for service in app.state.service.values(): + service.start() # Initialization Completed yield # Cleanup on shutdown - app.state.service_coordinator.clear() - app.state.service_coordinator.close() + logger.info("Shutting down services...") + for service in app.state.service.values(): + service.close() -app = FastAPI(lifespan=lifespan) +def custom_generate_unique_id(route: APIRoute) -> str: + """ + Custom function to generate unique IDs for API routes based on their tags and names. + """ + return f"{route.tags[0]}-{route.name}" + + +app = FastAPI( + lifespan=lifespan, + title=settings.PROJECT_NAME, # Title on generated documentation + openapi_url=f"{settings.BASE_URL}/openapi.json", # Path to generated OpenAPI documentation + generate_unique_id_function=custom_generate_unique_id, # Custom function for generating unique route IDs + version=settings.version, # Version of the API + debug=True if settings.ENVIRONMENT == "local" else False, +) + +# Register the exception handlers +register_exception_handlers(app) app.include_router(repository.router, prefix="/repository", tags=["repository"]) app.include_router(issue.router, prefix="/issue", tags=["issue"]) diff --git a/prometheus/app/middlewares/__init__.py b/prometheus/app/middlewares/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/prometheus/app/models/__init__.py b/prometheus/app/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/prometheus/app/models/requests/__init__.py b/prometheus/app/models/requests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/prometheus/app/models/requests/auth.py b/prometheus/app/models/requests/auth.py new file mode 100644 index 00000000..f0e8b8fa --- /dev/null +++ b/prometheus/app/models/requests/auth.py @@ -0,0 +1,32 @@ +import re + +from pydantic import BaseModel, Field, field_validator, model_validator + + +class LoginRequest(BaseModel): + username: str = Field(description="username of the user", max_length=20) + email: str = Field( + description="email of the user", + examples=["your_email@gmail.com"], + max_length=30, + ) + password: str = Field( + description="password of the user", + examples=["P@ssw0rd!"], + min_length=12, + max_length=30, + ) + + @classmethod + @field_validator("email", mode="after") + def validate_email_format(cls, v: str) -> str: + pattern = r"^[^@\s]+@[^@\s]+\.[^@\s]+$" + if not re.match(pattern, v): + raise ValueError("Invalid email format") + return v + + @model_validator(mode="after") + def check_username_or_email(self) -> "LoginRequest": + if not self.username and not self.email: + raise ValueError("At least one of 'username' or 'email' must be provided.") + return self diff --git a/prometheus/app/models/requests/issue.py b/prometheus/app/models/requests/issue.py new file mode 100644 index 00000000..f0d1f2b7 --- /dev/null +++ b/prometheus/app/models/requests/issue.py @@ -0,0 +1,78 @@ +from typing import Mapping, Optional, Sequence + +from pydantic import BaseModel, Field + +from prometheus.lang_graph.graphs.issue_state import IssueType + + +class IssueRequest(BaseModel): + issue_number: int = Field(description="The number of the issue", examples=[42]) + issue_title: str = Field( + description="The title of the issue", examples=["There is a memory leak"] + ) + issue_body: str = Field( + description="The description of the issue", examples=["foo/bar.c is causing a memory leak"] + ) + issue_comments: Optional[Sequence[Mapping[str, str]]] = Field( + default=None, + description="Comments on the issue", + examples=[ + [ + {"username": "user1", "comment": "I've experienced this issue as well."}, + { + "username": "user2", + "comment": "A potential fix is to adjust the memory settings.", + }, + ] + ], + ) + issue_type: IssueType = Field( + default=IssueType.AUTO, + description="The type of the issue, set to auto if you do not know", + examples=[IssueType.AUTO], + ) + run_build: Optional[bool] = Field( + default=False, + description="When editing the code, whenver we should run the build to verify the fix", + examples=[False], + ) + run_existing_test: Optional[bool] = Field( + default=False, + description="When editing the code, whenver we should run the existing test to verify the fix", + examples=[False], + ) + number_of_candidate_patch: Optional[int] = Field( + default=4, + description="When the patch is not verfied (through build or test), number of candidate patches we generate to select the best one", + examples=[4], + ) + dockerfile_content: Optional[str] = Field( + default=None, + description="Specify the containerized environment with dockerfile content", + examples=["FROM python:3.11\nWORKDIR /app\nCOPY . /app"], + ) + image_name: Optional[str] = Field( + default=None, + description="Specify the containerized environment with image name that should be pulled from dockerhub", + examples=["python:3.11-slim"], + ) + workdir: Optional[str] = Field( + default=None, + description="If you specified the container environment, you must also specify the workdir", + examples=["/app"], + ) + build_commands: Optional[Sequence[str]] = Field( + default=None, + description="If you specified dockerfile_content and run_build is True, you must also specify the build commands.", + examples=[["pip install -r requirements.txt", "python -m build"]], + ) + test_commands: Optional[Sequence[str]] = Field( + default=None, + description="If you specified dockerfile_content and run_test is True, you must also specify the test commands.", + examples=[["pytest ."]], + ) + push_to_remote: Optional[bool] = Field( + default=False, + description="When editing the code, whenver we should push the changes to a remote branch", + examples=[True], + ) diff --git a/prometheus/app/models/requests/user.py b/prometheus/app/models/requests/user.py new file mode 100644 index 00000000..6a0e3144 --- /dev/null +++ b/prometheus/app/models/requests/user.py @@ -0,0 +1,26 @@ +import re + +from pydantic import BaseModel, Field, field_validator + + +class CreateUserRequest(BaseModel): + username: str = Field(description="username of the user", max_length=20) + email: str = Field( + description="email of the user", + examples=["your_email@gmail.com"], + max_length=30, + ) + password: str = Field( + description="password of the user", + examples=["P@ssw0rd!"], + min_length=12, + max_length=30, + ) + github_token: str = Field(description="github token of the user", max_length=100) + + @field_validator("email", mode="after") + def validate_email_format(self, v: str) -> str: + pattern = r"^[^@\s]+@[^@\s]+\.[^@\s]+$" + if not re.match(pattern, v): + raise ValueError("Invalid email format") + return v diff --git a/prometheus/app/models/response/__init__.py b/prometheus/app/models/response/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/prometheus/app/models/response/auth.py b/prometheus/app/models/response/auth.py new file mode 100644 index 00000000..56efc9f6 --- /dev/null +++ b/prometheus/app/models/response/auth.py @@ -0,0 +1,9 @@ +from pydantic import BaseModel + + +class LoginResponse(BaseModel): + """ + Response model for user login. + """ + + access_token: str diff --git a/prometheus/app/models/response/issue.py b/prometheus/app/models/response/issue.py new file mode 100644 index 00000000..d06f46a7 --- /dev/null +++ b/prometheus/app/models/response/issue.py @@ -0,0 +1,10 @@ +from pydantic import BaseModel + + +class IssueResponse(BaseModel): + patch: str + passed_reproducing_test: bool + passed_build: bool + passed_existing_test: bool + issue_response: str + remote_branch_name: str | None = None diff --git a/prometheus/app/models/response/response.py b/prometheus/app/models/response/response.py new file mode 100644 index 00000000..09c7e595 --- /dev/null +++ b/prometheus/app/models/response/response.py @@ -0,0 +1,15 @@ +from typing import Generic, TypeVar + +from pydantic import BaseModel + +T = TypeVar("T") + + +class Response(BaseModel, Generic[T]): + """ + Generic response model for API responses. + """ + + code: int = 200 + message: str = "success" + data: T | None = None diff --git a/prometheus/app/services/base_service.py b/prometheus/app/services/base_service.py new file mode 100644 index 00000000..4bb9ff0d --- /dev/null +++ b/prometheus/app/services/base_service.py @@ -0,0 +1,18 @@ +class BaseService: + """ + Base class for all services in the Prometheus application. + """ + + def start(self): + """ + Start the service. + This method should be overridden by subclasses to implement specific startup logic. + """ + pass + + def close(self): + """ + Close the service and release any resources. + This method should be overridden by subclasses to implement specific cleanup logic. + """ + pass diff --git a/prometheus/app/services/database_service.py b/prometheus/app/services/database_service.py new file mode 100644 index 00000000..ae4cfc7b --- /dev/null +++ b/prometheus/app/services/database_service.py @@ -0,0 +1,30 @@ +import logging + +from sqlmodel import SQLModel, create_engine + +from prometheus.app.services.base_service import BaseService + + +class DatabaseService(BaseService): + def __init__(self, DATABASE_URL: str): + self.engine = create_engine(DATABASE_URL, echo=True) + self._logger = logging.getLogger("prometheus.app.services.database_service") + + # Create the database and tables + def create_db_and_tables(self): + SQLModel.metadata.create_all(self.engine) + + def start(self): + """ + Start the database service by creating the database and tables. + This method is called when the service is initialized. + """ + self.create_db_and_tables() + self._logger.info("Database and tables created successfully.") + + def close(self): + """ + Close the database connection and release any resources. + """ + self.engine.dispose() + self._logger.info("Database connection closed.") diff --git a/prometheus/app/services/issue_service.py b/prometheus/app/services/issue_service.py index cebf7a87..a66898bc 100644 --- a/prometheus/app/services/issue_service.py +++ b/prometheus/app/services/issue_service.py @@ -1,5 +1,10 @@ +import logging +import traceback +from datetime import datetime +from pathlib import Path from typing import Mapping, Optional, Sequence +from prometheus.app.services.base_service import BaseService from prometheus.app.services.knowledge_graph_service import KnowledgeGraphService from prometheus.app.services.llm_service import LLMService from prometheus.app.services.neo4j_service import Neo4jService @@ -10,7 +15,7 @@ from prometheus.lang_graph.graphs.issue_state import IssueType -class IssueService: +class IssueService(BaseService): def __init__( self, kg_service: KnowledgeGraphService, @@ -18,15 +23,20 @@ def __init__( neo4j_service: Neo4jService, llm_service: LLMService, max_token_per_neo4j_result: int, + working_directory: str, ): self.kg_service = kg_service self.repository_service = repository_service self.neo4j_service = neo4j_service self.llm_service = llm_service self.max_token_per_neo4j_result = max_token_per_neo4j_result + self.working_directory = working_directory + self.answer_issue_log_dir = Path(self.working_directory) / "answer_issue_logs" + self.answer_issue_log_dir.mkdir(parents=True, exist_ok=True) def answer_issue( self, + issue_number: int, issue_title: str, issue_body: str, issue_comments: Sequence[Mapping[str, str]], @@ -39,11 +49,13 @@ def answer_issue( workdir: Optional[str] = None, build_commands: Optional[Sequence[str]] = None, test_commands: Optional[Sequence[str]] = None, + push_to_remote: Optional[bool] = None, ): """ Processes an issue, generates patches if needed, runs optional builds and tests, and returning the results. Args: + issue_number (int): The number of the issue. issue_title (str): The title of the issue. issue_body (str): The body of the issue. issue_comments (Sequence[Mapping[str, str]]): Comments on the issue. @@ -56,6 +68,7 @@ def answer_issue( workdir (Optional[str]): Working directory for the container. build_commands (Optional[Sequence[str]]): Commands to build the project. test_commands (Optional[Sequence[str]]): Commands to test the project. + push_to_remote (Optional[bool]): Whether to push changes to a remote branch. Returns: Tuple containing: - edit_patch (str): The generated patch for the issue. @@ -64,58 +77,82 @@ def answer_issue( - passed_existing_test (bool): Whether the existing tests passed. - issue_response (str): Response generated for the issue. """ - # Construct the working directory - if dockerfile_content or image_name: - container = UserDefinedContainer( - self.kg_service.kg.get_local_path(), - workdir, - build_commands, - test_commands, - dockerfile_content, - image_name, - ) - else: - container = GeneralContainer(self.kg_service.kg.get_local_path()) - # Initialize the issue graph with the necessary services and parameters - issue_graph = IssueGraph( - advanced_model=self.llm_service.advanced_model, - base_model=self.llm_service.base_model, - kg=self.kg_service.kg, - git_repo=self.repository_service.git_repo, - neo4j_driver=self.neo4j_service.neo4j_driver, - max_token_per_neo4j_result=self.max_token_per_neo4j_result, - container=container, - build_commands=build_commands, - test_commands=test_commands, - ) - # Invoke the issue graph with the provided parameters - output_state = issue_graph.invoke( - issue_title, - issue_body, - issue_comments, - issue_type, - run_build, - run_existing_test, - number_of_candidate_patch, - ) + logger = logging.getLogger("prometheus") + formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + log_file = self.answer_issue_log_dir / f"{timestamp}.log" + file_handler = logging.FileHandler(log_file) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) - if output_state["issue_type"] == IssueType.BUG: - return ( - output_state["edit_patch"], - output_state["passed_reproducing_test"], - output_state["passed_build"], - output_state["passed_existing_test"], - output_state["issue_response"], + try: + # Construct the working directory + if dockerfile_content or image_name: + container = UserDefinedContainer( + self.kg_service.kg.get_local_path(), + workdir, + build_commands, + test_commands, + dockerfile_content, + image_name, + ) + else: + container = GeneralContainer(self.kg_service.kg.get_local_path()) + # Initialize the issue graph with the necessary services and parameters + issue_graph = IssueGraph( + advanced_model=self.llm_service.advanced_model, + base_model=self.llm_service.base_model, + kg=self.kg_service.kg, + git_repo=self.repository_service.git_repo, + neo4j_driver=self.neo4j_service.neo4j_driver, + max_token_per_neo4j_result=self.max_token_per_neo4j_result, + container=container, + build_commands=build_commands, + test_commands=test_commands, ) - elif output_state["issue_type"] == IssueType.QUESTION: - return ( - None, - False, - False, - False, - output_state["issue_response"], + # Invoke the issue graph with the provided parameters + output_state = issue_graph.invoke( + issue_title, + issue_body, + issue_comments, + issue_type, + run_build, + run_existing_test, + number_of_candidate_patch, ) - raise ValueError( - f"Unknown issue type: {output_state['issue_type']}. Expected BUG or QUESTION." - ) + if output_state["issue_type"] == IssueType.BUG: + # push to remote if requested + remote_branch_name = None + if output_state["edit_patch"] and push_to_remote: + remote_branch_name = self.repository_service.push_change_to_remote( + f"Fixes #{issue_number}", output_state["edit_patch"] + ) + + return ( + remote_branch_name, + output_state["edit_patch"], + output_state["passed_reproducing_test"], + output_state["passed_build"], + output_state["passed_existing_test"], + output_state["issue_response"], + ) + elif output_state["issue_type"] == IssueType.QUESTION: + return ( + None, + None, + False, + False, + False, + output_state["issue_response"], + ) + + raise ValueError( + f"Unknown issue type: {output_state['issue_type']}. Expected BUG or QUESTION." + ) + except Exception as e: + logger.error(f"Error in answer_issue: {str(e)}\n{traceback.format_exc()}") + return None, None, False, False, False, None + finally: + logger.removeHandler(file_handler) + file_handler.close() diff --git a/prometheus/app/services/knowledge_graph_service.py b/prometheus/app/services/knowledge_graph_service.py index 5ccaee41..f85751ec 100644 --- a/prometheus/app/services/knowledge_graph_service.py +++ b/prometheus/app/services/knowledge_graph_service.py @@ -3,12 +3,13 @@ from pathlib import Path from typing import Optional +from prometheus.app.services.base_service import BaseService from prometheus.app.services.neo4j_service import Neo4jService from prometheus.graph.knowledge_graph import KnowledgeGraph from prometheus.neo4j import knowledge_graph_handler -class KnowledgeGraphService: +class KnowledgeGraphService(BaseService): """Manages the lifecycle and operations of Knowledge Graphs. This service handles the creation, persistence, and management of Knowledge Graphs @@ -84,3 +85,7 @@ def exists(self) -> bool: def clear(self): self.kg_handler.clear_knowledge_graph() self.kg = None + + def close(self): + """Clear the knowledge graph before closing the service.""" + self.clear() diff --git a/prometheus/app/services/llm_service.py b/prometheus/app/services/llm_service.py index 33567807..143a5669 100644 --- a/prometheus/app/services/llm_service.py +++ b/prometheus/app/services/llm_service.py @@ -4,10 +4,11 @@ from langchain_core.language_models.chat_models import BaseChatModel from langchain_google_genai import ChatGoogleGenerativeAI +from prometheus.app.services.base_service import BaseService from prometheus.chat_models.custom_chat_openai import CustomChatOpenAI -class LLMService: +class LLMService(BaseService): def __init__( self, advanced_model_name: str, diff --git a/prometheus/app/services/neo4j_service.py b/prometheus/app/services/neo4j_service.py index 2ee7f08e..da27fea3 100644 --- a/prometheus/app/services/neo4j_service.py +++ b/prometheus/app/services/neo4j_service.py @@ -1,10 +1,15 @@ """Service for managing Neo4j database driver.""" +import logging + from neo4j import GraphDatabase +from prometheus.app.services.base_service import BaseService + -class Neo4jService: +class Neo4jService(BaseService): def __init__(self, neo4j_uri: str, neo4j_username: str, neo4j_password: str): + self._logger = logging.getLogger("prometheus.app.services.neo4j_service") self.neo4j_driver = GraphDatabase.driver( neo4j_uri, auth=(neo4j_username, neo4j_password), @@ -15,3 +20,4 @@ def __init__(self, neo4j_uri: str, neo4j_username: str, neo4j_password: str): def close(self): self.neo4j_driver.close() + self._logger.info("Neo4j driver connection closed.") diff --git a/prometheus/app/services/repository_service.py b/prometheus/app/services/repository_service.py index 8a02717d..5e6d88d6 100644 --- a/prometheus/app/services/repository_service.py +++ b/prometheus/app/services/repository_service.py @@ -5,11 +5,12 @@ from pathlib import Path from typing import Optional +from prometheus.app.services.base_service import BaseService from prometheus.app.services.knowledge_graph_service import KnowledgeGraphService from prometheus.git.git_repository import GitRepository -class RepositoryService: +class RepositoryService(BaseService): """Manages repository operations. This service provides functionality for Git repository operations including @@ -93,3 +94,7 @@ def clean(self): self.git_repo = None shutil.rmtree(self.target_directory) self.target_directory.mkdir(parents=True) + + def close(self): + """Cleans up the repository service before shutdown.""" + self.clean() diff --git a/prometheus/app/services/service_coordinator.py b/prometheus/app/services/service_coordinator.py deleted file mode 100644 index 05955eb2..00000000 --- a/prometheus/app/services/service_coordinator.py +++ /dev/null @@ -1,213 +0,0 @@ -"""Central coordinator for managing interactions between various prometheus services. - -This coordinator orchestrates the interactions between multiple specialized services -including knowledge graph management, LLM operations, database connections, and -repository management. It provides a unified interface for codebase analysis, -issue handling, and conversation management. -""" - -import logging -import traceback -from datetime import datetime -from pathlib import Path -from typing import Mapping, Optional, Sequence - -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 -from prometheus.app.services.neo4j_service import Neo4jService -from prometheus.app.services.repository_service import RepositoryService -from prometheus.lang_graph.graphs.issue_state import IssueType - - -class ServiceCoordinator: - """Coordinates operations between various prometheus services. - - This class serves as the central orchestrator for all service interactions, - managing the lifecycle of various operations including codebase analysis, - issue handling, and conversation management. It ensures proper initialization, - coordination, and cleanup of all dependent services. - """ - - def __init__( - self, - issue_service: IssueService, - knowledge_graph_service: KnowledgeGraphService, - llm_service: LLMService, - neo4j_service: Neo4jService, - repository_service: RepositoryService, - max_token_per_neo4j_result: int, - github_token: str, - working_directory: Path, - ): - """Initializes the service coordinator with required services. - - Args: - issue_service: Service for issue handling. - knowledge_graph_service: Service for knowledge graph operations. - llm_service: Service for language model operations. - neo4j_service: Service for Neo4j database operations. - repository_service: Service for repository management. - max_token_per_neo4j_result: Maximum number of tokens per Neo4j result. - github_token: GitHub access token for repository operations. - working_directory: Working directory for all Prometheus related files. - """ - self.issue_service = issue_service - self.knowledge_graph_service = knowledge_graph_service - self.llm_service = llm_service - self.neo4j_service = neo4j_service - self.repository_service = repository_service - self.max_token_per_neo4j_result = max_token_per_neo4j_result - self.github_token = github_token - self.working_directory = working_directory - self.answer_issue_log_dir = self.working_directory / "answer_issue_logs" - self.answer_issue_log_dir.mkdir(parents=True, exist_ok=True) - self._logger = logging.getLogger("prometheus.app.services.service_coordinator") - - if ( - self.knowledge_graph_service.get_local_path() - != self.repository_service.get_working_dir() - ): - self._logger.critical( - f"Knowledge graph and repository working directories do not match: {self.knowledge_graph_service.get_local_path()} vs {self.repository_service.get_working_dir()}. Resetting all services." - ) - self.clear() - - def answer_issue( - self, - issue_number: int, - issue_title: str, - issue_body: str, - issue_comments: Sequence[Mapping[str, str]], - issue_type: IssueType, - run_build: bool, - run_existing_test: bool, - number_of_candidate_patch: int, - dockerfile_content: Optional[str] = None, - image_name: Optional[str] = None, - workdir: Optional[str] = None, - build_commands: Optional[Sequence[str]] = None, - test_commands: Optional[Sequence[str]] = None, - push_to_remote: Optional[bool] = None, - ): - """ - Processes an issue, generates patches if needed, runs optional builds and tests, - and can push changes to a remote branch. - - Args: - issue_number: The issue number to process. - issue_title: Title of the issue. - issue_body: Body of the issue. - issue_comments: Comments on the issue. - issue_type: Type of the issue (e.g., bug, feature). - run_build: Whether to run a build after applying the patch. - run_existing_test: Whether to run existing tests after applying the patch. - number_of_candidate_patch: Number of candidate patches to generate. - dockerfile_content: Optional Dockerfile content for user-defined environment. - image_name: Optional name for the Docker image. - workdir: Working directory for the container. - build_commands: Commands to build the project. - test_commands: Commands to test the project. - push_to_remote: Whether to push changes to a remote branch. - Returns: - A tuple containing: - - remote_branch_name: Name of the remote branch if changes were pushed. - - patch: The generated patch for the issue. - - passed_reproducing_test: Whether the reproducing test passed. - - passed_build: Whether the build passed. - - passed_existing_test: Whether existing tests passed. - - issue_response: Response from the issue service after processing. - """ - logger = logging.getLogger("prometheus") - formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - log_file = self.answer_issue_log_dir / f"{timestamp}.log" - file_handler = logging.FileHandler(log_file) - file_handler.setFormatter(formatter) - logger.addHandler(file_handler) - - try: - # fix issue - patch, passed_reproducing_test, passed_build, passed_existing_test, issue_response = ( - self.issue_service.answer_issue( - issue_title=issue_title, - issue_body=issue_body, - issue_comments=issue_comments, - issue_type=issue_type, - run_build=run_build, - run_existing_test=run_existing_test, - number_of_candidate_patch=number_of_candidate_patch, - dockerfile_content=dockerfile_content, - image_name=image_name, - workdir=workdir, - build_commands=build_commands, - test_commands=test_commands, - ) - ) - # push to remote if requested - remote_branch_name = None - if patch and push_to_remote: - remote_branch_name = self.repository_service.push_change_to_remote( - f"Fixes #{issue_number}", patch - ) - return ( - remote_branch_name, - patch, - passed_reproducing_test, - passed_build, - passed_existing_test, - issue_response, - ) - except Exception as e: - logger.error(f"Error in answer_issue: {str(e)}\n{traceback.format_exc()}") - return None, None, False, False, False, None - finally: - logger.removeHandler(file_handler) - file_handler.close() - - def exists_knowledge_graph(self) -> bool: - return self.knowledge_graph_service.exists() - - def upload_local_repository(self, path: Path): - """Uploads a local repository. - - Args: - path: Path to the local repository directory. - """ - self.clear() - self.knowledge_graph_service.build_and_save_knowledge_graph(path) - - def upload_github_repository(self, https_url: str, commit_id: Optional[str] = None): - """Uploads a GitHub repository. - - Args: - https_url: HTTPS URL of the GitHub repository. - commit_id: Optional specific commit to analyze. - """ - # CLean the existing knowledge graph and repository state - self.clear() - # Clone the repository - saved_path = self.repository_service.clone_github_repo( - self.github_token, https_url, commit_id - ) - # Build and save the knowledge graph from the cloned repository - self.knowledge_graph_service.build_and_save_knowledge_graph( - saved_path, https_url, commit_id - ) - - def clear(self): - """Clears all service state and working directories. - - Resets the knowledge graph, cleans repository working directory, - and reinitializes subgraph services. - """ - self.knowledge_graph_service.clear() - self.repository_service.clean() - - def close(self): - """Closes all database connections and releases resources. - - This method should be called when the coordinator is no longer needed - to ensure proper cleanup of database connections and resources. - """ - self.neo4j_service.close() diff --git a/prometheus/app/services/user_service.py b/prometheus/app/services/user_service.py new file mode 100644 index 00000000..32cc9124 --- /dev/null +++ b/prometheus/app/services/user_service.py @@ -0,0 +1,104 @@ +import logging +from typing import Optional + +from argon2 import PasswordHasher +from sqlmodel import Session, or_, select + +from prometheus.app.entity.user import User +from prometheus.app.services.base_service import BaseService +from prometheus.app.services.database_service import DatabaseService +from prometheus.utils.jwt_utils import JWTUtils + + +class UserService(BaseService): + def __init__(self, database_service: DatabaseService): + self.database_service = database_service + self.engine = database_service.engine + self._logger = logging.getLogger("prometheus.app.services.user_service") + self.ph = PasswordHasher() + self.jwt_utils = JWTUtils() + + def create_user( + self, + username: str, + email: str, + password: str, + github_token: Optional[str] = None, + issue_credit: int = 0, + is_superuser: bool = False, + ) -> None: + """ + Create a new superuser and commit it to the database. + + Args: + username (str): Desired username. + email (str): Email address. + password (str): Plaintext password (will be hashed). + github_token (Optional[str]): Optional GitHub token. + issue_credit (int): Optional issue credit. + is_superuser (bool): Whether the user is a superuser. + Returns: + User: The created superuser instance. + """ + with Session(self.engine) as session: + statement = select(User).where(User.username == username) + if session.exec(statement).first(): + raise ValueError(f"Username '{username}' already exists") + statement = select(User).where(User.email == email) + if session.exec(statement).first(): + raise ValueError(f"Email '{email}' already exists") + + hashed_password = self.ph.hash(password) + + user = User( + username=username, + email=email, + password_hash=hashed_password, + github_token=github_token, + issue_credit=issue_credit, + is_superuser=is_superuser, + ) + session.add(user) + session.commit() + session.refresh(user) + + def login(self, username: str, email: str, password: str) -> str: + """ + Log in a user by verifying their credentials and return an access token. + + Args: + username (str): Username of the user. + email (str): Email address of the user. + password (str): Plaintext password. + """ + with Session(self.engine) as session: + statement = select(User).where(or_(User.username == username, User.email == email)) + user = session.exec(statement).first() + + if not user: + raise ValueError("Invalid username or email") + + if not self.ph.verify(user.password_hash, password): + raise ValueError("Invalid password") + + # Generate and return a JWT token for the user + token = self.jwt_utils.generate_token({"user_id": user.id}) + return token + + # Create a superuser and commit it to the database + def create_superuser( + self, + username: str, + email: str, + password: str, + github_token: Optional[str] = None, + ) -> None: + """ + Create a new superuser in the database. + + This method creates a superuser with the provided credentials and commits it to the database. + """ + self.create_user( + username, email, password, github_token, is_superuser=True, issue_credit=999999 + ) + self._logger.info(f"Superuser '{username}' created successfully.") diff --git a/prometheus/configuration/config.py b/prometheus/configuration/config.py index ea0b2860..f9320acb 100644 --- a/prometheus/configuration/config.py +++ b/prometheus/configuration/config.py @@ -1,3 +1,56 @@ -from dynaconf import Dynaconf +from typing import List, Literal -settings = Dynaconf(envvar_prefix="PROMETHEUS", load_dotenv=True) +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + model_config = SettingsConfigDict( + env_file=".env", env_file_encoding="utf-8", env_prefix="PROMETHEUS_" + ) + # General settings + version: str = "1.1" + BASE_URL: str = f"/api/v{version}" + ENVIRONMENT: Literal["local", "production"] + BACKEND_CORS_ORIGINS: List[str] + PROJECT_NAME: str = "Prometheus" + + # Logging + LOGGING_LEVEL: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] + + # Neo4j + NEO4J_URI: str + NEO4J_USERNAME: str + NEO4J_PASSWORD: str + NEO4J_BATCH_SIZE: int + + # Knowledge Graph + WORKING_DIRECTORY: str + KNOWLEDGE_GRAPH_MAX_AST_DEPTH: int + KNOWLEDGE_GRAPH_CHUNK_SIZE: int + KNOWLEDGE_GRAPH_CHUNK_OVERLAP: int + MAX_TOKEN_PER_NEO4J_RESULT: int + + # LLM models + ADVANCED_MODEL: str + BASE_MODEL: str + + # API Keys + ANTHROPIC_API_KEY: str + GEMINI_API_KEY: str + OPENAI_FORMAT_BASE_URL: str + OPENAI_FORMAT_API_KEY: str + + # Model parameters + MAX_INPUT_TOKENS: int + TEMPERATURE: float + MAX_OUTPUT_TOKENS: int + + # Database + DATABASE_URL: str + + # JWT Configuration + JWT_SECRET_KEY: str + ACCESS_TOKEN_EXPIRE_TIME: int = 7 # days + + +settings = Settings() diff --git a/prometheus/exceptions/file_operation_exceptions.py b/prometheus/exceptions/file_operation_exception.py similarity index 100% rename from prometheus/exceptions/file_operation_exceptions.py rename to prometheus/exceptions/file_operation_exception.py diff --git a/prometheus/exceptions/jwt_exception.py b/prometheus/exceptions/jwt_exception.py new file mode 100644 index 00000000..d5589110 --- /dev/null +++ b/prometheus/exceptions/jwt_exception.py @@ -0,0 +1,12 @@ +from prometheus.exceptions.server_exception import ServerException + + +class JWTException(ServerException): + """ + class for JWT exceptions. + This exception is raised when there is an issue with JWT operations, + such as token generation or validation. + """ + + def __init__(self, code: int = 401, message: str = "An error occurred with the JWT operation."): + super().__init__(code, message) diff --git a/prometheus/exceptions/server_exception.py b/prometheus/exceptions/server_exception.py new file mode 100644 index 00000000..5bf81cc7 --- /dev/null +++ b/prometheus/exceptions/server_exception.py @@ -0,0 +1,10 @@ +class ServerException(Exception): + """ + Base class for server exceptions. + This exception is raised when there is an issue with server operations. + """ + + def __init__(self, code: int, message: str): + super().__init__(message) + self.code = code + self.message = message diff --git a/prometheus/lang_graph/nodes/context_extraction_node.py b/prometheus/lang_graph/nodes/context_extraction_node.py index e1f8db73..cf331b34 100644 --- a/prometheus/lang_graph/nodes/context_extraction_node.py +++ b/prometheus/lang_graph/nodes/context_extraction_node.py @@ -5,7 +5,7 @@ from langchain_core.messages import SystemMessage from pydantic import BaseModel, Field -from prometheus.exceptions.file_operation_exceptions import FileOperationException +from prometheus.exceptions.file_operation_exception import FileOperationException from prometheus.lang_graph.subgraphs.context_retrieval_state import ContextRetrievalState from prometheus.models.context import Context from prometheus.utils.file_utils import read_file_with_line_numbers diff --git a/prometheus/script/create_db_and_tables.py b/prometheus/script/create_db_and_tables.py deleted file mode 100644 index 3147c90e..00000000 --- a/prometheus/script/create_db_and_tables.py +++ /dev/null @@ -1,4 +0,0 @@ -from prometheus.app.db import create_db_and_tables - -if __name__ == "__main__": - create_db_and_tables() diff --git a/prometheus/script/create_superuser.py b/prometheus/script/create_superuser.py index bf6bf5f9..20829308 100644 --- a/prometheus/script/create_superuser.py +++ b/prometheus/script/create_superuser.py @@ -1,6 +1,12 @@ import argparse +import re -from prometheus.app.db import create_superuser +from prometheus.app.services.database_service import DatabaseService +from prometheus.app.services.user_service import UserService +from prometheus.configuration.config import settings + +database_service: DatabaseService = DatabaseService(settings.DATABASE_URL) +user_service: UserService = UserService(database_service) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Create a superuser account.") @@ -11,7 +17,11 @@ args = parser.parse_args() - create_superuser( + pattern = r"^[^@\s]+@[^@\s]+\.[^@\s]+$" + if not re.match(pattern, args.email): + raise ValueError("Invalid email format") + + user_service.create_superuser( username=args.username, email=args.email, password=args.password, diff --git a/prometheus/script/generate_jwt_token.py b/prometheus/script/generate_jwt_token.py new file mode 100644 index 00000000..5e500c8f --- /dev/null +++ b/prometheus/script/generate_jwt_token.py @@ -0,0 +1,15 @@ +import base64 +import secrets + + +def generate_jwt_secret_token(length: int = 64) -> str: + """ + Generate a secure JWT secret token. + """ + token_bytes = secrets.token_bytes(length) + return base64.urlsafe_b64encode(token_bytes).decode("utf-8") + + +if __name__ == "__main__": + secret = generate_jwt_secret_token() + print(f"JWT_SECRET_TOKEN={secret}") diff --git a/prometheus/utils/file_utils.py b/prometheus/utils/file_utils.py index fb3252af..c63f79dd 100644 --- a/prometheus/utils/file_utils.py +++ b/prometheus/utils/file_utils.py @@ -1,7 +1,7 @@ import os from pathlib import Path -from prometheus.exceptions.file_operation_exceptions import FileOperationException +from prometheus.exceptions.file_operation_exception import FileOperationException def read_file_with_line_numbers( diff --git a/prometheus/utils/jwt_utils.py b/prometheus/utils/jwt_utils.py new file mode 100644 index 00000000..684652f5 --- /dev/null +++ b/prometheus/utils/jwt_utils.py @@ -0,0 +1,33 @@ +import datetime + +import jwt + +from prometheus.configuration.config import settings +from prometheus.exceptions.jwt_exception import JWTException + + +class JWTUtils: + """JWT Utility for token generation and validation""" + + def __init__(self, algorithm="HS256"): + """Initialize JWT utils with configuration""" + self.secret_key = settings.JWT_SECRET_KEY + self.algorithm = algorithm + self.expire_time = int(settings.ACCESS_TOKEN_EXPIRE_TIME) + + def generate_token(self, payload): + """Generate JWT token""" + payload_copy = payload.copy() + payload_copy["exp"] = datetime.datetime.now() + datetime.timedelta(days=self.expire_time) + token = jwt.encode(payload_copy, self.secret_key, algorithm=self.algorithm) + return token + + def decode_token(self, token): + """Decode and parse JWT token""" + try: + decoded = jwt.decode(token, self.secret_key, algorithms=[self.algorithm]) + return decoded + except jwt.ExpiredSignatureError: + raise JWTException(message="Token expired") + except jwt.InvalidTokenError: + raise JWTException(message="Invalid token") diff --git a/pyproject.toml b/pyproject.toml index 43acfb9b..1d163fcb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,9 +24,10 @@ dependencies = [ "dynaconf>=3.2.6", "docker>=7.1.0", "unidiff>=0.7.5", - "passlib[bcrypt]>=1.7.4", + "argon2-cffi>=23.1.0", "sqlmodel==0.0.24", - "psycopg2-binary" + "psycopg2-binary", + "pyjwt==2.6.0", ] requires-python = ">= 3.11" diff --git a/tests/app/api/test_auth.py b/tests/app/api/test_auth.py new file mode 100644 index 00000000..240cfe9e --- /dev/null +++ b/tests/app/api/test_auth.py @@ -0,0 +1,38 @@ +from unittest import mock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from prometheus.app.api import auth +from prometheus.app.exception_handler import register_exception_handlers + +app = FastAPI() +register_exception_handlers(app) +app.include_router(auth.router, prefix="/auth", tags=["auth"]) +client = TestClient(app) + + +@pytest.fixture +def mock_service(): + service = mock.MagicMock() + app.state.service = service + yield service + + +def test_login(mock_service): + mock_service["user_service"].login.return_value = "your_access_token" + response = client.post( + "/auth/login", + json={ + "username": "testuser", + "email": "test@gmail.com", + "password": "passwordpassword", + }, + ) + assert response.status_code == 200 + assert response.json() == { + "code": 200, + "message": "success", + "data": {"access_token": "your_access_token"}, + } diff --git a/tests/app/api/test_issue.py b/tests/app/api/test_issue.py index ae5e7c1a..5367c872 100644 --- a/tests/app/api/test_issue.py +++ b/tests/app/api/test_issue.py @@ -5,23 +5,25 @@ from fastapi.testclient import TestClient from prometheus.app.api import issue +from prometheus.app.exception_handler import register_exception_handlers from prometheus.lang_graph.graphs.issue_state import IssueType app = FastAPI() +register_exception_handlers(app) app.include_router(issue.router, prefix="/issue", tags=["issue"]) client = TestClient(app) @pytest.fixture -def mock_service_coordinator(): - service_coordinator = mock.MagicMock() - app.state.service_coordinator = service_coordinator - yield service_coordinator +def mock_service(): + service = mock.MagicMock() + app.state.service = service + yield service -def test_answer_issue(mock_service_coordinator): - mock_service_coordinator.exists_knowledge_graph.return_value = True - mock_service_coordinator.answer_issue.return_value = ( +def test_answer_issue(mock_service): + mock_service["knowledge_graph_service"].exists_knowledge_graph.return_value = True + mock_service["issue_service"].answer_issue.return_value = ( "feature/fix-42", # remote_branch_name "test patch", # patch True, # passed_reproducing_test @@ -41,17 +43,21 @@ def test_answer_issue(mock_service_coordinator): assert response.status_code == 200 assert response.json() == { - "remote_branch_name": "feature/fix-42", - "patch": "test patch", - "passed_reproducing_test": True, - "passed_build": True, - "passed_existing_test": True, - "issue_response": "Issue fixed", + "code": 200, + "message": "success", + "data": { + "remote_branch_name": "feature/fix-42", + "patch": "test patch", + "passed_reproducing_test": True, + "passed_build": True, + "passed_existing_test": True, + "issue_response": "Issue fixed", + }, } -def test_answer_issue_no_repository(mock_service_coordinator): - mock_service_coordinator.exists_knowledge_graph.return_value = False +def test_answer_issue_no_repository(mock_service): + mock_service["knowledge_graph_service"].exists.return_value = False response = client.post( "/issue/answer", @@ -61,8 +67,8 @@ def test_answer_issue_no_repository(mock_service_coordinator): assert response.status_code == 404 -def test_answer_issue_invalid_container_config(mock_service_coordinator): - mock_service_coordinator.exists_knowledge_graph.return_value = True +def test_answer_issue_invalid_container_config(mock_service): + mock_service["knowledge_graph_service"].exists.return_value = True response = client.post( "/issue/answer", @@ -78,9 +84,9 @@ def test_answer_issue_invalid_container_config(mock_service_coordinator): assert response.status_code == 400 -def test_answer_issue_with_container(mock_service_coordinator): - mock_service_coordinator.exists_knowledge_graph.return_value = True - mock_service_coordinator.answer_issue.return_value = ( +def test_answer_issue_with_container(mock_service): + mock_service["knowledge_graph_service"].exists.return_value = True + mock_service["issue_service"].answer_issue.return_value = ( "feature/fix-42", "test patch", True, @@ -102,7 +108,7 @@ def test_answer_issue_with_container(mock_service_coordinator): response = client.post("/issue/answer", json=test_payload) assert response.status_code == 200 - mock_service_coordinator.answer_issue.assert_called_once_with( + mock_service["issue_service"].answer_issue.assert_called_once_with( issue_number=42, issue_title="Test Issue", issue_body="Test description", diff --git a/tests/app/api/test_repository.py b/tests/app/api/test_repository.py index 3f983899..73f5bcf4 100644 --- a/tests/app/api/test_repository.py +++ b/tests/app/api/test_repository.py @@ -5,39 +5,23 @@ from fastapi.testclient import TestClient from prometheus.app.api import repository -from tests.test_utils import test_project_paths +from prometheus.app.exception_handler import register_exception_handlers app = FastAPI() +register_exception_handlers(app) app.include_router(repository.router, prefix="/repository", tags=["repository"]) client = TestClient(app) @pytest.fixture -def mock_service_coordinator(): - service_coordinator = mock.MagicMock() - app.state.service_coordinator = service_coordinator - yield service_coordinator +def mock_service(): + service = mock.MagicMock() + app.state.service = service + yield service -def test_upload_local_repository(mock_service_coordinator): - mock_service_coordinator.upload_local_repository.return_value = None - response = client.get( - "/repository/local", - params={"local_repository": test_project_paths.TEST_PROJECT_PATH.absolute().as_posix()}, - ) - assert response.status_code == 200 - - -def test_upload_fake_local_repository(): - response = client.get( - "/repository/local/", - params={"local_repository": "/foo/bar/"}, - ) - assert response.status_code == 404 - - -def test_delete(mock_service_coordinator): - mock_service_coordinator.exists_knowledge_graph.return_value = True - mock_service_coordinator.clear.return_value = None +def test_delete(mock_service): + mock_service["knowledge_graph_service"].exists.return_value = True + mock_service["knowledge_graph_service"].clear.return_value = None response = client.get("repository/delete") assert response.status_code == 200 diff --git a/tests/app/services/test_database_service.py b/tests/app/services/test_database_service.py new file mode 100644 index 00000000..5c7930b0 --- /dev/null +++ b/tests/app/services/test_database_service.py @@ -0,0 +1,17 @@ +import pytest + +from prometheus.app.services.database_service import DatabaseService +from tests.test_utils.fixtures import postgres_container_fixture # noqa: F401 + + +@pytest.mark.slow +def test_database_service(postgres_container_fixture): # noqa: F811 + url = postgres_container_fixture.get_connection_url() + database_service = DatabaseService(url) + assert database_service.engine is not None + + try: + database_service.engine.connect() + database_service.engine.dispose() # Ensure connection is valid + except Exception as e: + pytest.fail(f"Connection verification failed: {e}") diff --git a/tests/app/services/test_db_engine.py b/tests/app/services/test_db_engine.py deleted file mode 100644 index 86f3c949..00000000 --- a/tests/app/services/test_db_engine.py +++ /dev/null @@ -1,12 +0,0 @@ -import pytest -from sqlmodel import create_engine - -from prometheus.app.db import create_db_and_tables - - -@pytest.fixture(scope="session") -def test_engine(postgres_container_fixture): - url = postgres_container_fixture.get_connection_url() - engine = create_engine(url, echo=True) - create_db_and_tables() - yield engine diff --git a/tests/app/services/test_issue_service.py b/tests/app/services/test_issue_service.py index 861b608f..6e8f2224 100644 --- a/tests/app/services/test_issue_service.py +++ b/tests/app/services/test_issue_service.py @@ -49,6 +49,7 @@ def issue_service(mock_kg_service, mock_repository_service, mock_neo4j_service, mock_neo4j_service, mock_llm_service, max_token_per_neo4j_result=1000, + working_directory="/tmp/working_dir/", ) @@ -77,6 +78,7 @@ def test_answer_issue_with_general_container(issue_service, monkeypatch): # Exercise result = issue_service.answer_issue( + issue_number=-1, issue_title="Test Issue", issue_body="Test Body", issue_comments=[], @@ -101,7 +103,7 @@ def test_answer_issue_with_general_container(issue_service, monkeypatch): build_commands=None, test_commands=None, ) - assert result == ("test_patch", True, True, True, "test_response") + assert result == (None, "test_patch", True, True, True, "test_response") def test_answer_issue_with_user_defined_container(issue_service, monkeypatch): @@ -122,6 +124,7 @@ def test_answer_issue_with_user_defined_container(issue_service, monkeypatch): # Exercise result = issue_service.answer_issue( + issue_number=-1, issue_title="Test Issue", issue_body="Test Body", issue_comments=[], @@ -145,4 +148,4 @@ def test_answer_issue_with_user_defined_container(issue_service, monkeypatch): "FROM python:3.8", "test-image", ) - assert result == (None, False, False, False, "test_response") + assert result == (None, None, False, False, False, "test_response") diff --git a/tests/app/services/test_repository_service.py b/tests/app/services/test_repository_service.py index 99d1af27..6ce2ac65 100644 --- a/tests/app/services/test_repository_service.py +++ b/tests/app/services/test_repository_service.py @@ -32,7 +32,7 @@ def mock_knowledge_graph(): @pytest.fixture def service(mock_kg_service): - working_dir = Path("/test/working/dir") + working_dir = "/test/working/dir" # Don't mock GitRepository in the fixture anymore with patch("pathlib.Path.mkdir"): return RepositoryService( diff --git a/tests/app/services/test_service_coordinator.py b/tests/app/services/test_service_coordinator.py deleted file mode 100644 index 71652628..00000000 --- a/tests/app/services/test_service_coordinator.py +++ /dev/null @@ -1,115 +0,0 @@ -import tempfile -from pathlib import Path -from unittest.mock import Mock - -import pytest - -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 -from prometheus.app.services.neo4j_service import Neo4jService -from prometheus.app.services.repository_service import RepositoryService -from prometheus.app.services.service_coordinator import ServiceCoordinator - - -@pytest.fixture -def mock_services(): - llm_service = Mock(spec=LLMService) - llm_service.model = Mock() - - issue_service = Mock(sepc=IssueService) - knowledge_graph_service = Mock(spec=KnowledgeGraphService) - knowledge_graph_service.kg = None - knowledge_graph_service.get_local_path.return_value = None - neo4j_service = Mock(spec=Neo4jService) - repository_service = Mock(spec=RepositoryService) - repository_service.get_working_dir.return_value = None - - yield { - "issue_service": issue_service, - "knowledge_graph_service": knowledge_graph_service, - "llm_service": llm_service, - "neo4j_service": neo4j_service, - "max_token_per_neo4j_result": 1000, - "repository_service": repository_service, - "github_token": "test_token", - "working_directory": Path(tempfile.TemporaryDirectory().name), - } - - -@pytest.fixture -def service_coordinator(mock_services): - coordinator = ServiceCoordinator( - mock_services["issue_service"], - mock_services["knowledge_graph_service"], - mock_services["llm_service"], - mock_services["neo4j_service"], - mock_services["repository_service"], - mock_services["max_token_per_neo4j_result"], - mock_services["github_token"], - mock_services["working_directory"], - ) - return coordinator - - -def test_initialization(service_coordinator, mock_services): - """Test that services are properly initialized""" - assert service_coordinator.knowledge_graph_service == mock_services["knowledge_graph_service"] - assert service_coordinator.llm_service == mock_services["llm_service"] - assert service_coordinator.neo4j_service == mock_services["neo4j_service"] - assert service_coordinator.repository_service == mock_services["repository_service"] - - -def test_upload_local_repository(service_coordinator, mock_services): - """Test local repository upload process""" - test_path = Path("/test/path") - - service_coordinator.upload_local_repository(test_path) - - # Verify the sequence of operations - mock_services["knowledge_graph_service"].build_and_save_knowledge_graph.assert_called_once_with( - test_path - ) - - -def test_upload_github_repository(service_coordinator, mock_services): - """Test GitHub repository upload process""" - test_url = "https://github.com/test/repo" - test_commit = "abc123" - saved_path = Path("/saved/path") - mock_services["repository_service"].clone_github_repo.return_value = saved_path - - service_coordinator.upload_github_repository(test_url, test_commit) - - # Verify the sequence of operations - mock_services["repository_service"].clone_github_repo.assert_called_once_with( - mock_services["github_token"], test_url, test_commit - ) - mock_services["knowledge_graph_service"].build_and_save_knowledge_graph.assert_called_once_with( - saved_path, test_url, test_commit - ) - - -def test_clear(service_coordinator, mock_services): - """Test clearing of services""" - service_coordinator.clear() - - mock_services["knowledge_graph_service"].clear.assert_called_once() - mock_services["repository_service"].clean.assert_called_once() - - -def test_close(service_coordinator, mock_services): - """Test proper closing of services""" - service_coordinator.close() - - mock_services["neo4j_service"].close.assert_called_once() - - -def test_exists_knowledge_graph(service_coordinator, mock_services): - """Test knowledge graph existence check""" - mock_services["knowledge_graph_service"].exists.return_value = True - - result = service_coordinator.exists_knowledge_graph() - - mock_services["knowledge_graph_service"].exists.assert_called_once() - assert result is True diff --git a/tests/app/services/test_user_service.py b/tests/app/services/test_user_service.py new file mode 100644 index 00000000..3136972f --- /dev/null +++ b/tests/app/services/test_user_service.py @@ -0,0 +1,27 @@ +import pytest + +from prometheus.app.services.database_service import DatabaseService +from prometheus.app.services.user_service import UserService +from tests.test_utils.fixtures import postgres_container_fixture # noqa: F401 + + +@pytest.fixture +def mock_database_service(postgres_container_fixture): # noqa: F811 + service = DatabaseService(postgres_container_fixture.get_connection_url()) + service.start() + yield service + service.close() + + +def test_create_superuser(mock_database_service): + # Exercise + service = UserService(mock_database_service) + service.create_superuser("testuser", "test@gmail.com", "password123", github_token="gh_token") + + +def test_login(mock_database_service): + # Exercise + service = UserService(mock_database_service) + access_token = service.login("testuser", "test@gmail.com", "password123") + # Verify + assert access_token is not None diff --git a/tests/app/test_db.py b/tests/app/test_db.py deleted file mode 100644 index 589fca1b..00000000 --- a/tests/app/test_db.py +++ /dev/null @@ -1,49 +0,0 @@ -import pytest -from sqlmodel import Session, select - -from prometheus.app.db import create_superuser -from prometheus.app.entity.user import User -from tests.app.services.test_db_engine import test_engine # noqa: F401 -from tests.test_utils.fixtures import postgres_container_fixture # noqa: F401 - - -@pytest.mark.slow -def test_create_superuser(postgres_container_fixture, test_engine): # noqa: F811 - username = "admin" - email = "admin@example.com" - password = "strongpassword" - github_token = "ghp_123456" - - # Create superuser - create_superuser( - username=username, - email=email, - password=password, - github_token=github_token, - ) - - with Session(test_engine) as session: - user = session.exec(select(User).where(User.username == username)).first() - - assert user is not None - assert user.email == email - assert user.is_superuser is True - assert user.issue_credit == 999999 - assert user.github_token == github_token - assert user.password_hash != password # Password must be hashed - - -@pytest.mark.slow -def test_duplicate_superuser_raises(postgres_container_fixture, test_engine): # noqa: F811 - username = "admin2" - email = "admin2@example.com" - password = "password" - - create_superuser(username, email, password) - - # Trying again with same username or email - with pytest.raises(ValueError, match="Username 'admin2' already exists"): - create_superuser(username, "different@example.com", "pass") - - with pytest.raises(ValueError, match="Email 'admin2@example.com' already exists"): - create_superuser("different_user", email, "pass") diff --git a/tests/app/test_main.py b/tests/app/test_main.py index 8ad0c0c1..2a45d5f5 100644 --- a/tests/app/test_main.py +++ b/tests/app/test_main.py @@ -6,10 +6,10 @@ @pytest.fixture def mock_dependencies(): - """Mock the service coordinator dependencies""" - mock_coordinator = MagicMock() - with patch("prometheus.app.dependencies.initialize_services", return_value=mock_coordinator): - yield mock_coordinator + """Mock the service dependencies""" + mock_service = MagicMock() + with patch("prometheus.app.dependencies.initialize_services", return_value=mock_service): + yield mock_service @pytest.fixture @@ -24,5 +24,5 @@ def test_client(mock_dependencies): def test_app_initialization(test_client, mock_dependencies): """Test that the app initializes correctly with mocked dependencies""" - assert test_client.app.state.service_coordinator is not None - assert test_client.app.state.service_coordinator == mock_dependencies + assert test_client.app.state.service is not None + assert test_client.app.state.service == mock_dependencies diff --git a/tests/test_utils/fixtures.py b/tests/test_utils/fixtures.py index 4073a844..a0e7940a 100644 --- a/tests/test_utils/fixtures.py +++ b/tests/test_utils/fixtures.py @@ -51,17 +51,13 @@ def empty_neo4j_container_fixture(): @pytest.fixture(scope="session") def postgres_container_fixture(): - container = ( - PostgresContainer( - image=POSTGRES_IMAGE, - username=POSTGRES_USERNAME, - password=POSTGRES_PASSWORD, - dbname=POSTGRES_DB, - port=5432, - ) - .with_name(f"postgres_container_{uuid.uuid4().hex[:12]}") - .with_bind_ports(5432, 5432) # Bind to a specific port for testing - ) + container = PostgresContainer( + image=POSTGRES_IMAGE, + username=POSTGRES_USERNAME, + password=POSTGRES_PASSWORD, + dbname=POSTGRES_DB, + port=5432, + ).with_name(f"postgres_container_{uuid.uuid4().hex[:12]}") with container as postgres_container: yield postgres_container