Merge pull request #435 from TransformerOptimus/resource_manager_refa…

…ctoring1

Resource manager refactoring

requirements.txt

@@ -132,3 +132,4 @@ tiktoken==0.4.0
 psycopg2==2.9.6
 slack-sdk==3.21.3
 pytest==7.3.2
+pytest-cov==4.1.0

superagi/helper/resource_helper.py

@@ -8,14 +8,13 @@
 class ResourceHelper:
 
     @staticmethod
-    def make_written_file_resource(file_name: str, agent_id: int, file, channel):
+    def make_written_file_resource(file_name: str, agent_id: int, channel: str):
         """
         Function to create a Resource object for a written file.
 
         Args:
             file_name (str): The name of the file.
             agent_id (int): The ID of the agent.
-            file (FileStorage): The file.
             channel (str): The channel of the file.
 
         Returns:
@@ -32,29 +31,35 @@ def make_written_file_resource(file_name: str, agent_id: int, file, channel):
         else:
             file_type = "application/misc"
 
-        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
-
+        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]
-            if channel == "INPUT":
-                path = 'input'
-            else:
-                path = 'output'
+            file_name = file_name_parts[0] + '_' + str(datetime.datetime.now()).replace(' ', '') \
+                .replace('.', '').replace(':', '') + '.' + file_name_parts[1]
+            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,
                             type=file_type,
                             channel="OUTPUT",
                             agent_id=agent_id)
         return resource
+
+    @staticmethod
+    def get_resource_path(file_name: str):
+        """Get final path of the resource.
+
+        Args:
+            file_name (str): The name of the file.
+        """
+        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
+        return final_path

superagi/jobs/agent_executor.py

@@ -17,6 +17,7 @@
 from superagi.models.organisation import Organisation
 from superagi.models.project import Project
 from superagi.models.tool import Tool
+from superagi.resource_manager.manager import ResourceManager
 from superagi.tools.thinking.tools import ThinkingTool
 from superagi.vector_store.embedding.openai import OpenAiEmbedding
 from superagi.vector_store.vector_factory import VectorFactory
@@ -164,7 +165,7 @@ def execute_next_action(self, agent_execution_id):
         print(user_tools)
 
         tools = self.set_default_params_tools(tools, parsed_config, agent_execution.agent_id,
-                                              model_api_key=model_api_key)
+                                              model_api_key=model_api_key, session=session)
 
 
 
@@ -205,7 +206,7 @@ def execute_next_action(self, agent_execution_id):
         # finally:
         engine.dispose()
 
-    def set_default_params_tools(self, tools, parsed_config, agent_id, model_api_key):
+    def set_default_params_tools(self, tools, parsed_config, agent_id, model_api_key, session):
         """
         Set the default parameters for the tools.
 
@@ -232,6 +233,9 @@ def set_default_params_tools(self, tools, parsed_config, agent_id, model_api_key
                 tool.image_llm = OpenAi(model=parsed_config["model"], api_key=model_api_key)
             if hasattr(tool, 'agent_id'):
                 tool.agent_id = agent_id
+            if hasattr(tool, 'resource_manager'):
+                tool.resource_manager = ResourceManager(session=session)
+
             new_tools.append(tool)
         return tools
 

superagi/resource_manager/manager.py

@@ -0,0 +1,46 @@
+from sqlalchemy.orm import Session
+
+from superagi.helper.resource_helper import ResourceHelper
+from superagi.helper.s3_helper import S3Helper
+from superagi.lib.logger import logger
+
+
+class ResourceManager:
+    def __init__(self, session: Session):
+        self.session = session
+
+    def write_binary_file(self, file_name: str, data):
+        final_path = ResourceHelper.get_resource_path(file_name)
+
+        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")
+        except Exception as err:
+            return f"Error: {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,
+                                                                 agent_id=self.agent_id, channel="OUTPUT")
+            if resource is not None:
+                self.session.add(resource)
+                self.session.commit()
+                self.session.flush()
+                if resource.storage_type == "S3":
+                    s3_helper = S3Helper()
+                    s3_helper.upload_file(img, path=resource.path)
+
+    def write_file(self, file_name: str, content):
+        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")
+        except Exception as err:
+            return f"Error: {err}"

superagi/tools/email/send_email.py

@@ -57,7 +57,7 @@ def _execute(self, to: str, subject: str, body: str) -> str:
             body += f"\n{signature}"
         message.set_content(body)
         draft_folder = get_config('EMAIL_DRAFT_MODE_WITH_FOLDER')
-        send_to_draft = draft_folder is not None or draft_folder != "YOUR_DRAFTS_FOLDER"
+        send_to_draft = draft_folder is not None and draft_folder != "YOUR_DRAFTS_FOLDER"
         if message["To"] == "example@example.com" or send_to_draft:
             conn = ImapEmail().imap_open(draft_folder, email_sender, email_password)
             conn.append(

superagi/tools/file/append_file.py

@@ -3,6 +3,7 @@
 
 from pydantic import BaseModel, Field
 from superagi.config.config import get_config
+from superagi.helper.resource_helper import ResourceHelper
 
 from superagi.tools.base_tool import BaseTool
 
@@ -38,13 +39,7 @@ def _execute(self, file_name: str, content: str):
         Returns:
             file written to successfully. or error message.
         """
-        final_path = file_name
-        root_dir = get_config('RESOURCES_OUTPUT_ROOT_DIR')
-        if root_dir is not None:
-            root_dir = root_dir if root_dir.endswith("/") else root_dir + "/"
-            final_path = root_dir + file_name
-        else:
-            final_path = os.getcwd() + "/" + file_name
+        final_path = ResourceHelper.get_resource_path(file_name)
         try:
             directory = os.path.dirname(final_path)
             os.makedirs(directory, exist_ok=True)

superagi/tools/file/delete_file.py

@@ -3,6 +3,7 @@
 
 from pydantic import BaseModel, Field
 
+from superagi.helper.resource_helper import ResourceHelper
 from superagi.tools.base_tool import BaseTool
 from superagi.config.config import get_config
 
@@ -36,13 +37,7 @@ def _execute(self, file_name: str, content: str):
         Returns:
             file deleted successfully. or error message.
         """
-        final_path = file_name
-        root_dir = get_config('RESOURCES_INPUT_ROOT_DIR')
-        if root_dir is not None:
-            root_dir = root_dir if root_dir.endswith("/") else root_dir + "/"
-            final_path = root_dir + file_name
-        else:
-            final_path = os.getcwd() + "/" + file_name
+        final_path = ResourceHelper.get_resource_path(file_name)
         try:
             os.remove(final_path)
             return "File deleted successfully."

superagi/tools/file/write_file.py

@@ -1,16 +1,12 @@
-import os
-from typing import Type
+from typing import Type, Optional
+
 from pydantic import BaseModel, Field
+
+from superagi.resource_manager.manager import ResourceManager
 from superagi.tools.base_tool import BaseTool
-from superagi.config.config import get_config
-from sqlalchemy.orm import sessionmaker
-from superagi.models.db import connect_db
-from superagi.helper.resource_helper import ResourceHelper
-# from superagi.helper.s3_helper import upload_to_s3
-from superagi.helper.s3_helper import S3Helper
-from superagi.lib.logger import logger
 
 
+# from superagi.helper.s3_helper import upload_to_s3
 
 
 class WriteFileInput(BaseModel):
@@ -32,6 +28,10 @@ class WriteFileTool(BaseTool):
     args_schema: Type[BaseModel] = WriteFileInput
     description: str = "Writes text to a file"
     agent_id: int = None
+    resource_manager: Optional[ResourceManager] = None
+
+    class Config:
+        arbitrary_types_allowed = True
 
     def _execute(self, file_name: str, content: str):
         """
@@ -44,35 +44,5 @@ def _execute(self, file_name: str, content: str):
         Returns:
             file written to successfully. or error message.
         """
-        engine = connect_db()
-        Session = sessionmaker(bind=engine)
-        session = Session()
-
-        final_path = file_name
-        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
+        self.resource_manager.write_file(file_name, content)
 
-        try:
-            with open(final_path, 'w', encoding="utf-8") as file:
-                file.write(content)
-                file.close()
-            with open(final_path, 'rb') as file:
-                resource = ResourceHelper.make_written_file_resource(file_name=file_name,
-                                                      agent_id=self.agent_id,file=file,channel="OUTPUT")
-                if resource is not None:
-                    session.add(resource)
-                    session.commit()
-                    session.flush()
-                    if resource.storage_type == "S3":
-                        s3_helper = S3Helper()
-                        s3_helper.upload_file(file, path=resource.path)
-                        logger.info("Resource Uploaded to S3!")
-                session.close()
-            return f"File written to successfully - {file_name}"
-        except Exception as err:
-            return f"Error: {err}"

superagi/tools/image_generation/dalle_image_gen.py

@@ -1,92 +1,60 @@
 from typing import Type, Optional
+
+import requests
 from pydantic import BaseModel, Field
+
 from superagi.llms.base_llm import BaseLlm
+from superagi.resource_manager.manager import ResourceManager
 from superagi.tools.base_tool import BaseTool
-from superagi.config.config import get_config
-import os
-import requests
-from superagi.models.db import connect_db
-from superagi.helper.resource_helper import ResourceHelper
-from superagi.helper.s3_helper import S3Helper
-from sqlalchemy.orm import sessionmaker
-from superagi.lib.logger import logger
-
-
 
-class ImageGenInput(BaseModel):
+class DalleImageGenInput(BaseModel):
     prompt: str = Field(..., description="Prompt for Image Generation to be used by Dalle.")
     size: int = Field(..., description="Size of the image to be Generated. default size is 512")
     num: int = Field(..., description="Number of Images to be generated. default num is 2")
-    image_name: list = Field(..., description="Image Names for the generated images, example 'image_1.png'. Only include the image name. Don't include path.")
+    image_names: list = Field(..., description="Image Names for the generated images, example 'image_1.png'. Only include the image name. Don't include path.")
 
 
-class ImageGenTool(BaseTool):
+class DalleImageGenTool(BaseTool):
     """
     Dalle Image Generation tool
 
     Attributes:
-        name : The name.
-        description : The description.
-        args_schema : The args schema.
+        name : Name of the tool
+        description : The description
+        args_schema : The args schema
+        llm : The llm
+        agent_id : The agent id
+        resource_manager : Manages the file resources
     """
-    name: str = "Dalle Image Generation"
-    args_schema: Type[BaseModel] = ImageGenInput
+    name: str = "DalleImageGeneration"
+    args_schema: Type[BaseModel] = DalleImageGenInput
     description: str = "Generate Images using Dalle"
     llm: Optional[BaseLlm] = None
     agent_id: int = None
+    resource_manager: Optional[ResourceManager] = None
 
     class Config:
         arbitrary_types_allowed = True
 
-    def _execute(self, prompt: str, image_name: list, size: int = 512, num: int = 2):
+    def _execute(self, prompt: str, image_names: list, size: int = 512, num: int = 2):
         """
         Execute the Dalle Image Generation tool.
 
         Args:
             prompt : The prompt for image generation.
             size : The size of the image to be generated.
             num : The number of images to be generated.
-            image_name (list): The name of the image to be generated.
+            image_names (list): The name of the image to be generated.
 
         Returns:
             Image generated successfully. or error message.
         """
-        engine = connect_db()
-        Session = sessionmaker(bind=engine)
-        session = Session()
         if size not in [256, 512, 1024]:
             size = min([256, 512, 1024], key=lambda x: abs(x - size))
         response = self.llm.generate_image(prompt, size, num)
         response = response.__dict__
         response = response['_previous']['data']
         for i in range(num):
-            image = image_name[i]
-            final_path = image
-            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 + image
-            else:
-                final_path = os.getcwd() + "/" + image
-            url = response[i]['url']
-            data = requests.get(url).content
-            try:
-                with open(final_path, mode="wb") as img:
-                    img.write(data)
-                with open(final_path, 'rb') as img:
-                    resource = ResourceHelper.make_written_file_resource(file_name=image_name[i],
-                                                          agent_id=self.agent_id, file=img,channel="OUTPUT")
-                    if resource is not None:
-                        session.add(resource)
-                        session.commit()
-                        session.flush()
-                        if resource.storage_type == "S3":
-                            s3_helper = S3Helper()
-                            s3_helper.upload_file(img, path=resource.path)
-                logger.info(f"Image {image} saved successfully")
-            except Exception as err:
-                session.close()
-                return f"Error: {err}"
-        session.close()
+            data = requests.get(response[i]['url']).content
+            self.resource_manager.write_binary_file(image_names[i], data)
         return "Images downloaded successfully"