diff --git a/superagi/models/resource.py b/superagi/models/resource.py index 78713b690..1a519c012 100644 --- 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 index 3e6128b22..d88a452e3 100644 --- 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}" \ No newline at end of file