Refactor storage service constructors

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-27 23:16:26 -03:00
commit 8f7beb10b2
3 changed files with 17 additions and 17 deletions

View file

@ -1,33 +1,31 @@
from pathlib import Path
from platformdirs import user_data_dir
from .service import StorageService
class LocalStorageService(StorageService):
def __init__(self, session_service):
super().__init__(session_service)
self.data_dir = user_data_dir("langflow", "langflow")
def __init__(self, session_service, settings_service):
super().__init__(session_service, settings_service)
self.data_dir = settings_service.settings.CONFIG_DIR
self.set_ready()
def save_file(self, folder: str, file_name: str, data: bytes):
folder_path = Path(f"{self.data_dir}/{folder}")
def save_file(self, flow_id: str, file_name: str, data: bytes):
folder_path = Path(f"{self.data_dir}/{flow_id}")
folder_path.mkdir(parents=True, exist_ok=True)
with open(f"{self.data_dir}/{folder}/{file_name}", "wb") as f:
with open(folder_path / file_name, "wb") as f:
f.write(data)
def get_file(self, folder: str, file_name: str) -> bytes:
with open(f"{self.data_dir}/{folder}/{file_name}", "rb") as f:
def get_file(self, flow_id: str, file_name: str) -> bytes:
with open(f"{self.data_dir}/{flow_id}/{file_name}", "rb") as f:
return f.read()
def list_files(self, folder: str):
folder_path = Path(f"{self.data_dir}/{folder}")
def list_files(self, flow_id: str):
folder_path = Path(f"{self.data_dir}/{flow_id}")
return [file.name for file in folder_path.iterdir() if file.is_file()]
def delete_file(self, folder: str, file_name: str):
Path(f"{self.data_dir}/{folder}/{file_name}").unlink()
def delete_file(self, flow_id: str, file_name: str):
Path(f"{self.data_dir}/{flow_id}/{file_name}").unlink()
def teardown(self):
pass

View file

@ -5,8 +5,8 @@ from .service import StorageService
class S3StorageService(StorageService):
def __init__(self, session_service):
super().__init__(session_service)
def __init__(self, session_service, settings_service):
super().__init__(session_service, settings_service)
self.bucket = "langflow"
self.s3_client = boto3.client("s3")
self.set_ready()

View file

@ -5,12 +5,14 @@ from langflow.services.base import Service
if TYPE_CHECKING:
from langflow.services.session.service import SessionService
from langflow.services.settings.service import SettingsService
class StorageService(Service):
name = "storage_service"
def __init__(self, session_service: "SessionService"):
def __init__(self, session_service: "SessionService", settings_service: "SettingsService"):
self.settings_service = settings_service
self.session_service = session_service
self.set_ready()