From ab1c3f0b85922162ce9feb705bf27bdbe946491b Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Fri, 26 Jan 2024 22:40:58 -0300
Subject: [PATCH] Add constants and utility function for file storage

---
 .../langflow/services/storage/constants.py    | 28 +++++++++++++++++++
 .../langflow/services/storage/service.py      |  8 +++---
 .../langflow/services/storage/utils.py        |  5 ++++
 3 files changed, 37 insertions(+), 4 deletions(-)
 create mode 100644 src/backend/langflow/services/storage/constants.py
 create mode 100644 src/backend/langflow/services/storage/utils.py

diff --git a/src/backend/langflow/services/storage/constants.py b/src/backend/langflow/services/storage/constants.py
new file mode 100644
index 000000000..ddec29f3f
--- /dev/null
+++ b/src/backend/langflow/services/storage/constants.py
@@ -0,0 +1,28 @@
+EXTENSION_TO_CONTENT_TYPE = {
+    "json": "application/json",
+    "txt": "text/plain",
+    "csv": "text/csv",
+    "html": "text/html",
+    "pdf": "application/pdf",
+    "png": "image/png",
+    "jpg": "image/jpeg",
+    "jpeg": "image/jpeg",
+    "gif": "image/gif",
+    "svg": "image/svg+xml",
+    "mp3": "audio/mpeg",
+    "wav": "audio/wav",
+    "mp4": "video/mp4",
+    "webm": "video/webm",
+    "zip": "application/zip",
+    "tar": "application/x-tar",
+    "gz": "application/gzip",
+    "doc": "application/msword",
+    "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+    "xls": "application/vnd.ms-excel",
+    "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+    "ppt": "application/vnd.ms-powerpoint",
+    "pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
+    "xml": "application/xml",
+    "yaml": "application/x-yaml",
+    "yml": "application/x-yaml",
+}
diff --git a/src/backend/langflow/services/storage/service.py b/src/backend/langflow/services/storage/service.py
index 4e34c05bd..ee667ddbe 100644
--- a/src/backend/langflow/services/storage/service.py
+++ b/src/backend/langflow/services/storage/service.py
@@ -18,19 +18,19 @@ class StorageService(Service):
         self.ready = True
 
     @abstractmethod
-    def save_file(self, folder: str, file_name: str, data):
+    def save_file(self, folder: str, file_name: str, data) -> None:
         pass
 
     @abstractmethod
-    def get_file(self, folder: str, file_name: str):
+    def get_file(self, folder: str, file_name: str) -> bytes:
         pass
 
     @abstractmethod
-    def list_files(self, folder: str):
+    def list_files(self, folder: str) -> list[str]:
         pass
 
     @abstractmethod
-    def delete_file(self, folder: str, file_name: str):
+    def delete_file(self, folder: str, file_name: str) -> bool:
         pass
 
     def teardown(self):
diff --git a/src/backend/langflow/services/storage/utils.py b/src/backend/langflow/services/storage/utils.py
new file mode 100644
index 000000000..a5515171b
--- /dev/null
+++ b/src/backend/langflow/services/storage/utils.py
@@ -0,0 +1,5 @@
+from langflow.services.storage.constants import EXTENSION_TO_CONTENT_TYPE
+
+
+def build_content_type_from_extension(extension: str):
+    return EXTENSION_TO_CONTENT_TYPE.get(extension.lower(), "application/octet-stream")