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

Author: christinewangcw

src/codegen/extensions/langchain/tools.py

@@ -336,6 +336,29 @@ def _run(self, query: str, k: int = 5, preview_length: int = 200) -> str:
         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 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]:
     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),
     ]