Add storage service and factory classes
This commit is contained in:
parent
1a4362fe50
commit
e9fc9f631b
5 changed files with 144 additions and 0 deletions
0
src/backend/langflow/services/storage/__init__.py
Normal file
0
src/backend/langflow/services/storage/__init__.py
Normal file
31
src/backend/langflow/services/storage/factory.py
Normal file
31
src/backend/langflow/services/storage/factory.py
Normal file
|
|
@ -0,0 +1,31 @@
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from langflow.services.factory import ServiceFactory
|
||||||
|
from langflow.services.storage.service import StorageService
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from langflow.services.session.service import SessionService
|
||||||
|
from langflow.services.settings.service import SettingsService
|
||||||
|
|
||||||
|
|
||||||
|
class StorageServiceFactory(ServiceFactory):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__(StorageService)
|
||||||
|
|
||||||
|
def create(self, session_service: "SessionService", settings_service: "SettingsService"):
|
||||||
|
storage_type = settings_service.settings.STORAGE_TYPE
|
||||||
|
if storage_type.lower() == "local":
|
||||||
|
from .local import LocalStorageService
|
||||||
|
|
||||||
|
return LocalStorageService(session_service)
|
||||||
|
elif storage_type.lower() == "s3":
|
||||||
|
from .s3 import S3StorageService
|
||||||
|
|
||||||
|
return S3StorageService(session_service)
|
||||||
|
else:
|
||||||
|
logger.warning(f"Storage type {storage_type} not supported. Using local storage.")
|
||||||
|
from .local import LocalStorageService
|
||||||
|
|
||||||
|
return LocalStorageService(session_service)
|
||||||
33
src/backend/langflow/services/storage/local.py
Normal file
33
src/backend/langflow/services/storage/local.py
Normal file
|
|
@ -0,0 +1,33 @@
|
||||||
|
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")
|
||||||
|
|
||||||
|
self.set_ready()
|
||||||
|
|
||||||
|
def save_file(self, folder: str, file_name: str, data):
|
||||||
|
folder_path = Path(f"{self.data_dir}/{folder}")
|
||||||
|
folder_path.mkdir(parents=True, exist_ok=True)
|
||||||
|
with open(f"{self.data_dir}/{folder}/{file_name}", "w") as f:
|
||||||
|
f.write(data)
|
||||||
|
|
||||||
|
def get_file(self, folder: str, file_name: str):
|
||||||
|
with open(f"{self.data_dir}/{folder}/{file_name}", "r") as f:
|
||||||
|
return f.read()
|
||||||
|
|
||||||
|
def list_files(self, folder: str):
|
||||||
|
folder_path = Path(f"{self.data_dir}/{folder}")
|
||||||
|
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 teardown(self):
|
||||||
|
pass
|
||||||
43
src/backend/langflow/services/storage/s3.py
Normal file
43
src/backend/langflow/services/storage/s3.py
Normal file
|
|
@ -0,0 +1,43 @@
|
||||||
|
import boto3
|
||||||
|
from botocore.exceptions import ClientError, NoCredentialsError
|
||||||
|
|
||||||
|
from .service import StorageService
|
||||||
|
|
||||||
|
|
||||||
|
class S3StorageService(StorageService):
|
||||||
|
def __init__(self, session_service):
|
||||||
|
super().__init__(session_service)
|
||||||
|
self.bucket = "langflow"
|
||||||
|
self.s3_client = boto3.client("s3")
|
||||||
|
self.set_ready()
|
||||||
|
|
||||||
|
def save_file(self, folder: str, file_name: str, data):
|
||||||
|
try:
|
||||||
|
self.s3_client.put_object(Bucket=self.bucket, Key=f"{folder}/{file_name}", Body=data)
|
||||||
|
except NoCredentialsError:
|
||||||
|
raise Exception("Credentials not available for AWS S3.")
|
||||||
|
except ClientError as e:
|
||||||
|
raise Exception(f"An error occurred: {e}")
|
||||||
|
|
||||||
|
def get_file(self, folder: str, file_name: str):
|
||||||
|
try:
|
||||||
|
response = self.s3_client.get_object(Bucket=self.bucket, Key=f"{folder}/{file_name}")
|
||||||
|
return response["Body"].read()
|
||||||
|
except ClientError as e:
|
||||||
|
raise Exception(f"An error occurred: {e}")
|
||||||
|
|
||||||
|
def list_files(self, folder: str):
|
||||||
|
try:
|
||||||
|
response = self.s3_client.list_objects_v2(Bucket=self.bucket, Prefix=folder)
|
||||||
|
return [item["Key"] for item in response.get("Contents", []) if "/" not in item["Key"][len(folder) :]]
|
||||||
|
except ClientError as e:
|
||||||
|
raise Exception(f"An error occurred: {e}")
|
||||||
|
|
||||||
|
def delete_file(self, folder: str, file_name: str):
|
||||||
|
try:
|
||||||
|
self.s3_client.delete_object(Bucket=self.bucket, Key=f"{folder}/{file_name}")
|
||||||
|
except ClientError as e:
|
||||||
|
raise Exception(f"An error occurred: {e}")
|
||||||
|
|
||||||
|
def teardown(self):
|
||||||
|
pass
|
||||||
37
src/backend/langflow/services/storage/service.py
Normal file
37
src/backend/langflow/services/storage/service.py
Normal file
|
|
@ -0,0 +1,37 @@
|
||||||
|
from abc import abstractmethod
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from langflow.services.base import Service
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from langflow.services.session.service import SessionService
|
||||||
|
|
||||||
|
|
||||||
|
class StorageService(Service):
|
||||||
|
name = "storage_service"
|
||||||
|
|
||||||
|
def __init__(self, session_service: "SessionService"):
|
||||||
|
self.session_service = session_service
|
||||||
|
self.set_ready()
|
||||||
|
|
||||||
|
def set_ready(self):
|
||||||
|
self.ready = True
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def save_file(self, folder: str, file_name: str, data):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def get_file(self, folder: str, file_name: str):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def list_files(self, folder: str):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def delete_file(self, folder: str, file_name: str):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def teardown(self):
|
||||||
|
pass
|
||||||
Loading…
Add table
Add a link
Reference in a new issue