Update app.py

app.py

@@ -6,6 +6,7 @@ import os
 import schedule
 import time
 import threading
+import mimetypes
 
 app = FastAPI()
 
@@ -71,10 +72,22 @@ async def download_file(upload_id: str):
         raise HTTPException(status_code=404, detail="Upload not found or has expired")
     
     upload_data = data_storage[upload_id]
-    if upload_data['data_type'] == "file":
-        return Response(content=upload_data['file_data'], media_type="application/octet-stream", headers={"Content-Disposition": f"attachment; filename={upload_data['file_name']}"})
-    elif upload_data['data_type'] == "text":
-        return Response(content=upload_data['file_data'], media_type="text/plain", headers={"Content-Disposition": f"attachment; filename={upload_data['file_name']}"})
+    file_name = upload_data['file_name']
+    
+    # Guess the MIME type based on the file extension
+    mime_type, _ = mimetypes.guess_type(file_name)
+    if mime_type is None:  # Default to octet-stream if type is unknown
+        mime_type = "application/octet-stream"
+    
+    # For text uploads, ensure the MIME type is set correctly
+    if upload_data['data_type'] == "text":
+        mime_type = "text/plain"
+        # Optionally, force a `.txt` extension for text files if no extension is provided
+        if '.' not in file_name:
+            file_name += '.txt'
+    
+    return Response(content=upload_data['file_data'], media_type=mime_type, 
+                    headers={"Content-Disposition": f"attachment; filename={file_name}"})
 
 if __name__ == "__main__":
     scheduled_function()