From 281753fdcf3f3b2533fd3f6614cc12d9e6baa00b Mon Sep 17 00:00:00 2001 From: luciferlinx <129729795+luciferlinx101@users.noreply.github.com> Date: Wed, 19 Jul 2023 13:53:12 +0530 Subject: [PATCH] Resource duplicate fix (#803) --- superagi/helper/resource_helper.py | 40 +++++++++++++++---- superagi/resource_manager/file_manager.py | 15 +++---- .../unit_tests/helper/test_resource_helper.py | 23 +++++++++-- 3 files changed, 57 insertions(+), 21 deletions(-) diff --git a/superagi/helper/resource_helper.py b/superagi/helper/resource_helper.py index cd39eeb14..ec606e40e 100644 --- a/superagi/helper/resource_helper.py +++ b/superagi/helper/resource_helper.py @@ -10,7 +10,7 @@ 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. @@ -18,6 +18,7 @@ def make_written_file_resource(cls, file_name: str, agent: Agent, agent_executio 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: diff --git a/superagi/resource_manager/file_manager.py b/superagi/resource_manager/file_manager.py index c93c64c6d..6d681b5f5 100644 --- a/superagi/resource_manager/file_manager.py +++ b/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) diff --git a/tests/unit_tests/helper/test_resource_helper.py b/tests/unit_tests/helper/test_resource_helper.py index 4cc9a5f65..d1e9eb42a 100644 --- a/tests/unit_tests/helper/test_resource_helper.py +++ b/tests/unit_tests/helper/test_resource_helper.py @@ -1,8 +1,9 @@ -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): @@ -10,12 +11,28 @@ def test_make_written_file_resource(mocker): 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'