Supercoder Improve tool addition (#755)

superagi/resource_manager/file_manager.py

@@ -1,21 +1,17 @@
 import csv
-
 from sqlalchemy.orm import Session
-
+import os
 from superagi.helper.resource_helper import ResourceHelper
 from superagi.helper.s3_helper import S3Helper
 from superagi.lib.logger import logger
 from superagi.models.agent import Agent
 from superagi.models.agent_execution import AgentExecution
 from superagi.types.storage_types import StorageType
-
-
 class FileManager:
     def __init__(self, session: Session, agent_id: int = None, agent_execution_id: int = None):
         self.session = session
         self.agent_id = agent_id
         self.agent_execution_id = agent_execution_id
-
     def write_binary_file(self, file_name: str, data):
         if self.agent_id is not None:
             final_path = ResourceHelper.get_agent_write_resource_path(file_name,
@@ -35,7 +31,6 @@ def write_binary_file(self, file_name: str, data):
             return f"Binary {file_name} saved successfully"
         except Exception as err:
             return f"Error write_binary_file: {err}"
-
     def write_to_s3(self, file_name, final_path):
         with open(final_path, 'rb') as img:
             resource = ResourceHelper.make_written_file_resource(file_name=file_name,
@@ -68,7 +63,6 @@ def write_file(self, file_name: str, content):
             return f"{file_name} - File written successfully"
         except Exception as err:
             return f"Error write_file: {err}"
-
     def write_csv_file(self, file_name: str, csv_data):
         if self.agent_id is not None:
             final_path = ResourceHelper.get_agent_write_resource_path(file_name,
@@ -79,7 +73,6 @@ def write_csv_file(self, file_name: str, csv_data):
                                                                                                    self.agent_execution_id))
         else:
             final_path = ResourceHelper.get_resource_path(file_name)
-
         try:
             with open(final_path, mode="w") as file:
                 writer = csv.writer(file, lineterminator="\n")
@@ -90,10 +83,42 @@ def write_csv_file(self, file_name: str, csv_data):
             return f"{file_name} - File written successfully"
         except Exception as err:
             return f"Error write_csv_file: {err}"
+
 
     def get_agent_resource_path(self, file_name: str):
         return ResourceHelper.get_agent_write_resource_path(file_name, agent=Agent.get_agent_from_id(self.session,
                                                                                                      self.agent_id),
                                                             agent_execution=AgentExecution
                                                             .get_agent_execution_from_id(self.session,
                                                                                          self.agent_execution_id))
+    def read_file(self, file_name: str):
+        if self.agent_id is not None:
+            final_path = self.get_agent_resource_path(file_name)
+        else:
+            final_path = ResourceHelper.get_resource_path(file_name)
+
+        try:
+            with open(final_path, mode="r") as file:
+                content = file.read()
+            logger.info(f"{file_name} - File read successfully")
+            return content
+        except Exception as err:
+            return f"Error while reading file {file_name}: {err}"
+    def get_files(self):
+        """
+        Gets all file names generated by the CodingTool.
+        Returns:
+            A list of file names.
+        """
+
+        if self.agent_id is not None:
+            final_path = self.get_agent_resource_path("")
+        else:
+            final_path = ResourceHelper.get_resource_path("")
+        try:
+            # List all files in the directory
+            files = os.listdir(final_path)
+        except Exception as err:
+            logger.error(f"Error while accessing files in {final_path}: {err}")
+            files = []
+        return files

superagi/tools/code/coding_toolkit.py

@@ -5,14 +5,15 @@
 from superagi.tools.code.write_code import CodingTool
 from superagi.tools.code.write_spec import WriteSpecTool
 from superagi.tools.code.write_test import WriteTestTool
+from superagi.tools.code.improve_code import ImproveCodeTool
 
 
 class CodingToolkit(BaseToolkit, ABC):
     name: str = "CodingToolkit"
     description: str = "Coding Tool kit contains all tools related to coding tasks"
 
     def get_tools(self) -> List[BaseTool]:
-        return [CodingTool(), WriteSpecTool(), WriteTestTool()]
+        return [CodingTool(), WriteSpecTool(), WriteTestTool(), ImproveCodeTool()]
 
     def get_env_keys(self) -> List[str]:
         return []

superagi/tools/code/improve_code.py

@@ -0,0 +1,99 @@
+import re
+from typing import Type, Optional, List
+
+from pydantic import BaseModel, Field
+
+from superagi.agent.agent_prompt_builder import AgentPromptBuilder
+from superagi.helper.prompt_reader import PromptReader
+from superagi.helper.token_counter import TokenCounter
+from superagi.lib.logger import logger
+from superagi.llms.base_llm import BaseLlm
+from superagi.resource_manager.file_manager import FileManager
+from superagi.tools.base_tool import BaseTool
+from superagi.tools.tool_response_query_manager import ToolResponseQueryManager
+
+
+class ImproveCodeSchema(BaseModel):
+    pass
+
+
+class ImproveCodeTool(BaseTool):
+    """
+    Used to improve the already generated code by reading the code from the files
+
+    Attributes:
+        llm: LLM used for code generation.
+        name : The name of the tool.
+        description : The description of the tool.
+        resource_manager: Manages the file resources.
+    """
+    llm: Optional[BaseLlm] = None
+    agent_id: int = None
+    name = "ImproveCodeTool"
+    description = (
+        "This tool improves the generated code."
+    )
+    args_schema: Type[ImproveCodeSchema] = ImproveCodeSchema
+    resource_manager: Optional[FileManager] = None
+    tool_response_manager: Optional[ToolResponseQueryManager] = None
+    goals: List[str] = []
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def _execute(self) -> str:
+        """
+        Execute the improve code tool.
+
+        Returns:
+            Improved code or error message.
+        """
+        # Get all file names that the CodingTool has written
+        file_names = self.resource_manager.get_files()
+        logger.info(file_names)
+        # Loop through each file
+        for file_name in file_names:
+            if '.txt' not in file_name and '.sh' not in file_name and '.json' not in file_name:
+                # Read the file content
+                content = self.resource_manager.read_file(file_name)
+
+                # Generate a prompt from improve_code.txt
+                prompt = PromptReader.read_tools_prompt(__file__, "improve_code.txt")
+
+                # Combine the hint from the file, goals, and content
+                prompt = prompt.replace("{goals}", AgentPromptBuilder.add_list_items_to_string(self.goals))
+                prompt = prompt.replace("{content}", content)
+
+                # Add the file content to the chat completion prompt
+                prompt = prompt + "\nOriginal Code:\n```\n" + content + "\n```"
+
+
+
+                # Use LLM to generate improved code
+                result = self.llm.chat_completion([{'role': 'system', 'content': prompt}])
+
+                # Extract the response first
+                response = result.get('response')
+                if not response: 
+                    logger.info("RESPONSE NOT AVAILABLE")
+
+                # Now extract the choices from response
+                choices = response.get('choices')
+                if not choices: 
+                    logger.info("CHOICES NOT AVAILABLE")
+
+                # Now you can safely extract the message content
+                improved_content = choices[0]["message"]["content"]
+                # improved_content = result["messages"][0]["content"]
+                parsed_content = re.findall("```(?:\w*\n)?(.*?)```", improved_content, re.DOTALL)
+                parsed_content_code = "\n".join(parsed_content)
+
+                # Rewrite the file with the improved content
+                save_result = self.resource_manager.write_file(file_name, parsed_content_code)
+
+                if save_result.startswith("Error"):
+                    return save_result
+            else:
+                continue
+
+        return f"All codes improved and saved successfully in: " + " ".join(file_names)

superagi/tools/code/prompts/improve_code.txt

@@ -0,0 +1,26 @@
+You are a super smart developer. You have been tasked with fixing and filling the function and classes where only the description of code is written without the actual code . There might be placeholders in the code you have to fill in.
+You provide fully functioning, well formatted code with few comments, that works and has no bugs.
+If the code is already correct and doesn't need change, just return the same code
+However, make sure that you only return the improved code, without any additional content.
+
+
+Please structure the improved code as follows:
+
+```
+CODE
+```
+
+Please return the full new code in same format as the original code
+Don't write any explanation or description in your response other than the actual code
+
+Your high-level goal is:
+{goals}
+
+The content of the file you need to improve is:
+{content}
+
+Only return the code and not any other line
+
+To start, first analyze the existing code. Check for any function with missing logic inside it and fill the function. 
+Make sure, that not a single function is empty or contains just comments, there should be function logic inside it
+Return fully completed functions by filling the placeholders

superagi/tools/code/prompts/write_code.txt

@@ -1,4 +1,5 @@
 You are a super smart developer who practices good Development for writing code according to a specification.
+Please note that the code should be fully functional. There should be no placeholder in functions or classes in any file.
 
 Your high-level goal is:
 {goals}
@@ -23,7 +24,6 @@ Each file must strictly follow a markdown code block format, where the following
 ```
 
 You will start with the "entrypoint" file, then go to the ones that are imported by that file, and so on.
-Please note that the code should be fully functional. No placeholders.
 
 Follow a language and framework appropriate best practice file naming convention.
 Make sure that files contain all imports, types etc. Make sure that code in different files are compatible with each other.

superagi/tools/code/write_code.py

@@ -40,7 +40,7 @@ class CodingTool(BaseTool):
         "Make sure that every detail of the architecture is, in the end, implemented as code. "
         "Think step by step and reason yourself to the right decisions to make sure we get it right. "
         "You will first lay out the names of the core classes, functions, methods that will be necessary, "
-        "as well as a quick comment on their purpose. Then you will output the content of each file including ALL code."
+        "as well as a quick comment on their purpose. Then you will output the content of each file including each function and class and ALL code."
     )
     args_schema: Type[CodingSchema] = CodingSchema
     goals: List[str] = []
@@ -84,6 +84,8 @@ def _execute(self, code_description: str) -> str:
         for match in matches:
             # Get the filename
             file_name = re.sub(r'[<>"|?*]', "", match.group(1))
+            if not file_name[0].isalnum():
+                file_name = file_name[1:-1]
 
             # Get the code
             code = match.group(2)

tests/unit_tests/tools/code/test_improve_code.py

@@ -0,0 +1,52 @@
+import pytest
+from unittest.mock import Mock, MagicMock
+from superagi.tools.code.improve_code import ImproveCodeTool
+
+@pytest.fixture
+def mock_improve_code_tool():
+    improve_code_tool = ImproveCodeTool()
+    improve_code_tool.resource_manager = Mock()
+    improve_code_tool.llm = Mock()
+    return improve_code_tool
+
+def test_execute(mock_improve_code_tool):
+    mock_improve_code_tool.resource_manager.get_files.return_value = ['test1', 'test2']
+    mock_improve_code_tool.resource_manager.read_file.return_value = "test file content"
+    mock_improve_code_tool.llm.chat_completion.return_value = {
+        "response":
+            {
+                "choices":
+                    [
+                        {
+                            "message":
+                                {
+                                    "content": "```\nimproved code\n```"
+                                }
+                        }
+                    ]
+            }
+    }
+    mock_improve_code_tool.resource_manager.write_file.return_value = "file saved successfully"
+
+    assert mock_improve_code_tool._execute() == "All codes improved and saved successfully in: test1 test2"
+
+def test_execute_with_error(mock_improve_code_tool):
+    mock_improve_code_tool.resource_manager.get_files.return_value = ['test1']
+    mock_improve_code_tool.resource_manager.read_file.return_value = "test file content"
+    mock_improve_code_tool.llm.chat_completion.return_value = {
+        "response":
+            {
+                "choices":
+                    [
+                        {
+                            "message":
+                                {
+                                    "content": "```\nimproved code\n```"
+                                }
+                        }
+                    ]
+            }
+    }
+    mock_improve_code_tool.resource_manager.write_file.return_value = "Error: Could not save file"
+
+    assert mock_improve_code_tool._execute() == "Error: Could not save file"