Resource Manager Delete File issue resolved

TransformerOptimus · Sep 19, 2023 · 5f6f2a2 · 5f6f2a2
1 parent 2ccbd4b
commit 5f6f2a2
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/superagi/models/resource.py b/superagi/models/resource.py
@@ -63,6 +63,14 @@ def validate_resource_type(storage_type):
     def find_by_run_ids(cls, session, run_ids: list):
         db_resources_arr=session.query(Resource).filter(Resource.agent_execution_id.in_(run_ids)).all()
         return db_resources_arr
+
+    @classmethod
+    def delete_resource(cls, session, file_name, agent_id, agent_execution_id):
+        deleted = session.query(Resource).filter(Resource.name == file_name, Resource.agent_id == agent_id,
+                                                 Resource.agent_execution_id == agent_execution_id).delete()
+        session.commit()
+
+        return deleted
 
 class InvalidResourceType(Exception):
     """Custom exception for invalid resource type"""
diff --git a/superagi/tools/file/delete_file.py b/superagi/tools/file/delete_file.py
@@ -10,6 +10,7 @@
 from superagi.types.storage_types import StorageType
 from superagi.config.config import get_config
 from superagi.helper.s3_helper import S3Helper
+from superagi.models.resource import Resource
 
 
 class DeleteFileInput(BaseModel):
@@ -29,7 +30,7 @@ class DeleteFileTool(BaseTool):
     """
     name: str = "Delete File"
     agent_id: int = None
-    agent_execution_id:int = None
+    agent_execution_id: int = None
     args_schema: Type[BaseModel] = DeleteFileInput
     description: str = "Delete a file"
 
@@ -53,12 +54,14 @@ def _execute(self, file_name: str):
         if StorageType.get_storage_type(get_config("STORAGE_TYPE", StorageType.FILE.value)) == StorageType.S3:
             try:
                 S3Helper().delete_file(final_path)
+                Resource.delete_resource(self.toolkit_config.session, file_name, self.agent_id, self.agent_execution_id)
                 return "File deleted successfully."
             except Exception as err:
                 return f"Error: {err}"
         else:
             try:
                 os.remove(final_path)
+                Resource.delete_resource(self.toolkit_config.session, file_name, self.agent_id, self.agent_execution_id)
                 return "File deleted successfully."
             except Exception as err:
                 return f"Error: {err}"