Resource duplicate fix (#803)

superagi/helper/resource_helper.py

@@ -10,14 +10,15 @@
 
 class ResourceHelper:
     @classmethod
-    def make_written_file_resource(cls, file_name: str, agent: Agent, agent_execution: AgentExecution):
+    def make_written_file_resource(cls, file_name: str, agent: Agent, agent_execution: AgentExecution, session):
         """
         Function to create a Resource object for a written file.
 
         Args:
             file_name (str): The name of the file.
             agent (Agent): Agent related to resource.
             agent_execution(AgentExecution): Agent Execution related to a resource
+            session (Session): The database session.
 
         Returns:
             Resource: The Resource object.
@@ -46,13 +47,36 @@ def make_written_file_resource(cls, file_name: str, agent: Agent, agent_executio
         logger.info("make_written_file_resource:", final_path)
         if StorageType.get_storage_type(get_config("STORAGE_TYPE", StorageType.FILE.value)) == StorageType.S3:
             file_path = "resources" + file_path
-        resource = Resource(name=file_name, path=file_path, storage_type=storage_type.value,
-                            size=file_size,
-                            type=file_type,
-                            channel="OUTPUT",
-                            agent_id=agent.id,
-                            agent_execution_id=agent_execution.id)
-        return resource
+        existing_resource = session.query(Resource).filter_by(
+            name=file_name,
+            path=file_path,
+            storage_type=storage_type.value,
+            type=file_type,
+            channel="OUTPUT",
+            agent_id=agent.id,
+            agent_execution_id=agent_execution.id
+        ).first()
+
+        if existing_resource:
+            # Update the existing resource attributes
+            existing_resource.size = file_size
+            session.commit()
+            session.flush()
+            return existing_resource
+        else:
+            resource = Resource(
+                name=file_name,
+                path=file_path,
+                storage_type=storage_type.value,
+                size=file_size,
+                type=file_type,
+                channel="OUTPUT",
+                agent_id=agent.id,
+                agent_execution_id=agent_execution.id
+            )
+            session.add(resource)
+            session.commit()
+            return resource
 
     @classmethod
     def get_formatted_agent_level_path(cls, agent: Agent, path) -> object:

superagi/resource_manager/file_manager.py

@@ -26,7 +26,6 @@ def write_binary_file(self, file_name: str, data):
                                                                           self.agent_execution_id))
         else:
             final_path = ResourceHelper.get_resource_path(file_name)
-
         try:
             with open(final_path, mode="wb") as img:
                 img.write(data)
@@ -44,14 +43,11 @@ def write_to_s3(self, file_name, final_path):
                                                                                                self.agent_id),
                                                                  agent_execution=AgentExecution
                                                                  .get_agent_execution_from_id(self.session,
-                                                                                              self.agent_execution_id))
-            if resource is not None:
-                self.session.add(resource)
-                self.session.commit()
-                self.session.flush()
-                if resource.storage_type == StorageType.S3.value:
-                    s3_helper = S3Helper()
-                    s3_helper.upload_file(img, path=resource.path)
+                                                                                              self.agent_execution_id),
+                                                                 session=self.session)
+            if resource.storage_type == StorageType.S3.value:
+                s3_helper = S3Helper()
+                s3_helper.upload_file(img, path=resource.path)
 
     def write_file(self, file_name: str, content):
         if self.agent_id is not None:
@@ -63,7 +59,6 @@ def write_file(self, file_name: str, content):
                                                                                                    self.agent_execution_id))
         else:
             final_path = ResourceHelper.get_resource_path(file_name)
-
         try:
             with open(final_path, mode="w") as file:
                 file.write(content)

tests/unit_tests/helper/test_resource_helper.py

@@ -1,21 +1,38 @@
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 
 from superagi.helper.resource_helper import ResourceHelper
 from superagi.models.agent import Agent
 from superagi.models.agent_execution import AgentExecution
+from superagi.models.resource import Resource
 
 
 def test_make_written_file_resource(mocker):
     mocker.patch('os.getcwd', return_value='/')
     mocker.patch('os.makedirs', return_value=None)
     mocker.patch('os.path.getsize', return_value=1000)
     mocker.patch('os.path.splitext', return_value=("", ".txt"))
-    mocker.patch('superagi.helper.resource_helper.get_config', side_effect=['FILE','/','/','FILE'])
+    mocker.patch('superagi.helper.resource_helper.get_config', side_effect=['FILE', '/', '/', 'FILE'])
     mock_agent = Agent(id=1, name='TestAgent')
     mock_agent_execution = AgentExecution(id=1, name='TestExecution')
+    session = MagicMock()
 
     with patch('superagi.helper.resource_helper.logger') as logger_mock:
-        result = ResourceHelper.make_written_file_resource('test.txt', mock_agent, mock_agent_execution)
+        session.query.return_value.filter_by.return_value.first.return_value = None
+        # Create a Resource object
+        resource = Resource(
+            name='test.txt',
+            path='/test.txt',
+            storage_type='FILE',
+            size=1000,
+            type='application/txt',
+            channel='OUTPUT',
+            agent_id=1,
+            agent_execution_id=1
+        )
+
+        # Mock the session.add() method to return the created Resource object
+        session.add.return_value = resource
+        result = ResourceHelper.make_written_file_resource('test.txt', mock_agent, mock_agent_execution, session)
 
     assert result.name == 'test.txt'
     assert result.path == '/test.txt'