feat: adds extra tools to codebase agent (#594)

# Motivation

<!-- Why is this change necessary? -->

# Content

<!-- Please include a summary of the change -->

# Testing

<!-- How was the change tested? -->

# Please check the following before marking your PR as ready for review

- [x ] I have added tests for my changes
- [ x] I have updated the documentation or added new documentation as
needed

---------

Co-authored-by: kopekC <[email protected]>

Author: kopekC

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

@@ -1318,6 +1318,10 @@ def get_modified_symbols_in_pr(self, pr_id: int) -> tuple[list[Symbol], str]:
         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