Repository: Chainlit/chainlit (language: python)

Issue:
Mixed dependencies in the Storage Clients module for the Data Layer
**Describe the bug**
When importing a certain storage client for usage in the data layer (e.g. `S3StorageClient`), one might encounter import errors due to mixed dependencies in the module, such as: **Azure** packages missing.

**To Reproduce**
Steps to reproduce the behavior:

1. Install Chainlit with relevant packages for using the `S3StorageClient` for the data layer (i.e. `boto3`)
2. Make sure you have no Azure packages installed in your Python runtime env. 
3. Try importing the client with `from chainlit.data.storage_clients import S3StorageClient`
4. You should encounter import errors because the `storage_clients` module has top-level imports for Azure:
```python
import boto3  # type: ignore
from azure.storage.filedatalake import (
    ContentSettings,
    DataLakeFileClient,
    DataLakeServiceClient,
    FileSystemClient,
)
```
**(Note the top-level import of both boto3 and azure.storage.filedatalake within the module, even though we only choose one client out of the two)**
6. Same issue would apply if you had Azure dependencies installed, but not S3 (`boto3`) 

**Expected behavior**
The storage client should be importable and ready for instantiation based on the type you require, without having to install all other storage clients' dependencies. 

**Desktop (please complete the following information):**

- OS: iOS
- Version: 1.2.0

---

I can imagine that splitting the `storage_clients` module to separate sub-modules based on the specific client type (e.g. `S3`, `Azure`, `GCS` etc.) would resolve this issue, as each would import **only** what that specific client needs. 



Discussion / hints:
> I can imagine that splitting the storage_clients module to separate sub-modules based on the specific client type (e.g. S3, Azure, GCS etc.) would resolve this issue, as each would import only what that specific client needs.

I have done this locally (creating a new script for only S3StorageClient and import it from there into your main app.py file) and it works like a charm.
I could do a PR to handle this maybe in the next days. 
> > I can imagine that splitting the storage_clients module to separate sub-modules based on the specific client type (e.g. S3, Azure, GCS etc.) would resolve this issue, as each would import only what that specific client needs.
> 
> I have done this locally (creating a new script for only S3StorageClient and import it from there into your main app.py file) and it works like a charm. I could do a PR to handle this maybe in the next days.

Yes, I have done so too and it works fine, but I imagine some people would not appreciate having to do that, or not realize that this is a solution. 

Bunching all the imports together in the same module for all the storage clients just seems like a simple oversight to me. I'm not sure why the data layers were properly distributed to modules (e.g. dynamodb in its own module) but the storage clients weren't.
> I can imagine that splitting the `storage_clients` module to separate sub-modules based on the specific client type (e.g. `S3`, `Azure`, `GCS` etc.) would resolve this issue, as each would import **only** what that specific client needs.

Let us know if you are doing a PR for this (to avoid double work). I imagine a `storage_clients` module in the `data` module with a `base.py` and files for each implementation (and, preferably, no code in `__init__.py`.

I love how active this community is! ❤️ 

P.S. 

> Bunching all the imports together in the same module for all the storage clients just seems like a simple oversight to me. I'm not sure why the data layers were properly distributed to modules (e.g. dynamodb in its own module) but the storage clients weren't.

One step at a time!
I will look into doing one tomorrow, unless someone beats me to it :)
Hi guys, I am trying doing the PR and I am working on Windows. 
I have set up everything but once running `pnpm run buildUi` I am encountering errors (tsup: command not found).

Should I open a new issue on this? Is Windows supported? 

> Hi guys, I am trying doing the PR and I am working on Windows. I have set up everything but once running `pnpm run buildUi` I am encountering errors (tsup: command not found).
> 
> Should I open a new issue on this? Is Windows supported?

Did you do `pnpm install` before running the build? 
@hadarsharon @ndricca No need to work both on the same PR, plenty of other bugs and cleanup left! :')

Fix patch:
diff --git a/backend/chainlit/data/base.py b/backend/chainlit/data/base.py
index d34eaaa899..d0df7187cd 100644
--- a/backend/chainlit/data/base.py
+++ b/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
diff --git a/backend/chainlit/data/dynamodb.py b/backend/chainlit/data/dynamodb.py
index b3524b0734..e067378b0e 100644
--- a/backend/chainlit/data/dynamodb.py
+++ b/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
diff --git a/backend/chainlit/data/sql_alchemy.py b/backend/chainlit/data/sql_alchemy.py
index f8ea7d7d52..79596b7359 100644
--- a/backend/chainlit/data/sql_alchemy.py
+++ b/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
diff --git a/backend/chainlit/data/storage_clients/__init__.py b/backend/chainlit/data/storage_clients/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/backend/chainlit/data/storage_clients.py b/backend/chainlit/data/storage_clients/azure.py
similarity index 69%
rename from backend/chainlit/data/storage_clients.py
rename to backend/chainlit/data/storage_clients/azure.py
index 42d5f5e40a..fac2164e8a 100644
--- a/backend/chainlit/data/storage_clients.py
+++ b/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 {}
diff --git a/backend/chainlit/data/storage_clients/base.py b/backend/chainlit/data/storage_clients/base.py
new file mode 100644
index 0000000000..2d65d68093
--- /dev/null
+++ b/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
diff --git a/backend/chainlit/data/storage_clients/s3.py b/backend/chainlit/data/storage_clients/s3.py
new file mode 100644
index 0000000000..4cfd6923c6
--- /dev/null
+++ b/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 {}
diff --git a/backend/poetry.lock b/backend/poetry.lock
index 4c717509ad..4ce5808c61 100644
--- a/backend/poetry.lock
+++ b/backend/poetry.lock
@@ -1855,6 +1855,23 @@ files = [
 [package.extras]
 colors = ["colorama (>=0.4.6)"]
 
+[[package]]
+name = "jinja2"
+version = "3.1.4"
+description = "A very fast and expressive template engine."
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "jinja2-3.1.4-py3-none-any.whl", hash = "sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d"},
+    {file = "jinja2-3.1.4.tar.gz", hash = "sha256:4a3aee7acbbe7303aede8e9648d13b8bf88a429282aa6122a993f0ac800cb369"},
+]
+
+[package.dependencies]
+MarkupSafe = ">=2.0"
+
+[package.extras]
+i18n = ["Babel (>=2.7)"]
+
 [[package]]
 name = "jiter"
 version = "0.5.0"
@@ -2613,6 +2630,75 @@ files = [
 [package.dependencies]
 llama-index-core = ">=0.10.29"
 
+[[package]]
+name = "markupsafe"
+version = "2.1.5"
+description = "Safely add untrusted strings to HTML/XML markup."
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a17a92de5231666cfbe003f0e4b9b3a7ae3afb1ec2845aadc2bacc93ff85febc"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:72b6be590cc35924b02c78ef34b467da4ba07e4e0f0454a2c5907f473fc50ce5"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e61659ba32cf2cf1481e575d0462554625196a1f2fc06a1c777d3f48e8865d46"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2174c595a0d73a3080ca3257b40096db99799265e1c27cc5a610743acd86d62f"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ae2ad8ae6ebee9d2d94b17fb62763125f3f374c25618198f40cbb8b525411900"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:075202fa5b72c86ad32dc7d0b56024ebdbcf2048c0ba09f1cde31bfdd57bcfff"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:598e3276b64aff0e7b3451b72e94fa3c238d452e7ddcd893c3ab324717456bad"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:fce659a462a1be54d2ffcacea5e3ba2d74daa74f30f5f143fe0c58636e355fdd"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-win32.whl", hash = "sha256:d9fad5155d72433c921b782e58892377c44bd6252b5af2f67f16b194987338a4"},
+    {file = "MarkupSafe-2.1.5-cp310-cp310-win_amd64.whl", hash = "sha256:bf50cd79a75d181c9181df03572cdce0fbb75cc353bc350712073108cba98de5"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:629ddd2ca402ae6dbedfceeba9c46d5f7b2a61d9749597d4307f943ef198fc1f"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5b7b716f97b52c5a14bffdf688f971b2d5ef4029127f1ad7a513973cfd818df2"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ec585f69cec0aa07d945b20805be741395e28ac1627333b1c5b0105962ffced"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b91c037585eba9095565a3556f611e3cbfaa42ca1e865f7b8015fe5c7336d5a5"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7502934a33b54030eaf1194c21c692a534196063db72176b0c4028e140f8f32c"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:0e397ac966fdf721b2c528cf028494e86172b4feba51d65f81ffd65c63798f3f"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:c061bb86a71b42465156a3ee7bd58c8c2ceacdbeb95d05a99893e08b8467359a"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:3a57fdd7ce31c7ff06cdfbf31dafa96cc533c21e443d57f5b1ecc6cdc668ec7f"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-win32.whl", hash = "sha256:397081c1a0bfb5124355710fe79478cdbeb39626492b15d399526ae53422b906"},
+    {file = "MarkupSafe-2.1.5-cp311-cp311-win_amd64.whl", hash = "sha256:2b7c57a4dfc4f16f7142221afe5ba4e093e09e728ca65c51f5620c9aaeb9a617"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:8dec4936e9c3100156f8a2dc89c4b88d5c435175ff03413b443469c7c8c5f4d1"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:3c6b973f22eb18a789b1460b4b91bf04ae3f0c4234a0a6aa6b0a92f6f7b951d4"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ac07bad82163452a6884fe8fa0963fb98c2346ba78d779ec06bd7a6262132aee"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f5dfb42c4604dddc8e4305050aa6deb084540643ed5804d7455b5df8fe16f5e5"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ea3d8a3d18833cf4304cd2fc9cbb1efe188ca9b5efef2bdac7adc20594a0e46b"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:d050b3361367a06d752db6ead6e7edeb0009be66bc3bae0ee9d97fb326badc2a"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:bec0a414d016ac1a18862a519e54b2fd0fc8bbfd6890376898a6c0891dd82e9f"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:58c98fee265677f63a4385256a6d7683ab1832f3ddd1e66fe948d5880c21a169"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-win32.whl", hash = "sha256:8590b4ae07a35970728874632fed7bd57b26b0102df2d2b233b6d9d82f6c62ad"},
+    {file = "MarkupSafe-2.1.5-cp312-cp312-win_amd64.whl", hash = "sha256:823b65d8706e32ad2df51ed89496147a42a2a6e01c13cfb6ffb8b1e92bc910bb"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:c8b29db45f8fe46ad280a7294f5c3ec36dbac9491f2d1c17345be8e69cc5928f"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ec6a563cff360b50eed26f13adc43e61bc0c04d94b8be985e6fb24b81f6dcfdf"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a549b9c31bec33820e885335b451286e2969a2d9e24879f83fe904a5ce59d70a"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4f11aa001c540f62c6166c7726f71f7573b52c68c31f014c25cc7901deea0b52"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:7b2e5a267c855eea6b4283940daa6e88a285f5f2a67f2220203786dfa59b37e9"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:2d2d793e36e230fd32babe143b04cec8a8b3eb8a3122d2aceb4a371e6b09b8df"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:ce409136744f6521e39fd8e2a24c53fa18ad67aa5bc7c2cf83645cce5b5c4e50"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-win32.whl", hash = "sha256:4096e9de5c6fdf43fb4f04c26fb114f61ef0bf2e5604b6ee3019d51b69e8c371"},
+    {file = "MarkupSafe-2.1.5-cp37-cp37m-win_amd64.whl", hash = "sha256:4275d846e41ecefa46e2015117a9f491e57a71ddd59bbead77e904dc02b1bed2"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:656f7526c69fac7f600bd1f400991cc282b417d17539a1b228617081106feb4a"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:97cafb1f3cbcd3fd2b6fbfb99ae11cdb14deea0736fc2b0952ee177f2b813a46"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f3fbcb7ef1f16e48246f704ab79d79da8a46891e2da03f8783a5b6fa41a9532"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa9db3f79de01457b03d4f01b34cf91bc0048eb2c3846ff26f66687c2f6d16ab"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ffee1f21e5ef0d712f9033568f8344d5da8cc2869dbd08d87c84656e6a2d2f68"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:5dedb4db619ba5a2787a94d877bc8ffc0566f92a01c0ef214865e54ecc9ee5e0"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:30b600cf0a7ac9234b2638fbc0fb6158ba5bdcdf46aeb631ead21248b9affbc4"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:8dd717634f5a044f860435c1d8c16a270ddf0ef8588d4887037c5028b859b0c3"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-win32.whl", hash = "sha256:daa4ee5a243f0f20d528d939d06670a298dd39b1ad5f8a72a4275124a7819eff"},
+    {file = "MarkupSafe-2.1.5-cp38-cp38-win_amd64.whl", hash = "sha256:619bc166c4f2de5caa5a633b8b7326fbe98e0ccbfacabd87268a2b15ff73a029"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7a68b554d356a91cce1236aa7682dc01df0edba8d043fd1ce607c49dd3c1edcf"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:db0b55e0f3cc0be60c1f19efdde9a637c32740486004f20d1cff53c3c0ece4d2"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3e53af139f8579a6d5f7b76549125f0d94d7e630761a2111bc431fd820e163b8"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:17b950fccb810b3293638215058e432159d2b71005c74371d784862b7e4683f3"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4c31f53cdae6ecfa91a77820e8b151dba54ab528ba65dfd235c80b086d68a465"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:bff1b4290a66b490a2f4719358c0cdcd9bafb6b8f061e45c7a2460866bf50c2e"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:bc1667f8b83f48511b94671e0e441401371dfd0f0a795c7daa4a3cd1dde55bea"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5049256f536511ee3f7e1b3f87d1d1209d327e818e6ae1365e8653d7e3abb6a6"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-win32.whl", hash = "sha256:00e046b6dd71aa03a41079792f8473dc494d564611a8f89bbbd7cb93295ebdcf"},
+    {file = "MarkupSafe-2.1.5-cp39-cp39-win_amd64.whl", hash = "sha256:fa173ec60341d6bb97a89f5ea19c85c5643c1e7dedebc22f5181eb73573142c5"},
+    {file = "MarkupSafe-2.1.5.tar.gz", hash = "sha256:d283d37a890ba4c1ae73ffadf8046435c76e7bc2247bbb63c00bd1a709c6544b"},
+]
+
 [[package]]
 name = "marshmallow"
 version = "3.22.0"
@@ -2718,6 +2804,51 @@ files = [
     {file = "more_itertools-10.5.0-py3-none-any.whl", hash = "sha256:037b0d3203ce90cca8ab1defbbdac29d5f993fc20131f3664dc8d6acfa872aef"},
 ]
 
+[[package]]
+name = "moto"
+version = "5.0.14"
+description = ""
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "moto-5.0.14-py2.py3-none-any.whl", hash = "sha256:c738ffe85d3844ef37b865951736c4faf2e0f3e4f05db87bdad97a6c01b88174"},
+    {file = "moto-5.0.14.tar.gz", hash = "sha256:0f849243269fd03372426c302b18cb605302da32620d7f0266be6a40735b2acd"},
+]
+
+[package.dependencies]
+boto3 = ">=1.9.201"
+botocore = ">=1.14.0"
+cryptography = ">=3.3.1"
+Jinja2 = ">=2.10.1"
+python-dateutil = ">=2.1,<3.0.0"
+requests = ">=2.5"
+responses = ">=0.15.0"
+werkzeug = ">=0.5,<2.2.0 || >2.2.0,<2.2.1 || >2.2.1"
+xmltodict = "*"
+
+[package.extras]
+all = ["PyYAML (>=5.1)", "antlr4-python3-runtime", "aws-xray-sdk (>=0.93,!=0.96)", "cfn-lint (>=0.40.0)", "docker (>=3.0.0)", "graphql-core", "joserfc (>=0.9.0)", "jsondiff (>=1.1.2)", "jsonpath-ng", "multipart", "openapi-spec-validator (>=0.5.0)", "py-partiql-parser (==0.5.6)", "pyparsing (>=3.0.7)", "setuptools"]
+apigateway = ["PyYAML (>=5.1)", "joserfc (>=0.9.0)", "openapi-spec-validator (>=0.5.0)"]
+apigatewayv2 = ["PyYAML (>=5.1)", "openapi-spec-validator (>=0.5.0)"]
+appsync = ["graphql-core"]
+awslambda = ["docker (>=3.0.0)"]
+batch = ["docker (>=3.0.0)"]
+cloudformation = ["PyYAML (>=5.1)", "aws-xray-sdk (>=0.93,!=0.96)", "cfn-lint (>=0.40.0)", "docker (>=3.0.0)", "graphql-core", "joserfc (>=0.9.0)", "jsondiff (>=1.1.2)", "openapi-spec-validator (>=0.5.0)", "py-partiql-parser (==0.5.6)", "pyparsing (>=3.0.7)", "setuptools"]
+cognitoidp = ["joserfc (>=0.9.0)"]
+dynamodb = ["docker (>=3.0.0)", "py-partiql-parser (==0.5.6)"]
+dynamodbstreams = ["docker (>=3.0.0)", "py-partiql-parser (==0.5.6)"]
+events = ["jsonpath-ng"]
+glue = ["pyparsing (>=3.0.7)"]
+iotdata = ["jsondiff (>=1.1.2)"]
+proxy = ["PyYAML (>=5.1)", "antlr4-python3-runtime", "aws-xray-sdk (>=0.93,!=0.96)", "cfn-lint (>=0.40.0)", "docker (>=2.5.1)", "graphql-core", "joserfc (>=0.9.0)", "jsondiff (>=1.1.2)", "jsonpath-ng", "multipart", "openapi-spec-validator (>=0.5.0)", "py-partiql-parser (==0.5.6)", "pyparsing (>=3.0.7)", "setuptools"]
+resourcegroupstaggingapi = ["PyYAML (>=5.1)", "cfn-lint (>=0.40.0)", "docker (>=3.0.0)", "graphql-core", "joserfc (>=0.9.0)", "jsondiff (>=1.1.2)", "openapi-spec-validator (>=0.5.0)", "py-partiql-parser (==0.5.6)", "pyparsing (>=3.0.7)"]
+s3 = ["PyYAML (>=5.1)", "py-partiql-parser (==0.5.6)"]
+s3crc32c = ["PyYAML (>=5.1)", "crc32c", "py-partiql-parser (==0.5.6)"]
+server = ["PyYAML (>=5.1)", "antlr4-python3-runtime", "aws-xray-sdk (>=0.93,!=0.96)", "cfn-lint (>=0.40.0)", "docker (>=3.0.0)", "flask (!=2.2.0,!=2.2.1)", "flask-cors", "graphql-core", "joserfc (>=0.9.0)", "jsondiff (>=1.1.2)", "jsonpath-ng", "openapi-spec-validator (>=0.5.0)", "py-partiql-parser (==0.5.6)", "pyparsing (>=3.0.7)", "setuptools"]
+ssm = ["PyYAML (>=5.1)"]
+stepfunctions = ["antlr4-python3-runtime", "jsonpath-ng"]
+xray = ["aws-xray-sdk (>=0.93,!=0.96)", "setuptools"]
+
 [[package]]
 name = "msal"
 version = "1.31.0"
@@ -4210,6 +4341,25 @@ requests = ">=2.0.0"
 [package.extras]
 rsa = ["oauthlib[signedtoken] (>=3.0.0)"]
 
+[[package]]
+name = "responses"
+version = "0.25.3"
+description = "A utility library for mocking out the `requests` Python library."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "responses-0.25.3-py3-none-any.whl", hash = "sha256:521efcbc82081ab8daa588e08f7e8a64ce79b91c39f6e62199b19159bea7dbcb"},
+    {file = "responses-0.25.3.tar.gz", hash = "sha256:617b9247abd9ae28313d57a75880422d55ec63c29d33d629697590a034358dba"},
+]
+
+[package.dependencies]
+pyyaml = "*"
+requests = ">=2.30.0,<3.0"
+urllib3 = ">=1.25.10,<3.0"
+
+[package.extras]
+tests = ["coverage (>=6.0.0)", "flake8", "mypy", "pytest (>=7.0.0)", "pytest-asyncio", "pytest-cov", "pytest-httpserver", "tomli", "tomli-w", "types-PyYAML", "types-requests"]
+
 [[package]]
 name = "rpds-py"
 version = "0.20.0"
@@ -5320,6 +5470,23 @@ files = [
 [package.dependencies]
 anyio = ">=3.0.0"
 
+[[package]]
+name = "werkzeug"
+version = "3.0.4"
+description = "The comprehensive WSGI web application library."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "werkzeug-3.0.4-py3-none-any.whl", hash = "sha256:02c9eb92b7d6c06f31a782811505d2157837cea66aaede3e217c7c27c039476c"},
+    {file = "werkzeug-3.0.4.tar.gz", hash = "sha256:34f2371506b250df4d4f84bfe7b0921e4762525762bbd936614909fe25cd7306"},
+]
+
+[package.dependencies]
+MarkupSafe = ">=2.1.1"
+
+[package.extras]
+watchdog = ["watchdog (>=2.3)"]
+
 [[package]]
 name = "wrapt"
 version = "1.16.0"
@@ -5413,6 +5580,17 @@ files = [
 [package.dependencies]
 h11 = ">=0.9.0,<1"
 
+[[package]]
+name = "xmltodict"
+version = "0.13.0"
+description = "Makes working with XML feel like you are working with JSON"
+optional = false
+python-versions = ">=3.4"
+files = [
+    {file = "xmltodict-0.13.0-py2.py3-none-any.whl", hash = "sha256:aa89e8fd76320154a40d19a0df04a4695fb9dc5ba977cbb68ab3e4eb225e7852"},
+    {file = "xmltodict-0.13.0.tar.gz", hash = "sha256:341595a488e3e01a85a9d8911d8912fd922ede5fecc4dce437eb4b6c8d037e56"},
+]
+
 [[package]]
 name = "yarl"
 version = "1.11.1"
@@ -5540,4 +5718,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9,<4.0.0"
-content-hash = "ef9341345f921f6b78cccbcf94fd539d9d0814428d40de0ade9a244362616e04"
+content-hash = "eda02c1c17a62c92406f37a9a5a81cbebf039f63bf3200258122ac564ec47de9"
diff --git a/backend/pyproject.toml b/backend/pyproject.toml
index 64c1b61e63..9f22ab9cf8 100644
--- a/backend/pyproject.toml
+++ b/backend/pyproject.toml
@@ -80,6 +80,7 @@ slack_bolt = "^1.18.1"
 discord = "^2.3.2"
 botbuilder-core = "^4.15.0"
 aiosqlite = "^0.20.0"
+moto = "^5.0.14"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.8.0"
@@ -119,6 +120,7 @@ ignore_missing_imports = true
 
 
 
+
 [tool.poetry.group.custom-data]
 optional = true
 


Test patch:
diff --git a/backend/tests/data/conftest.py b/backend/tests/data/conftest.py
index 19177756fc..e1fe80faf7 100644
--- a/backend/tests/data/conftest.py
+++ b/backend/tests/data/conftest.py
@@ -2,7 +2,7 @@
 
 from unittest.mock import AsyncMock
 
-from chainlit.data.base import BaseStorageClient
+from chainlit.data.storage_clients.base import BaseStorageClient
 from chainlit.user import User
 
 
diff --git a/backend/tests/data/storage_clients/test_s3.py b/backend/tests/data/storage_clients/test_s3.py
new file mode 100644
index 0000000000..38804f8c67
--- /dev/null
+++ b/backend/tests/data/storage_clients/test_s3.py
@@ -0,0 +1,46 @@
+import os
+
+import boto3  # type: ignore
+import pytest
+from chainlit.data.storage_clients.s3 import S3StorageClient
+from moto import mock_aws
+
+
+# Fixtures for setting up the DynamoDB table
+@pytest.fixture
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+
+
+@pytest.fixture
+def s3_mock(aws_credentials):
+    """Moto mock S3 setup."""
+    with mock_aws():
+        s3 = boto3.client("s3", region_name="us-east-1")
+        # Create a mock bucket
+        s3.create_bucket(Bucket="my-test-bucket")
+        yield s3
+
+
+@pytest.mark.asyncio
+async def test_upload_file(s3_mock):
+    # Initialize the S3StorageClient with the mock bucket
+    client = S3StorageClient(bucket="my-test-bucket")
+
+    # Call the upload_file method and await the result
+    result = await client.upload_file(
+        object_key="test.txt", data="This is a test file", mime="text/plain"
+    )
+
+    # Assert that the file upload returned the correct URL
+    assert result["object_key"] == "test.txt"
+    assert result["url"] == "https://my-test-bucket.s3.amazonaws.com/test.txt"
+
+    # Verify that the file exists in the mock S3
+    response = s3_mock.get_object(Bucket="my-test-bucket", Key="test.txt")
+    assert response["Body"].read().decode() == "This is a test file"
diff --git a/backend/tests/data/test_sql_alchemy.py b/backend/tests/data/test_sql_alchemy.py
index 49b2619181..b94f174512 100644
--- a/backend/tests/data/test_sql_alchemy.py
+++ b/backend/tests/data/test_sql_alchemy.py
@@ -6,7 +6,7 @@
 from sqlalchemy.ext.asyncio import create_async_engine
 
 from chainlit import User
-from chainlit.data.base import BaseStorageClient
+from chainlit.data.storage_clients.base import BaseStorageClient
 from chainlit.data.sql_alchemy import SQLAlchemyDataLayer
 from chainlit.element import Text
 
diff --git a/cypress/e2e/custom_data_layer/sql_alchemy.py b/cypress/e2e/custom_data_layer/sql_alchemy.py
index 8f06b4bcf3..430b02a58e 100644
--- a/cypress/e2e/custom_data_layer/sql_alchemy.py
+++ b/cypress/e2e/custom_data_layer/sql_alchemy.py
@@ -2,7 +2,7 @@
 
 import chainlit.data as cl_data
 from chainlit.data.sql_alchemy import SQLAlchemyDataLayer
-from chainlit.data.storage_clients import AzureStorageClient
+from chainlit.data.storage_clients.azure import AzureStorageClient
 
 import chainlit as cl