From 3345027932df3c941d6367dd46ff7e7b1d3907fc Mon Sep 17 00:00:00 2001
From: tawsifkamal <tawsifkamal123@gmail.com>
Date: Thu, 20 Mar 2025 12:30:00 -0700
Subject: [PATCH] done

---
 src/codegen/agents/chat_agent.py          |  2 +-
 src/codegen/agents/code_agent.py          |  2 +-
 src/codegen/extensions/langchain/agent.py | 81 ++++++++++-------------
 3 files changed, 38 insertions(+), 47 deletions(-)

diff --git a/src/codegen/agents/chat_agent.py b/src/codegen/agents/chat_agent.py
index 24ecada26..a64d36c95 100644
--- a/src/codegen/agents/chat_agent.py
+++ b/src/codegen/agents/chat_agent.py
@@ -29,7 +29,7 @@ def __init__(self, codebase: "Codebase", model_provider: str = "anthropic", mode
                 - max_tokens: Maximum number of tokens to generate
         """
         self.codebase = codebase
-        self.agent = create_chat_agent(self.codebase, model_provider=model_provider, model_name=model_name, memory=memory, additional_tools=tools, **kwargs)
+        self.agent = create_chat_agent(self.codebase, model_provider=model_provider, model_name=model_name, memory=memory, tools=tools, **kwargs)
 
     def run(self, prompt: str, thread_id: Optional[str] = None) -> str:
         """Run the agent with a prompt.
diff --git a/src/codegen/agents/code_agent.py b/src/codegen/agents/code_agent.py
index 693c0cd44..73709f667 100644
--- a/src/codegen/agents/code_agent.py
+++ b/src/codegen/agents/code_agent.py
@@ -70,7 +70,7 @@ def __init__(
             model_provider=model_provider,
             model_name=model_name,
             memory=memory,
-            additional_tools=tools,
+            tools=tools,
             config=agent_config,
             **kwargs,
         )
diff --git a/src/codegen/extensions/langchain/agent.py b/src/codegen/extensions/langchain/agent.py
index 42493c1fc..111f182e2 100644
--- a/src/codegen/extensions/langchain/agent.py
+++ b/src/codegen/extensions/langchain/agent.py
@@ -40,7 +40,7 @@ def create_codebase_agent(
     system_message: SystemMessage = SystemMessage(REASONER_SYSTEM_MESSAGE),
     memory: bool = True,
     debug: bool = False,
-    additional_tools: list[BaseTool] | None = None,
+    tools: list[BaseTool] | None = None,
     config: AgentConfig | None = None,
     **kwargs,
 ) -> CompiledGraph:
@@ -63,35 +63,28 @@ def create_codebase_agent(
     """
     llm = LLM(model_provider=model_provider, model_name=model_name, **kwargs)
 
-    # Get all codebase tools
-    tools = [
-        ViewFileTool(codebase),
-        ListDirectoryTool(codebase),
-        RipGrepTool(codebase),
-        # EditFileTool(codebase),
-        CreateFileTool(codebase),
-        DeleteFileTool(codebase),
-        RenameFileTool(codebase),
-        # MoveSymbolTool(codebase),
-        # RevealSymbolTool(codebase),
-        # SemanticEditTool(codebase),
-        ReplacementEditTool(codebase),
-        RelaceEditTool(codebase),
-        ReflectionTool(codebase),
-        SearchFilesByNameTool(codebase),
-        GlobalReplacementEditTool(codebase),
-        # SemanticSearchTool(codebase),
-        # =====[ Github Integration ]=====
-        # Enable Github integration
-        # GithubCreatePRTool(codebase),
-        # GithubViewPRTool(codebase),
-        # GithubCreatePRCommentTool(codebase),
-        # GithubCreatePRReviewCommentTool(codebase),
-    ]
-
-    # Add additional tools if provided
-    if additional_tools:
-        tools.extend(additional_tools)
+    if not tools:
+        tools = [
+            ViewFileTool(codebase),
+            ListDirectoryTool(codebase),
+            RipGrepTool(codebase),
+            # EditFileTool(codebase),
+            CreateFileTool(codebase),
+            DeleteFileTool(codebase),
+            RenameFileTool(codebase),
+            ReplacementEditTool(codebase),
+            RelaceEditTool(codebase),
+            ReflectionTool(codebase),
+            SearchFilesByNameTool(codebase),
+            GlobalReplacementEditTool(codebase),
+            # SemanticSearchTool(codebase),
+            # =====[ Github Integration ]=====
+            # Enable Github integration
+            # GithubCreatePRTool(codebase),
+            # GithubViewPRTool(codebase),
+            # GithubCreatePRCommentTool(codebase),
+            # GithubCreatePRReviewCommentTool(codebase),
+        ]
 
     memory = MemorySaver() if memory else None
 
@@ -105,7 +98,7 @@ def create_chat_agent(
     system_message: SystemMessage = SystemMessage(REASONER_SYSTEM_MESSAGE),
     memory: bool = True,
     debug: bool = False,
-    additional_tools: list[BaseTool] | None = None,
+    tools: list[BaseTool] | None = None,
     config: dict[str, Any] | None = None,  # over here you can pass in the max length of the number of messages
     **kwargs,
 ) -> CompiledGraph:
@@ -128,20 +121,18 @@ def create_chat_agent(
     """
     llm = LLM(model_provider=model_provider, model_name=model_name, **kwargs)
 
-    tools = [
-        ViewFileTool(codebase),
-        ListDirectoryTool(codebase),
-        RipGrepTool(codebase),
-        CreateFileTool(codebase),
-        DeleteFileTool(codebase),
-        RenameFileTool(codebase),
-        MoveSymbolTool(codebase),
-        RevealSymbolTool(codebase),
-        RelaceEditTool(codebase),
-    ]
-
-    if additional_tools:
-        tools.extend(additional_tools)
+    if not tools:
+        tools = [
+            ViewFileTool(codebase),
+            ListDirectoryTool(codebase),
+            RipGrepTool(codebase),
+            CreateFileTool(codebase),
+            DeleteFileTool(codebase),
+            RenameFileTool(codebase),
+            MoveSymbolTool(codebase),
+            RevealSymbolTool(codebase),
+            RelaceEditTool(codebase),
+        ]
 
     memory = MemorySaver() if memory else None