Update app.py

app.py

@@ -97,16 +97,27 @@ class DownloadResponse(BaseModel):
     file_size: int = None  # Include file size in the response
 
 
-def wait_for_file_to_complete(file_path, initial_size):
+
+
+def wait_for_file_to_complete(file_path, initial_size, max_wait_time=30):
     """
-    Wait until the file size stops increasing, indicating the download has completed.
+    Wait until the file size stops increasing for a set amount of time,
+    indicating the download has completed.
     """
+    wait_time = 0
     while True:
         current_size = os.path.getsize(file_path)
         if current_size == initial_size:
-            break
+            wait_time += 1
+            if wait_time >= max_wait_time:
+                # File size hasn't changed for max_wait_time seconds
+                break
+        else:
+            # Reset wait time if the file is still growing
+            wait_time = 0
         initial_size = current_size
-        time.sleep(1)  # Wait for 1 second before checking again
+        time.sleep(1)  # Check every second
+
 
 
 
@@ -160,7 +171,8 @@ async def download_file(request: DownloadRequest):
         logger.info(f"Downloaded file: {downloaded_file}")
         
         # Ensure the file is fully written (wait for download to complete)
-        wait_for_file_to_complete(downloaded_file, initial_size=os.path.getsize(downloaded_file))
+        initial_size = os.path.getsize(downloaded_file)
+        wait_for_file_to_complete(downloaded_file, initial_size)
         
         # Get the file size
         file_size = os.path.getsize(downloaded_file)