Resource manager s3 fix (TransformerOptimus#935)

GuardingUs-AI · Aug 2, 2023 · 40a23be · 40a23be
1 parent 7900f4d
commit 40a23be
Show file tree

Hide file tree

Showing 11 changed files with 125 additions and 49 deletions.
diff --git a/gui/pages/Content/Marketplace/MarketKnowledge.js b/gui/pages/Content/Marketplace/MarketKnowledge.js
@@ -18,7 +18,7 @@ export default function MarketKnowledge() {
 
     if (window.location.href.toLowerCase().includes('marketplace')) {
       setShowMarketplace(true);
-      axios.get(`https://app.superagi.com/api/knowledge/get/list?page=0`)
+      axios.get(`https://app.superagi.com/api/knowledges/get/list?page=0`)
         .then((response) => {
           const data = response.data || [];
           setKnowledgeTemplates(data);

diff --git a/superagi/helper/github_helper.py b/superagi/helper/github_helper.py
@@ -3,6 +3,12 @@
 
 import requests
 from superagi.lib.logger import logger
+from superagi.helper.resource_helper import ResourceHelper
+from superagi.models.agent import Agent
+from superagi.models.agent_execution import AgentExecution
+from superagi.types.storage_types import StorageType
+from superagi.config.config import get_config
+from superagi.helper.s3_helper import S3Helper
 
 
 class GithubHelper:
@@ -198,7 +204,7 @@ def delete_file(self, repository_name, file_name, folder_path, commit_message, h
         return file_response.status_code
 
     def add_file(self, repository_owner, repository_name, file_name, folder_path, head_branch, base_branch, headers,
-                 body, commit_message):
+                 commit_message, agent_id, agent_execution_id, session):
         """
         Adds a file to the given repository.
 
@@ -213,7 +219,7 @@ def add_file(self, repository_owner, repository_name, file_name, folder_path, he
         Returns:
             None
         """
-
+        body = self._get_file_contents(file_name, agent_id, agent_execution_id, session)
         body_bytes = body.encode("ascii")
         base64_bytes = base64.b64encode(body_bytes)
         file_content = base64_bytes.decode("ascii")
@@ -317,3 +323,15 @@ def validate_github_link(cls, link: str) -> bool:
             return True
 
         return False
+
+    def _get_file_contents(self, file_name, agent_id, agent_execution_id, session):
+        final_path = ResourceHelper().get_agent_read_resource_path(file_name,
+                                                                    agent=Agent.get_agent_from_id(session, agent_id),
+                                                                    agent_execution=AgentExecution.get_agent_execution_from_id(
+                                                                  session, agent_execution_id))
+        if StorageType.get_storage_type(get_config("STORAGE_TYPE", StorageType.FILE.value)) == StorageType.S3:
+                attachment_data = S3Helper().read_from_s3(final_path)
+        else:
+            with open(final_path, "r") as file:
+                attachment_data = file.read().decode('utf-8')
+        return attachment_data
diff --git a/superagi/helper/s3_helper.py b/superagi/helper/s3_helper.py
@@ -60,6 +60,14 @@ def read_from_s3(self, file_path):
             return response['Body'].read().decode('utf-8')
         raise Exception(f"Error read_from_s3: {response}")
 
+    def read_binary_from_s3(self, file_path):
+        file_path = "resources" + file_path
+        logger.info(f"Reading file from s3: {file_path}")
+        response = self.s3.get_object(Bucket=get_config("BUCKET_NAME"), Key=file_path)
+        if response['ResponseMetadata']['HTTPStatusCode'] == 200:
+            return response['Body'].read()
+        raise Exception(f"Error read_from_s3: {response}")
+
     def get_json_file(self, path):
         """
         Get a JSON file from S3.
@@ -76,3 +84,22 @@ def get_json_file(self, path):
             return json.loads(s3_response)
         except:
             raise HTTPException(status_code=500, detail="AWS credentials not found. Check your configuration.")
+
+    def delete_file(self, path):
+        """
+        Delete a file from S3.
+
+        Args:
+            path (str): The path to the file to delete.
+
+        Raises:
+            HTTPException: If the AWS credentials are not found.
+
+        Returns:
+            None
+        """
+        try:
+            self.s3.delete_object(Bucket=self.bucket_name, Key=path)
+            logger.info("File deleted from S3 successfully!")
+        except:
+            raise HTTPException(status_code=500, detail="AWS credentials not found. Check your configuration.")
diff --git a/superagi/helper/twitter_helper.py b/superagi/helper/twitter_helper.py
@@ -7,6 +7,9 @@
 from superagi.helper.resource_helper import ResourceHelper
 from superagi.models.agent import Agent
 from superagi.models.agent_execution import AgentExecution
+from superagi.types.storage_types import StorageType
+from superagi.config.config import get_config
+from superagi.helper.s3_helper import S3Helper
 
 
 class TwitterHelper:
@@ -19,7 +22,7 @@ def get_media_ids(self, session, media_files, creds, agent_id, agent_execution_i
                        resource_owner_secret=creds.oauth_token_secret)
         for file in media_files:
             file_path = self.get_file_path(session, file, agent_id, agent_execution_id)
-            image_data = open(file_path, 'rb').read()
+            image_data = self._get_image_data(file_path)
             b64_image = base64.b64encode(image_data)
             upload_endpoint = 'https://upload.twitter.com/1.1/media/upload.json'
             headers = {'Authorization': 'application/octet-stream'}
@@ -31,12 +34,20 @@ def get_media_ids(self, session, media_files, creds, agent_id, agent_execution_i
         return media_ids
 
     def get_file_path(self, session, file_name, agent_id, agent_execution_id):
-        final_path = ResourceHelper().get_agent_write_resource_path(file_name,
+        final_path = ResourceHelper().get_agent_read_resource_path(file_name,
                                                                     agent=Agent.get_agent_from_id(session, agent_id),
                                                                     agent_execution=AgentExecution.get_agent_execution_from_id(
                                                                   session, agent_execution_id))
         return final_path
 
+    def _get_image_data(self, file_path):
+        if StorageType.get_storage_type(get_config("STORAGE_TYPE", StorageType.FILE.value)) == StorageType.S3:
+                attachment_data = S3Helper().read_binary_from_s3(file_path)
+        else:
+            with open(file_path, "rb") as file:
+                attachment_data = file.read()
+        return attachment_data
+
     def send_tweets(self, params, creds):
         tweet_endpoint = "https://api.twitter.com/2/tweets"
         oauth = OAuth1Session(creds.api_key,

diff --git a/superagi/tools/email/send_email_attachment.py b/superagi/tools/email/send_email_attachment.py
@@ -10,8 +10,11 @@
 from superagi.helper.imap_email import ImapEmail
 from superagi.tools.base_tool import BaseTool
 from superagi.helper.resource_helper import ResourceHelper
+from superagi.helper.s3_helper import S3Helper
 from superagi.models.agent import Agent
 from superagi.models.agent_execution import AgentExecution
+from superagi.config.config import get_config
+from superagi.types.storage_types import StorageType
 
 
 class SendEmailAttachmentInput(BaseModel):
@@ -94,8 +97,12 @@ def send_email_with_attachment(self, to, subject, body, attachment_path, attachm
             if ctype is None or encoding is not None:
                 ctype = "application/octet-stream"
             maintype, subtype = ctype.split("/", 1)
-            with open(attachment_path, "rb") as file:
-                message.add_attachment(file.read(), maintype=maintype, subtype=subtype, filename=attachment)
+            if StorageType.get_storage_type(get_config("STORAGE_TYPE", StorageType.FILE.value)) == StorageType.S3:
+                attachment_data = S3Helper().read_binary_from_s3(attachment_path)
+            else:
+                with open(attachment_path, "rb") as file:
+                    attachment_data = file.read()
+            message.add_attachment(attachment_data, maintype=maintype, subtype=subtype, filename=attachment)
 
         send_to_draft = self.get_tool_config('EMAIL_DRAFT_MODE') or "FALSE"
         if send_to_draft.upper() == "TRUE":

diff --git a/superagi/tools/file/delete_file.py b/superagi/tools/file/delete_file.py
@@ -7,6 +7,9 @@
 from superagi.models.agent_execution import AgentExecution
 from superagi.models.agent import Agent
 from superagi.tools.base_tool import BaseTool
+from superagi.types.storage_types import StorageType
+from superagi.config.config import get_config
+from superagi.helper.s3_helper import S3Helper
 
 
 class DeleteFileInput(BaseModel):
@@ -46,8 +49,16 @@ def _execute(self, file_name: str):
           AgentExecution.get_agent_execution_from_id(
               session=self.toolkit_config.session,
               agent_execution_id=self.agent_execution_id))
-        try:
-            os.remove(final_path)
-            return "File deleted successfully."
-        except Exception as err:
-            return f"Error: {err}"
+
+        if StorageType.get_storage_type(get_config("STORAGE_TYPE", StorageType.FILE.value)) == StorageType.S3:
+            try:
+                S3Helper().delete_file(final_path)
+                return "File deleted successfully."
+            except Exception as err:
+                return f"Error: {err}"
+        else:
+            try:
+                os.remove(final_path)
+                return "File deleted successfully."
+            except Exception as err:
+                return f"Error: {err}"
diff --git a/superagi/tools/github/add_file.py b/superagi/tools/github/add_file.py
@@ -24,10 +24,6 @@ class GithubAddFileSchema(BaseModel):
         ...,
         description="folder path for the file to be stored",
     )
-    body: str = Field(
-        ...,
-        description="content to be stored",
-    )
     commit_message: str = Field(
         ...,
         description="clear description of the contents of file",
@@ -50,16 +46,17 @@ class GithubAddFileTool(BaseTool):
     name: str = "Github Add File"
     args_schema: Type[BaseModel] = GithubAddFileSchema
     description: str = "Add a file or folder to a particular github repository"
+    agent_id: int = None
+    agent_execution_id: int = None
 
-    def _execute(self, repository_name: str, base_branch: str, body: str, commit_message: str, repository_owner: str,
+    def _execute(self, repository_name: str, base_branch: str, commit_message: str, repository_owner: str,
                  file_name='.gitkeep', folder_path=None) -> str:
         """
         Execute the add file tool.
 
         Args:
             repository_name : The name of the repository to add file to.
             base_branch : The branch to interact with.
-            body : The content to be stored.
             commit_message : Clear description of the contents of file.
             repository_owner : Owner of the GitHub repository.
             file_name : The name of the file to add.
@@ -68,6 +65,7 @@ def _execute(self, repository_name: str, base_branch: str, body: str, commit_mes
         Returns:
             Pull request success message if pull request is created successfully else error message.
         """
+        session = self.toolkit_config.session
         try:
             github_access_token = self.get_tool_config("GITHUB_ACCESS_TOKEN")
             github_username = self.get_tool_config("GITHUB_USERNAME")
@@ -82,7 +80,7 @@ def _execute(self, repository_name: str, base_branch: str, body: str, commit_mes
 
             branch_response = github_helper.create_branch(repository_name, base_branch, head_branch, headers)
             file_response = github_helper.add_file(repository_owner, repository_name, file_name, folder_path,
-                                                   head_branch, base_branch, headers, body, commit_message)
+                                                   head_branch, base_branch, headers, commit_message, self.agent_id, self.agent_execution_id, session)
             pr_response = github_helper.create_pull_request(repository_owner, repository_name, head_branch, base_branch,
                                                             headers)
             if (pr_response == 201 or pr_response == 422) and (file_response == 201 or file_response == 422):

diff --git a/superagi/tools/github/search_repo.py b/superagi/tools/github/search_repo.py
@@ -55,6 +55,8 @@ def _execute(self, repository_owner: str, repository_name: str, file_name: str,
         github_access_token = self.get_tool_config("GITHUB_ACCESS_TOKEN")
         github_username = self.get_tool_config("GITHUB_USERNAME")
         github_repo_search = GithubHelper(github_access_token, github_username)
-        content = github_repo_search.get_content_in_file(repository_owner, repository_name, file_name, folder_path)
-
-        return content
+        try:
+            content = github_repo_search.get_content_in_file(repository_owner, repository_name, file_name, folder_path)
+            return content
+        except:
+            return "File not found"
diff --git a/tests/unit_tests/helper/test_github_helper.py b/tests/unit_tests/helper/test_github_helper.py
@@ -132,28 +132,6 @@ def test_delete_file(self, mock_delete):
             headers={'header': 'value'}
         )
 
-    @patch('requests.put')
-    def test_add_file(self, mock_put):
-        # Create response mock
-        mock_resp = MagicMock()
-        mock_resp.status_code = 201
-        mock_put.return_value = mock_resp
-
-        gh = GithubHelper('access_token', 'username')
-        status_code = gh.add_file('owner', 'repo', 'test.txt', 'path', 'head', 'base', {'header': 'value'}, 'body',
-                                  'message')
-
-        self.assertEqual(status_code, 201)
-        mock_put.assert_called_once_with(
-            'https://api.github.com/repos/username/repo/contents/path/test.txt',
-            json={
-                'message': 'message',
-                'content': base64.b64encode('body'.encode("ascii")).decode("ascii"),
-                'branch': 'head'
-            },
-            headers={'header': 'value'}
-        )
-
     @patch('requests.post')
     def test_create_pull_request(self, mock_post):
         # Create response mock

diff --git a/tests/unit_tests/helper/test_s3_helper.py b/tests/unit_tests/helper/test_s3_helper.py
@@ -74,4 +74,32 @@ def test_read_from_s3(s3helper_object, http_status, expected_result, raises):
         with pytest.raises(Exception):
             s3helper_object.read_from_s3('path')
     else:
-        assert s3helper_object.read_from_s3('path') == expected_result
+        assert s3helper_object.read_from_s3('path') == expected_result
+
+@pytest.mark.parametrize('http_status, expected_result, raises',
+                         [(200, b'file_content', False),
+                          (500, None, True)])
+def test_read_binary_from_s3(s3helper_object, http_status, expected_result, raises):
+    s3helper_object.s3.get_object = MagicMock(
+        return_value={'ResponseMetadata': {'HTTPStatusCode': http_status},
+                      'Body': MagicMock(read=lambda: (expected_result))}
+    )
+
+    if raises:
+        with pytest.raises(Exception):
+            s3helper_object.read_binary_from_s3('path')
+    else:
+        assert s3helper_object.read_binary_from_s3('path') == expected_result
+
+def test_delete_file_success(s3helper_object):
+    s3helper_object.s3.delete_object = MagicMock()
+    try:
+        s3helper_object.delete_file('path')
+    except:
+        pytest.fail("Unexpected Exception !")
+
+def test_delete_file_fail(s3helper_object):
+    s3helper_object.s3.delete_object = MagicMock(side_effect=Exception())
+    with pytest.raises(HTTPException):
+        s3helper_object.delete_file('path')
+
diff --git a/tests/unit_tests/tools/github/test_add_file.py b/tests/unit_tests/tools/github/test_add_file.py
@@ -11,7 +11,6 @@ def test_github_add_file_schema():
         base_branch="main",
         file_name="test_file",
         folder_path="test_folder",
-        body="test_content",
         commit_message="test_commit",
         repository_owner="test_owner"
     )
@@ -20,7 +19,6 @@ def test_github_add_file_schema():
     assert schema.base_branch == "main"
     assert schema.file_name == "test_file"
     assert schema.folder_path == "test_folder"
-    assert schema.body == "test_content"
     assert schema.commit_message == "test_commit"
     assert schema.repository_owner == "test_owner"
 
@@ -45,7 +43,6 @@ def test_github_add_file_tool_execute(mock_make_fork, mock_create_branch, mock_a
     response = github_add_file_tool._execute(
         repository_name="test_repo",
         base_branch="main",
-        body="test_content",
         commit_message="test_commit",
         repository_owner="test_owner",
         file_name="test_file",
@@ -62,7 +59,6 @@ def test_github_add_file_tool_execute(mock_make_fork, mock_create_branch, mock_a
     response = github_add_file_tool._execute(
         repository_name="test_repo",
         base_branch="main",
-        body="test_content",
         commit_message="test_commit",
         repository_owner="test_owner",
         file_name="test_file",