feat: adds extra tools to codebase agent

src/codegen/agents/code_agent.py

@@ -1,16 +1,18 @@
 from typing import Optional
 from uuid import uuid4
 
+from langchain.tools import BaseTool
+
 from codegen.extensions.langchain.agent import create_codebase_agent
 from codegen.sdk.core.codebase import Codebase
 
 
 class CodeAgent:
     """Agent for interacting with a codebase."""
 
-    def __init__(self, codebase: Codebase):
+    def __init__(self, codebase: Codebase, tools: Optional[list[BaseTool]] = None):
         self.codebase = codebase
-        self.agent = create_codebase_agent(self.codebase)
+        self.agent = create_codebase_agent(self.codebase, additional_tools=tools)
 
     def run(self, prompt: str, session_id: Optional[str] = None) -> str:
         if session_id is None:

src/codegen/extensions/langchain/agent.py

@@ -1,5 +1,7 @@
 """Demo implementation of an agent with Codegen tools."""
 
+from typing import Optional
+
 from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
 from langchain.hub import pull
@@ -34,6 +36,7 @@ def create_codebase_agent(
     temperature: float = 0,
     verbose: bool = True,
     chat_history: list[BaseMessage] = [],
+    additional_tools: Optional[list[BaseTool]] = None,
 ) -> RunnableWithMessageHistory:
     """Create an agent with all codebase tools.
 
@@ -42,6 +45,8 @@ def create_codebase_agent(
         model_name: Name of the model to use (default: gpt-4)
         temperature: Model temperature (default: 0)
         verbose: Whether to print agent's thought process (default: True)
+        chat_history: Optional list of messages to initialize chat history with
+        additional_tools: Optional list of additional tools to provide to the agent
 
     Returns:
         Initialized agent with message history
@@ -78,6 +83,9 @@ def create_codebase_agent(
         # GithubCreatePRCommentTool(codebase),
         # GithubCreatePRReviewCommentTool(codebase),
     ]
+    # Add additional tools if provided
+    if additional_tools:
+        tools.extend(additional_tools)
 
     prompt = ChatPromptTemplate.from_messages(
         [

src/codegen/sdk/core/codebase.py

@@ -104,8 +104,8 @@
 TExport = TypeVar("TExport", bound="Export", default=Export)
 TSGlobalVar = TypeVar("TSGlobalVar", bound="Assignment", default=Assignment)
 PyGlobalVar = TypeVar("PyGlobalVar", bound="Assignment", default=Assignment)
 TSDirectory = Directory[TSFile, TSSymbol, TSImportStatement, TSGlobalVar, TSClass, TSFunction, TSImport]
 PyDirectory = Directory[PyFile, PySymbol, PyImportStatement, PyGlobalVar, PyClass, PyFunction, PyImport]
 
 
 @apidoc
@@ -182,13 +182,13 @@
             main_project = ProjectConfig.from_path(repo_path, programming_language=ProgrammingLanguage(language.upper()) if language else None)
             projects = [main_project]
         else:
             main_project = projects[0]
 
         # Initialize codebase
         self._op = main_project.repo_operator
         self.viz = VisualizationManager(op=self._op)
         self.repo_path = Path(self._op.repo_path)
         self.ctx = CodebaseContext(projects, config=config, secrets=secrets, io=io, progress=progress)
         self.console = Console(record=True, soft_wrap=True)
 
     @noapidoc
@@ -204,7 +204,7 @@
         yield "nodes", len(self.ctx.nodes)
         yield "edges", len(self.ctx.edges)
 
     __rich_repr__.angular = ANGULAR_STYLE
 
     @property
     @deprecated("Please do not use the local repo operator directly")
@@ -246,8 +246,8 @@
 
     @noapidoc
     def _symbols(self, symbol_type: SymbolType | None = None) -> list[TSymbol | TClass | TFunction | TGlobalVar]:
         matches: list[Symbol] = self.ctx.get_nodes(NodeType.SYMBOL)
         return [x for x in matches if x.is_top_level and (symbol_type is None or x.symbol_type == symbol_type)]
 
     # =====[ Node Types ]=====
     @overload
@@ -256,7 +256,7 @@
     def files(self, *, extensions: Literal["*"]) -> list[File]: ...
     @overload
     def files(self, *, extensions: None = ...) -> list[TSourceFile]: ...
     @proxy_property
     def files(self, *, extensions: list[str] | Literal["*"] | None = None) -> list[TSourceFile] | list[File]:
         """A list property that returns all files in the codebase.
 
@@ -291,13 +291,13 @@
         return sort_editables(files, alphabetical=True, dedupe=False)
 
     @cached_property
     def codeowners(self) -> list["CodeOwner[TSourceFile]"]:
         """List all CodeOnwers in the codebase.
 
         Returns:
             list[CodeOwners]: A list of CodeOwners objects in the codebase.
         """
         if self.G.codeowners_parser is None:
             return []
         return CodeOwner.from_parser(self.G.codeowners_parser, lambda *args, **kwargs: self.files(*args, **kwargs))
 
@@ -1318,6 +1318,10 @@
         patch = cg_pr.get_pr_diff()
         return cg_pr.modified_symbols, patch
 
+    def create_pr_comment(self, pr_number: int, body: str) -> None:
+        """Create a comment on a pull request"""
+        return self._op.create_pr_comment(pr_number, body)
+
 
 # The last 2 lines of code are added to the runner. See codegen-backend/cli/generate/utils.py
 # Type Aliases