devai-api/app/gitlab_utils.py

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime
import gitlab
import os
from git import Repo

from langchain.agents import AgentType, initialize_agent
from langchain_community.agent_toolkits.gitlab.toolkit import GitLabToolkit
from langchain_community.utilities.gitlab import GitLabAPIWrapper

from google.cloud.aiplatform import telemetry
from langchain_google_vertexai import ChatVertexAI
from vertexai.generative_models import GenerativeModel

from .github_utils import delete_folder
from .constants import USER_AGENT, MODEL_NAME
from .file_processor import format_files_as_string


LLM_INSTRUCTION_TEMPLATE = """You are principal software engineer and given requirements below for implementation.
        You must follow rules when generating implementation:
        - you must use existing codebase from the context below
        - in your response, generate complete source code files and not diffs
        - in your response, include full filepath with name before file content, excluding repository name
        - must return response using sample format
        
        REQUIREMENTS:
        {prompt}

        SAMPLE RESPONSE FORMAT:
        menu-service/src/main/java/org/google/demo/Menu.java
        OLD <<<<
        existing code from the context below for the file
        >>>> OLD
        NEW <<<<
        new generated code by LLM
        >>>> NEW

        CONTEXT:
        {codebase}
        """

PR_PROMPT_TEMPLATE = """Create GitLab merge request using provided details below.
        Update existing files or Create new files, commit them and push them to opened merge request.
        
        DETAILS: 
        {response_text}
        """

class MergeRequestError(Exception):
    """Custom exception for merge request creation."""
    pass

def _create_branch() -> str:
    """Creates a new branch in GitLab for a merge request.

    Returns:
        str: The name of the newly created branch.
    """
    gitlab_url = os.environ["GITLAB_URL"]
    gitlab_base_branch = os.environ["GITLAB_BASE_BRANCH"]
    gitlab_repo_name = os.environ["GITLAB_REPOSITORY"]
    gitlab_access_token = os.environ["GITLAB_PERSONAL_ACCESS_TOKEN"]

    gl = gitlab.Gitlab(gitlab_url, private_token=gitlab_access_token)
    project = gl.projects.get(gitlab_repo_name)

    new_branch_name = f'feature/generated-{datetime.datetime.now().strftime("%m%d%Y-%H%M")}'
    project.branches.create({'branch': new_branch_name, 'ref': gitlab_base_branch})
    print("Created new branch:", new_branch_name)

    return new_branch_name

def _clone_repo(repo_name: str) -> Repo:
    """Clones a GitLab repository to the local file system.

    Args:
        repo_name (str): The name of the repository to clone.

    Returns:
        Repo: The cloned repository object.
    """
    gitlab_repo_name = os.environ["GITLAB_REPOSITORY"]
    gitlab_access_token = os.environ["GITLAB_PERSONAL_ACCESS_TOKEN"]

    repo = Repo.clone_from(
        f"https://oauth2:{gitlab_access_token}@gitlab.com/{gitlab_repo_name}.git",
        repo_name,
    )

    return repo

def _init_agent(new_gitlab_branch: str):
    """Initializes an agent with the GitLab toolkit.

    Args:
        new_gitlab_branch (str): The name of the branch to use for the agent.

    Returns:
        Agent: The initialized agent.
    """
    gitlab = GitLabAPIWrapper(gitlab_branch=new_gitlab_branch)
    toolkit = GitLabToolkit.from_gitlab_api_wrapper(gitlab)

    with telemetry.tool_context_manager(USER_AGENT):
        llm = ChatVertexAI(model_name=MODEL_NAME,
            convert_system_message_to_human=True,
            temperature=0.2,
            max_output_tokens=8192)

    agent = initialize_agent(
        toolkit.get_tools(), 
        llm, 
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, 
        verbose=True,
        handle_parsing_errors=True,
        max_iterations=10,
        return_intermediate_steps=True,
        early_stopping_method="generate",
    )

    return agent

def _generate_llm_instructions(prompt: str, codebase: str) -> str:
    """Generates instructions for the Language Model (LLM).

    Args:
        prompt (str): The user prompt.
        codebase (str): The codebase context.

    Returns:
        str: The formatted instructions for the LLM.
    """
    return LLM_INSTRUCTION_TEMPLATE.format(prompt=prompt, codebase=codebase)

def _get_llm_response(instructions: str, repo_name: str) -> str:
    """Sends instructions to the LLM and retrieves its response.

    Args:
        instructions (str): The instructions for the LLM.
        repo_name (str): The name of the repository.

    Returns:
        str: The response from the LLM.
    """
    code_chat_model = GenerativeModel(MODEL_NAME)

    with telemetry.tool_context_manager(USER_AGENT):
        code_chat = code_chat_model.start_chat(response_validation=False)
        response = code_chat.send_message(instructions)
    
    # Remove repo name from the response
    return response.text.replace(f"{repo_name}/", "")

def _create_gitlab_merge_request(response_text: str, agent) -> None:
    """Creates a GitLab merge request using the provided response text and agent.

    Args:
        response_text (str): The response text to use for the merge request.
        agent: The agent to use for creating the merge request.
    """
    pr_prompt = PR_PROMPT_TEMPLATE.format(response_text=response_text)
    agent.invoke(pr_prompt)

def create_merge_request(prompt: str) -> str:
    """Creates a new GitLab merge request.

    Args:
        prompt (str): The prompt describing the changes to be made.

    Returns:
        str: The implementation details returned by the LLM.

    Raises:
        MergeRequestError: If an error occurs during the merge request creation.
    """
    _, repo_name = get_repo_details()
    try:
        delete_folder(repo_name)    
        _clone_repo(repo_name)
        codebase = load_codebase(repo_name, prompt)

        instructions = _generate_llm_instructions(prompt, codebase)
        implementation_details = _get_llm_response(instructions, repo_name)

        new_gitlab_branch = _create_branch()
        agent = _init_agent(new_gitlab_branch)
        
        _create_gitlab_merge_request(implementation_details, agent)
        
        return implementation_details
    except Exception as e:
         raise MergeRequestError(f"Failed to create merge request: {e}") from e
    finally:
        delete_folder(repo_name)

def get_repo_details() -> tuple[str, str]:
    """Extracts the repository owner and name from the GITLAB_REPOSITORY environment variable.

    Returns:
        tuple[str, str]: A tuple containing the repository owner and name.
    """
    gitlab_repo_name = os.environ["GITLAB_REPOSITORY"]
    repo = gitlab_repo_name.split("/")
    return (repo[0], repo[1])

def load_codebase(repo_name: str, prompt: str) -> str:
    """Loads the codebase from the specified repository.

    Args:
        repo_name (str): The name of the repository.
        prompt (str): The user prompt.

    Returns:
        str: The formatted codebase as a string.
    """
    # Defaults to repo root
    service = ""

    if "menu service" in prompt.lower():
        service = "menu-service"
    if "customer service" in prompt.lower():
        service = "customer-service/src"
    if "customer ui" in prompt.lower():
        service = "customer-ui/src"
    if "inventory service" in prompt.lower():
        service = "inventory-service/spanner"
    if "order-service" in prompt.lower():
        service = "order-service"

    code_path = f"{repo_name}/{service}"

    return format_files_as_string(code_path)