Merge pull request #446 from TransformerOptimus/supercoder

Supercoder

superagi/helper/resource_helper.py

@@ -20,7 +20,7 @@ def make_written_file_resource(file_name: str, agent_id: int, channel: str):
         Returns:
             Resource: The Resource object.
         """
-        path = get_config("RESOURCES_OUTPUT_ROOT_DIR")
+        path = ResourceHelper.get_root_dir()
         storage_type = get_config("STORAGE_TYPE")
         file_extension = os.path.splitext(file_name)[1][1:]
 
@@ -31,17 +31,21 @@ def make_written_file_resource(file_name: str, agent_id: int, channel: str):
         else:
             file_type = "application/misc"
 
-        final_path = ResourceHelper.get_resource_path(file_name)
+        if agent_id is not None:
+            final_path = ResourceHelper.get_agent_resource_path(file_name, agent_id)
+            path = path + str(agent_id) + "/"
+        else:
+            final_path = ResourceHelper.get_resource_path(file_name)
         file_size = os.path.getsize(final_path)
 
         if storage_type == "S3":
             file_name_parts = file_name.split('.')
             file_name = file_name_parts[0] + '_' + str(datetime.datetime.now()).replace(' ', '') \
                 .replace('.', '').replace(':', '') + '.' + file_name_parts[1]
-            path = 'input' if (channel == "INPUT") else 'output'
+            path = 'input/' if (channel == "INPUT") else 'output/'
 
-        logger.info(path + "/" + file_name)
-        resource = Resource(name=file_name, path=path + "/" + file_name, storage_type=storage_type, size=file_size,
+        logger.info(final_path)
+        resource = Resource(name=file_name, path=path + file_name, storage_type=storage_type, size=file_size,
                             type=file_type,
                             channel="OUTPUT",
                             agent_id=agent_id)
@@ -54,12 +58,32 @@ def get_resource_path(file_name: str):
         Args:
             file_name (str): The name of the file.
         """
+        return ResourceHelper.get_root_dir() + file_name
+
+    @staticmethod
+    def get_root_dir():
+        """Get root dir of the resource.
+        """
         root_dir = get_config('RESOURCES_OUTPUT_ROOT_DIR')
 
         if root_dir is not None:
             root_dir = root_dir if root_dir.startswith("/") else os.getcwd() + "/" + root_dir
             root_dir = root_dir if root_dir.endswith("/") else root_dir + "/"
-            final_path = root_dir + file_name
         else:
-            final_path = os.getcwd() + "/" + file_name
+            root_dir = os.getcwd() + "/"
+        return root_dir
+
+    @staticmethod
+    def get_agent_resource_path(file_name: str, agent_id: int):
+        """Get agent resource path
+
+        Args:
+            file_name (str): The name of the file.
+        """
+        root_dir = ResourceHelper.get_root_dir()
+        if agent_id is not None:
+            directory = os.path.dirname(root_dir + str(agent_id) + "/")
+            os.makedirs(directory, exist_ok=True)
+            root_dir = root_dir + str(agent_id) + "/"
+        final_path = root_dir + file_name
         return final_path

superagi/jobs/agent_executor.py

@@ -19,6 +19,7 @@
 from superagi.models.tool import Tool
 from superagi.resource_manager.manager import ResourceManager
 from superagi.tools.thinking.tools import ThinkingTool
+from superagi.tools.tool_response_query_manager import ToolResponseQueryManager
 from superagi.vector_store.embedding.openai import OpenAiEmbedding
 from superagi.vector_store.vector_factory import VectorFactory
 from superagi.helper.encyption_helper import decrypt_data
@@ -234,7 +235,10 @@ def set_default_params_tools(self, tools, parsed_config, agent_id, model_api_key
             if hasattr(tool, 'agent_id'):
                 tool.agent_id = agent_id
             if hasattr(tool, 'resource_manager'):
-                tool.resource_manager = ResourceManager(session=session)
+                tool.resource_manager = ResourceManager(session=session, agent_id=agent_id)
+            if hasattr(tool, 'tool_response_manager'):
+                tool.tool_response_manager = ToolResponseQueryManager(session=session, agent_execution_id=parsed_config[
+                    "agent_execution_id"])
 
             new_tools.append(tool)
         return tools

superagi/models/agent_execution_feed.py

@@ -1,4 +1,5 @@
 from sqlalchemy import Column, Integer, Text, String
+from sqlalchemy.orm import Session
 
 from superagi.models.base_model import DBBaseModel
 
@@ -36,3 +37,16 @@ def __repr__(self):
         return f"AgentExecutionFeed(id={self.id}, " \
                f"agent_execution_id={self.agent_execution_id}, " \
                f"feed='{self.feed}', role='{self.role}', extra_info={self.extra_info})"
+
+    @classmethod
+    def get_last_tool_response(cls, session: Session, agent_execution_id: int, tool_name: str = None):
+        agent_execution_feeds = session.query(AgentExecutionFeed).filter(
+            AgentExecutionFeed.agent_execution_id == agent_execution_id,
+            AgentExecutionFeed.role == "system").order_by(AgentExecutionFeed.created_at.desc()).all()
+
+        for agent_execution_feed in agent_execution_feeds:
+            if tool_name and not agent_execution_feed.feed.startswith("Tool " + tool_name):
+                continue
+            if agent_execution_feed.feed.startswith("Tool"):
+                return agent_execution_feed.feed
+        return ""

superagi/resource_manager/manager.py

@@ -3,21 +3,30 @@
 from superagi.helper.resource_helper import ResourceHelper
 from superagi.helper.s3_helper import S3Helper
 from superagi.lib.logger import logger
+import os
 
 
 class ResourceManager:
-    def __init__(self, session: Session):
+    def __init__(self, session: Session, agent_id: int = None):
         self.session = session
+        self.agent_id = agent_id
 
     def write_binary_file(self, file_name: str, data):
-        final_path = ResourceHelper.get_resource_path(file_name)
-
+        if self.agent_id is not None:
+            final_path = ResourceHelper.get_agent_resource_path(file_name, self.agent_id)
+        else:
+            final_path = ResourceHelper.get_resource_path(file_name)
+
+        # if self.agent_id is not None:
+        #     directory = os.path.dirname(final_path + "/" + str(self.agent_id) + "/")
+        #     os.makedirs(directory, exist_ok=True)
         try:
             with open(final_path, mode="wb") as img:
                 img.write(data)
                 img.close()
             self.write_to_s3(file_name, final_path)
             logger.info(f"Binary {file_name} saved successfully")
+            return f"Binary {file_name} saved successfully"
         except Exception as err:
             return f"Error: {err}"
 
@@ -34,13 +43,17 @@ def write_to_s3(self, file_name, final_path):
                     s3_helper.upload_file(img, path=resource.path)
 
     def write_file(self, file_name: str, content):
-        final_path = ResourceHelper.get_resource_path(file_name)
+        if self.agent_id is not None:
+            final_path = ResourceHelper.get_agent_resource_path(file_name, self.agent_id)
+        else:
+            final_path = ResourceHelper.get_resource_path(file_name)
 
         try:
             with open(final_path, mode="w") as file:
                 file.write(content)
                 file.close()
             self.write_to_s3(file_name, final_path)
             logger.info(f"{file_name} saved successfully")
+            return f"{file_name} saved successfully"
         except Exception as err:
-            return f"Error: {err}"
+            return f"Error: {err}"

superagi/tools/code/write_code.py

@@ -0,0 +1,143 @@
+import re
+from typing import Type, Optional, List
+
+from pydantic import BaseModel, Field
+
+from superagi.agent.agent_prompt_builder import AgentPromptBuilder
+from superagi.helper.token_counter import TokenCounter
+from superagi.lib.logger import logger
+from superagi.llms.base_llm import BaseLlm
+from superagi.resource_manager.manager import ResourceManager
+from superagi.tools.base_tool import BaseTool
+from superagi.tools.tool_response_query_manager import ToolResponseQueryManager
+
+
+class CodingSchema(BaseModel):
+    code_description: str = Field(
+        ...,
+        description="Description of the coding task",
+    )
+
+class CodingTool(BaseTool):
+    """
+    Used to generate code.
+
+    Attributes:
+        llm: LLM used for code generation.
+        name : The name of tool.
+        description : The description of tool.
+        args_schema : The args schema.
+        goals : The goals.
+        resource_manager: Manages the file resources
+    """
+    llm: Optional[BaseLlm] = None
+    agent_id: int = None
+    name = "CodingTool"
+    description = (
+        "You will get instructions for code to write. You will write a very long answer. "
+        "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."
+    )
+    args_schema: Type[CodingSchema] = CodingSchema
+    goals: List[str] = []
+    resource_manager: Optional[ResourceManager] = None
+    tool_response_manager: Optional[ToolResponseQueryManager] = None
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+    def _execute(self, code_description: str) -> str:
+        """
+        Execute the write_code tool.
+
+        Args:
+            code_description : The coding task description.
+            code_file_name: The name of the file where the generated codes will be saved.
+
+        Returns:
+            Generated codes files or error message.
+        """
+        try:
+            prompt = """You are a super smart developer who practices good Development for writing code according to a specification.
+
+            Your high-level goal is:
+            {goals}
+            
+            Coding task description:
+            {code_description}
+
+            {spec}
+
+            You will get instructions for code to write.
+            You need to write a detailed answer. Make sure all parts of the architecture are turned into code.
+            Think carefully about each step and make good choices to get it right. First, list the main classes, 
+            functions, methods you'll use and a quick comment on their purpose.
+
+            Then you will output the content of each file including ALL code.
+            Each file must strictly follow a markdown code block format, where the following tokens must be replaced such that
+            [FILENAME] is the lowercase file name including the file extension,
+            [LANG] is the markup code block language for the code's language, and [CODE] is the code: 
+            [FILENAME]
+            ```[LANG]
+            [CODE]
+            ```
+
+            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.
+            Ensure to implement all code, if you are unsure, write a plausible implementation.
+            Include module dependency or package manager dependency definition file.
+            Before you finish, double check that all parts of the architecture is present in the files.
+            """
+            prompt = prompt.replace("{goals}", AgentPromptBuilder.add_list_items_to_string(self.goals))
+            prompt = prompt.replace("{code_description}", code_description)
+            spec_response = self.tool_response_manager.get_last_response("WriteSpecTool")
+            if spec_response != "":
+                prompt = prompt.replace("{spec}", "Use this specs for generating the code:\n" + spec_response)
+            logger.info(prompt)
+            messages = [{"role": "system", "content": prompt}]
+
+            total_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
+            token_limit = TokenCounter.token_limit(self.llm.get_model())
+            result = self.llm.chat_completion(messages, max_tokens=(token_limit - total_tokens - 100))
+
+            # Get all filenames and corresponding code blocks
+            regex = r"(\S+?)\n```\S*\n(.+?)```"
+            matches = re.finditer(regex, result["content"], re.DOTALL)
+
+            file_names = []
+            # Save each file
+
+            for match in matches:
+                # Get the filename
+                file_name = re.sub(r'[<>"|?*]', "", match.group(1))
+
+                # Get the code
+                code = match.group(2)
+
+                # Ensure file_name is not empty
+                if not file_name.strip():
+                    continue
+
+                file_names.append(file_name)
+                save_result = self.resource_manager.write_file(file_name, code)
+                if save_result.startswith("Error"):
+                    return save_result
+
+            # Get README contents and save
+            split_result = result["content"].split("```")
+            if len(split_result) > 0:
+                readme = split_result[0]
+                save_readme_result = self.resource_manager.write_file("README.md", readme)
+                if save_readme_result.startswith("Error"):
+                    return save_readme_result
+
+            return result["content"] + "\n Codes generated and saved successfully in " + ", ".join(file_names)
+        except Exception as e:
+            logger.error(e)
+            return f"Error generating codes: {e}"