From de4262faa667b5d975d480e4413e6f4f9f77d375 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 25 Sep 2023 18:59:21 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(endpoints.py):=20fix=20argum?=
 =?UTF-8?q?ent=20passed=20to=20save=5Fuploaded=5Ffile=20function=20to=20co?=
 =?UTF-8?q?rrectly=20save=20the=20uploaded=20file=20=F0=9F=90=9B=20fix(uti?=
 =?UTF-8?q?ls.py):=20fix=20argument=20type=20in=20save=5Fuploaded=5Ffile?=
 =?UTF-8?q?=20function=20to=20correctly=20handle=20UploadFile=20object?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/api/v1/endpoints.py     |  2 +-
 src/backend/langflow/services/cache/utils.py | 15 +++++++++------
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/backend/langflow/api/v1/endpoints.py b/src/backend/langflow/api/v1/endpoints.py
index 870f91d28..0d1ad751f 100644
--- a/src/backend/langflow/api/v1/endpoints.py
+++ b/src/backend/langflow/api/v1/endpoints.py
@@ -159,7 +159,7 @@ async def process_flow(
 async def create_upload_file(file: UploadFile, flow_id: str):
     # Cache file
     try:
-        file_path = save_uploaded_file(file.file, folder_name=flow_id)
+        file_path = save_uploaded_file(file, folder_name=flow_id)
 
         return UploadFileResponse(
             flowId=flow_id,
diff --git a/src/backend/langflow/services/cache/utils.py b/src/backend/langflow/services/cache/utils.py
index a36243b75..4342914cd 100644
--- a/src/backend/langflow/services/cache/utils.py
+++ b/src/backend/langflow/services/cache/utils.py
@@ -8,6 +8,7 @@ from collections import OrderedDict
 from pathlib import Path
 from typing import Any, Dict
 from appdirs import user_cache_dir
+from fastapi import UploadFile
 from langflow.services.database.models.base import orjson_dumps
 
 CACHE: Dict[str, Any] = {}
@@ -152,7 +153,7 @@ def save_binary_file(content: str, file_name: str, accepted_types: list[str]) ->
 
 
 @create_cache_folder
-def save_uploaded_file(file, folder_name):
+def save_uploaded_file(file: UploadFile, folder_name):
     """
     Save an uploaded file to the specified folder with a hash of its content as the file name.
 
@@ -165,6 +166,8 @@ def save_uploaded_file(file, folder_name):
     """
     cache_path = Path(CACHE_DIR)
     folder_path = cache_path / folder_name
+    file_extension = Path(file.filename).suffix
+    file_object = file.file
 
     # Create the folder if it doesn't exist
     if not folder_path.exists():
@@ -173,22 +176,22 @@ def save_uploaded_file(file, folder_name):
     # Create a hash of the file content
     sha256_hash = hashlib.sha256()
     # Reset the file cursor to the beginning of the file
-    file.seek(0)
+    file_object.seek(0)
     # Iterate over the uploaded file in small chunks to conserve memory
-    while chunk := file.read(8192):  # Read 8KB at a time (adjust as needed)
+    while chunk := file_object.read(8192):  # Read 8KB at a time (adjust as needed)
         sha256_hash.update(chunk)
 
     # Use the hex digest of the hash as the file name
     hex_dig = sha256_hash.hexdigest()
-    file_name = hex_dig
+    file_name = f"{hex_dig}{file_extension}"
 
     # Reset the file cursor to the beginning of the file
-    file.seek(0)
+    file_object.seek(0)
 
     # Save the file with the hash as its name
     file_path = folder_path / file_name
     with open(file_path, "wb") as new_file:
-        while chunk := file.read(8192):
+        while chunk := file_object.read(8192):
             new_file.write(chunk)
 
     return file_path