Factor storage clients into separate modules. Closes #1355. (#1363)

* Refactor storage_clients.py from a script to a module inside chainlit.data
* Add pytest on S3StorageClient + add moto in tests dependencies
* Moved BaseStorageClient from from chainlit.data.base.py to chainlit.data.storage_clients.base.py

Co-authored-by: jl1andricca <[email protected]>

backend/chainlit/data/base.py

@@ -105,17 +105,3 @@ async def update_thread(
     @abstractmethod
     async def build_debug_url(self) -> str:
         pass
-
-
-class BaseStorageClient(ABC):
-    """Base class for non-text data persistence like Azure Data Lake, S3, Google Storage, etc."""
-
-    @abstractmethod
-    async def upload_file(
-        self,
-        object_key: str,
-        data: Union[bytes, str],
-        mime: str = "application/octet-stream",
-        overwrite: bool = True,
-    ) -> Dict[str, Any]:
-        pass

backend/chainlit/data/dynamodb.py

@@ -12,7 +12,8 @@
 import boto3  # type: ignore
 from boto3.dynamodb.types import TypeDeserializer, TypeSerializer
 from chainlit.context import context
-from chainlit.data.base import BaseDataLayer, BaseStorageClient
+from chainlit.data.base import BaseDataLayer
+from chainlit.data.storage_clients.base import BaseStorageClient
 from chainlit.data.utils import queue_until_user_message
 from chainlit.element import ElementDict
 from chainlit.logger import logger

backend/chainlit/data/sql_alchemy.py

@@ -8,7 +8,8 @@
 import aiofiles
 import aiohttp
 
-from chainlit.data.base import BaseDataLayer, BaseStorageClient
+from chainlit.data.base import BaseDataLayer
+from chainlit.data.storage_clients.base import BaseStorageClient
 from chainlit.data.utils import queue_until_user_message
 from chainlit.element import ElementDict
 from chainlit.logger import logger

backend/chainlit/data/storage_clients.py → backend/chainlit/data/storage_clients/azure.py

@@ -1,13 +1,12 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
-import boto3  # type: ignore
 from azure.storage.filedatalake import (
     ContentSettings,
     DataLakeFileClient,
     DataLakeServiceClient,
     FileSystemClient,
 )
-from chainlit.data.base import BaseStorageClient
+from chainlit.data.storage_clients.base import BaseStorageClient
 from chainlit.logger import logger
 
 if TYPE_CHECKING:
@@ -79,34 +78,3 @@ async def upload_file(
         except Exception as e:
             logger.warn(f"AzureStorageClient, upload_file error: {e}")
             return {}
-
-
-class S3StorageClient(BaseStorageClient):
-    """
-    Class to enable Amazon S3 storage provider
-    """
-
-    def __init__(self, bucket: str):
-        try:
-            self.bucket = bucket
-            self.client = boto3.client("s3")
-            logger.info("S3StorageClient initialized")
-        except Exception as e:
-            logger.warn(f"S3StorageClient initialization error: {e}")
-
-    async def upload_file(
-        self,
-        object_key: str,
-        data: Union[bytes, str],
-        mime: str = "application/octet-stream",
-        overwrite: bool = True,
-    ) -> Dict[str, Any]:
-        try:
-            self.client.put_object(
-                Bucket=self.bucket, Key=object_key, Body=data, ContentType=mime
-            )
-            url = f"https://{self.bucket}.s3.amazonaws.com/{object_key}"
-            return {"object_key": object_key, "url": url}
-        except Exception as e:
-            logger.warn(f"S3StorageClient, upload_file error: {e}")
-            return {}

backend/chainlit/data/storage_clients/base.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Union
+
+
+class BaseStorageClient(ABC):
+    """Base class for non-text data persistence like Azure Data Lake, S3, Google Storage, etc."""
+
+    @abstractmethod
+    async def upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        pass

backend/chainlit/data/storage_clients/s3.py

@@ -0,0 +1,36 @@
+from typing import Any, Dict, Union
+
+import boto3  # type: ignore
+from chainlit.data.storage_clients.base import BaseStorageClient
+from chainlit.logger import logger
+
+
+class S3StorageClient(BaseStorageClient):
+    """
+    Class to enable Amazon S3 storage provider
+    """
+
+    def __init__(self, bucket: str):
+        try:
+            self.bucket = bucket
+            self.client = boto3.client("s3")
+            logger.info("S3StorageClient initialized")
+        except Exception as e:
+            logger.warn(f"S3StorageClient initialization error: {e}")
+
+    async def upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        try:
+            self.client.put_object(
+                Bucket=self.bucket, Key=object_key, Body=data, ContentType=mime
+            )
+            url = f"https://{self.bucket}.s3.amazonaws.com/{object_key}"
+            return {"object_key": object_key, "url": url}
+        except Exception as e:
+            logger.warn(f"S3StorageClient, upload_file error: {e}")
+            return {}