community: add method to create branch and list files for gitlab tool

libs/community/langchain_community/agent_toolkits/gitlab/toolkit.py

@@ -1,6 +1,6 @@
 """GitLab Toolkit."""
 
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 from langchain_core.tools import BaseTool
 from langchain_core.tools.base import BaseToolkit
@@ -9,15 +9,35 @@
     COMMENT_ON_ISSUE_PROMPT,
     CREATE_FILE_PROMPT,
     CREATE_PULL_REQUEST_PROMPT,
+    CREATE_REPO_BRANCH,
     DELETE_FILE_PROMPT,
     GET_ISSUE_PROMPT,
     GET_ISSUES_PROMPT,
+    GET_REPO_FILES_FROM_DIRECTORY,
+    GET_REPO_FILES_IN_BOT_BRANCH,
+    GET_REPO_FILES_IN_MAIN,
+    LIST_REPO_BRANCES,
     READ_FILE_PROMPT,
+    SET_ACTIVE_BRANCH,
     UPDATE_FILE_PROMPT,
 )
 from langchain_community.tools.gitlab.tool import GitLabAction
 from langchain_community.utilities.gitlab import GitLabAPIWrapper
 
+# only include a subset of tools by default to avoid a breaking change, where
+# new tools are added to the toolkit and the user's code breaks because of
+# the new tools
+DEFAULT_INCLUDED_TOOLS = [
+    "get_issues",
+    "get_issue",
+    "comment_on_issue",
+    "create_pull_request",
+    "create_file",
+    "read_file",
+    "update_file",
+    "delete_file",
+]
+
 
 class GitLabToolkit(BaseToolkit):
     """GitLab Toolkit.
@@ -39,7 +59,10 @@ class GitLabToolkit(BaseToolkit):
 
     @classmethod
     def from_gitlab_api_wrapper(
-        cls, gitlab_api_wrapper: GitLabAPIWrapper
+        cls,
+        gitlab_api_wrapper: GitLabAPIWrapper,
+        *,
+        included_tools: Optional[List[str]] = None,
     ) -> "GitLabToolkit":
         """Create a GitLabToolkit from a GitLabAPIWrapper.
 
@@ -50,6 +73,10 @@ def from_gitlab_api_wrapper(
             GitLabToolkit. The GitLab toolkit.
         """
 
+        tools_to_include = (
+            included_tools if included_tools is not None else DEFAULT_INCLUDED_TOOLS
+        )
+
         operations: List[Dict] = [
             {
                 "mode": "get_issues",
@@ -91,6 +118,41 @@ def from_gitlab_api_wrapper(
                 "name": "Delete File",
                 "description": DELETE_FILE_PROMPT,
             },
+            {
+                "mode": "create_branch",
+                "name": "Create a new branch",
+                "description": CREATE_REPO_BRANCH,
+            },
+            {
+                "mode": "list_branches_in_repo",
+                "name": "Get the list of branches",
+                "description": LIST_REPO_BRANCES,
+            },
+            {
+                "mode": "set_active_branch",
+                "name": "Change the active branch",
+                "description": SET_ACTIVE_BRANCH,
+            },
+            {
+                "mode": "list_files_in_main_branch",
+                "name": "Overview of existing files in Main branch",
+                "description": GET_REPO_FILES_IN_MAIN,
+            },
+            {
+                "mode": "list_files_in_bot_branch",
+                "name": "Overview of files in current working branch",
+                "description": GET_REPO_FILES_IN_BOT_BRANCH,
+            },
+            {
+                "mode": "list_files_from_directory",
+                "name": "Overview of files in current working branch from a specific path",  # noqa: E501
+                "description": GET_REPO_FILES_FROM_DIRECTORY,
+            },
+        ]
+        operations_filtered = [
+            operation
+            for operation in operations
+            if operation["mode"] in tools_to_include
         ]
         tools = [
             GitLabAction(
@@ -99,7 +161,7 @@ def from_gitlab_api_wrapper(
                 mode=action["mode"],
                 api_wrapper=gitlab_api_wrapper,
             )
-            for action in operations
+            for action in operations_filtered
         ]
         return cls(tools=tools)  # type: ignore[arg-type]
 

libs/community/langchain_community/tools/gitlab/prompt.py

@@ -68,3 +68,27 @@
 DELETE_FILE_PROMPT = """
 This tool is a wrapper for the GitLab API, useful when you need to delete a file in a GitLab repository. Simply pass in the full file path of the file you would like to delete. **IMPORTANT**: the path must not start with a slash
 """
+
+GET_REPO_FILES_IN_MAIN = """
+This tool will provide an overview of all existing files in the main branch of the GitLab repository repository. It will list the file names. No input parameters are required.
+"""
+
+GET_REPO_FILES_IN_BOT_BRANCH = """
+This tool will provide an overview of all files in your current working branch where you should implement changes. No input parameters are required.
+"""
+
+GET_REPO_FILES_FROM_DIRECTORY = """
+This tool will provide an overview of all files in your current working branch from a specific directory. **VERY IMPORTANT**: You must specify the path of the directory as a string input parameter.
+"""
+
+LIST_REPO_BRANCES = """
+This tool is a wrapper for the GitLab API, useful when you need to read the branches names in a GitLab repository. No input parameters are required.
+"""
+
+CREATE_REPO_BRANCH = """
+This tool will create a new branch in the repository. **VERY IMPORTANT**: You must specify the name of the new branch as a string input parameter.
+"""
+
+SET_ACTIVE_BRANCH = """
+This tool will set the active branch in the repository, similar to `git checkout <branch_name>` and `git switch -c <branch_name>`. **VERY IMPORTANT**: You must specify the name of the branch as a string input parameter.
+"""

libs/community/langchain_community/utilities/gitlab.py

@@ -325,6 +325,161 @@ def delete_file(self, file_path: str) -> str:
         except Exception as e:
             return "Unable to delete file due to error:\n" + str(e)
 
+    def list_files_in_main_branch(self) -> str:
+        """
+        Get the list of files in the main branch of the repository
+
+        Returns:
+            str: A plaintext report containing the list of files
+            in the repository in the main branch
+        """
+        if self.gitlab_base_branch is None:
+            return "No base branch set. Please set a base branch."
+        return self._list_files(self.gitlab_base_branch)
+
+    def list_files_in_bot_branch(self) -> str:
+        """
+        Get the list of files in the active branch of the repository
+
+        Returns:
+            str: A plaintext report containing the list of files
+            in the repository in the active branch
+        """
+        if self.gitlab_branch is None:
+            return "No active branch set. Please set a branch."
+        return self._list_files(self.gitlab_branch)
+
+    def list_files_from_directory(self, path: str) -> str:
+        """
+        Get the list of files in the active branch of the repository
+        from a specific directory
+
+        Returns:
+            str: A plaintext report containing the list of files
+            in the repository in the active branch from the specified directory
+        """
+        if self.gitlab_branch is None:
+            return "No active branch set. Please set a branch."
+        return self._list_files(
+            branch=self.gitlab_branch,
+            path=path,
+        )
+
+    def _list_files(self, branch: str, path: str = "") -> str:
+        try:
+            files = self._get_repository_files(
+                branch=branch,
+                path=path,
+            )
+            if files:
+                files_str = "\n".join(files)
+                return f"Found {len(files)} files in branch `{branch}`:\n{files_str}"
+            else:
+                return f"No files found in branch: `{branch}`"
+        except Exception as e:
+            return f"Error: {e}"
+
+    def _get_repository_files(self, branch: str, path: str = "") -> List[str]:
+        repo_contents = self.gitlab_repo_instance.repository_tree(ref=branch, path=path)
+
+        files: List[str] = []
+        for content in repo_contents:
+            if content["type"] == "tree":
+                files.extend(self._get_repository_files(branch, content["path"]))
+            else:
+                files.append(content["path"])
+
+        return files
+
+    def create_branch(self, proposed_branch_name: str) -> str:
+        """
+        Create a new branch in the repository and set it as the active branch
+
+        Parameters:
+            proposed_branch_name (str): The name of the new branch to be created
+        Returns:
+            str: A success or failure message
+        """
+        from gitlab import GitlabCreateError
+
+        max_attempts = 100
+        new_branch_name = proposed_branch_name
+        for i in range(max_attempts):
+            try:
+                response = self.gitlab_repo_instance.branches.create(
+                    {
+                        "branch": new_branch_name,
+                        "ref": self.gitlab_branch,
+                    }
+                )
+
+                self.gitlab_branch = response.name
+                return (
+                    f"Branch '{response.name}' "
+                    "created successfully, and set as current active branch."
+                )
+
+            except GitlabCreateError as e:
+                if (
+                    e.response_code == 400
+                    and "Branch already exists" in e.error_message
+                ):
+                    i += 1
+                    new_branch_name = f"{proposed_branch_name}_v{i}"
+                else:
+                    # Handle any other exceptions
+                    print(f"Failed to create branch. Error: {e}")  # noqa: T201
+                    raise Exception(
+                        "Unable to create branch name from proposed_branch_name: "
+                        f"{proposed_branch_name}"
+                    )
+
+        return (
+            f"Unable to create branch. At least {max_attempts} branches exist "
+            f"with named derived from "
+            f"proposed_branch_name: `{proposed_branch_name}`"
+        )
+
+    def list_branches_in_repo(self) -> str:
+        """
+        Get the list of branches in the repository
+
+        Returns:
+            str: A plaintext report containing the number of branches
+            and each branch name
+        """
+        branches = [
+            branch.name for branch in self.gitlab_repo_instance.branches.list(all=True)
+        ]
+        if branches:
+            branches_str = "\n".join(branches)
+            return (
+                f"Found {str(len(branches))} branches in the repository:"
+                f"\n{branches_str}"
+            )
+        return "No branches found in the repository"
+
+    def set_active_branch(self, branch_name: str) -> str:
+        """Equivalent to `git checkout branch_name` for this Agent.
+        Clones formatting from Gitlab.
+
+        Returns an Error (as a string) if branch doesn't exist.
+        """
+        curr_branches = [
+            branch.name
+            for branch in self.gitlab_repo_instance.branches.list(
+                all=True,
+            )
+        ]
+        if branch_name in curr_branches:
+            self.gitlab_branch = branch_name
+            return f"Switched to branch `{branch_name}`"
+        else:
+            return (
+                f"Error {branch_name} does not exist,"
+                f"in repo with current branches: {str(curr_branches)}"
+            )
+
     def run(self, mode: str, query: str) -> str:
         if mode == "get_issues":
             return self.get_issues()
@@ -342,5 +497,17 @@ def run(self, mode: str, query: str) -> str:
             return self.update_file(query)
         elif mode == "delete_file":
             return self.delete_file(query)
+        elif mode == "create_branch":
+            return self.create_branch(query)
+        elif mode == "list_branches_in_repo":
+            return self.list_branches_in_repo()
+        elif mode == "set_active_branch":
+            return self.set_active_branch(query)
+        elif mode == "list_files_in_main_branch":
+            return self.list_files_in_main_branch()
+        elif mode == "list_files_in_bot_branch":
+            return self.list_files_in_bot_branch()
+        elif mode == "list_files_from_directory":
+            return self.list_files_from_directory(query)
         else:
             raise ValueError("Invalid mode" + mode)