chore(agi): CG-10772 add create PR tool

src/codegen/extensions/langchain/tools.py

@@ -33,9 +33,9 @@
 class ViewFileTool(BaseTool):
     """Tool for viewing file contents and metadata."""
 
     name: ClassVar[str] = "view_file"
     description: ClassVar[str] = "View the contents and metadata of a file in the codebase"
     args_schema: ClassVar[type[BaseModel]] = ViewFileInput
     codebase: Codebase = Field(exclude=True)
 
     def __init__(self, codebase: Codebase) -> None:
@@ -56,9 +56,9 @@
 class ListDirectoryTool(BaseTool):
     """Tool for listing directory contents."""
 
     name: ClassVar[str] = "list_directory"
     description: ClassVar[str] = "List contents of a directory in the codebase"
     args_schema: ClassVar[type[BaseModel]] = ListDirectoryInput
     codebase: Codebase = Field(exclude=True)
 
     def __init__(self, codebase: Codebase) -> None:
@@ -79,9 +79,9 @@
 class SearchTool(BaseTool):
     """Tool for searching the codebase."""
 
     name: ClassVar[str] = "search"
     description: ClassVar[str] = "Search the codebase using text search"
     args_schema: ClassVar[type[BaseModel]] = SearchInput
     codebase: Codebase = Field(exclude=True)
 
     def __init__(self, codebase: Codebase) -> None:
@@ -102,7 +102,7 @@
 class EditFileTool(BaseTool):
     """Tool for editing files."""
 
     name: ClassVar[str] = "edit_file"
     description: ClassVar[str] = "Edit a file by replacing its entire content"
     args_schema: ClassVar[type[BaseModel]] = EditFileInput
     codebase: Codebase = Field(exclude=True)
@@ -336,6 +336,29 @@
         return json.dumps(result, indent=2)
 
 
+class CreatePRInput(BaseModel):
+    """Input for creating a PR"""
+
+    title: str = Field(..., description="The title of the PR")
+    body: str = Field(..., description="The body of the PR")
+
+
+class CreatePRTool(BaseTool):
+    """Tool for creating a PR."""
+
+    name: ClassVar[str] = "create_pr"
+    description: ClassVar[str] = "Create a PR for the current branch"
+    args_schema: ClassVar[type[BaseModel]] = CreatePRInput
+    codebase: Codebase = Field(exclude=True)
+
+    def __init__(self, codebase: Codebase) -> None:
+        super().__init__(codebase=codebase)
+
+    def _run(self, title: str, body: str) -> str:
+        pr = self.codebase.create_pr(title=title, body=body)
+        return pr.html_url
+
+
 def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]:
     """Get all workspace tools initialized with a codebase.
 
@@ -345,26 +368,15 @@
     Returns:
         List of initialized Langchain tools
     """
-    from .tools import (
-        CommitTool,
-        CreateFileTool,
-        DeleteFileTool,
-        EditFileTool,
-        ListDirectoryTool,
-        RevealSymbolTool,
-        SearchTool,
-        SemanticEditTool,
-        ViewFileTool,
-    )
-
     return [
-        ViewFileTool(codebase),
-        ListDirectoryTool(codebase),
-        SearchTool(codebase),
-        EditFileTool(codebase),
+        CommitTool(codebase),
         CreateFileTool(codebase),
+        CreatePRTool(codebase),
         DeleteFileTool(codebase),
-        CommitTool(codebase),
+        EditFileTool(codebase),
+        ListDirectoryTool(codebase),
         RevealSymbolTool(codebase),
+        SearchTool(codebase),
         SemanticEditTool(codebase),
+        ViewFileTool(codebase),
     ]