aboutsummaryrefslogtreecommitdiff
path: root/autogpts/forge/forge/file_storage/__init__.py
blob: 9d5df553e2aaf5b61d4564c65e82346311a64421 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import enum
from pathlib import Path

from .base import FileStorage, FileStorageConfiguration
from .gcs import GCSFileStorage, GCSFileStorageConfiguration
from .local import LocalFileStorage
from .s3 import S3FileStorage, S3FileStorageConfiguration


class FileStorageBackendName(str, enum.Enum):
    LOCAL = "local"
    GCS = "gcs"
    S3 = "s3"


def get_storage(
    backend: FileStorageBackendName,
    root_path: Path = ".",
    restrict_to_root: bool = True,
) -> FileStorage:
    match backend:
        case FileStorageBackendName.LOCAL:
            from .local import FileStorageConfiguration, LocalFileStorage

            config = FileStorageConfiguration.from_env()
            config.root = root_path
            config.restrict_to_root = restrict_to_root
            return LocalFileStorage(config)
        case FileStorageBackendName.S3:
            from .s3 import S3FileStorage, S3FileStorageConfiguration

            config = S3FileStorageConfiguration.from_env()
            config.root = root_path
            return S3FileStorage(config)
        case FileStorageBackendName.GCS:
            from .gcs import GCSFileStorage, GCSFileStorageConfiguration

            config = GCSFileStorageConfiguration.from_env()
            config.root = root_path
            return GCSFileStorage(config)