Update app.py

app.py

@@ -97,6 +97,19 @@ class DownloadResponse(BaseModel):
     file_size: int = None  # Include file size in the response
 
 
+def wait_for_file_to_complete(file_path, initial_size):
+    """
+    Wait until the file size stops increasing, indicating the download has completed.
+    """
+    while True:
+        current_size = os.path.getsize(file_path)
+        if current_size == initial_size:
+            break
+        initial_size = current_size
+        time.sleep(1)  # Wait for 1 second before checking again
+
+
+
 @app.post("/download", response_model=DownloadResponse)
 async def download_file(request: DownloadRequest):
     try:
@@ -142,14 +155,13 @@ async def download_file(request: DownloadRequest):
         
         if not files:
             raise Exception("No files found in download directory after download attempt")
-        # Ensure the file is fully written 
-
-        time.sleep(2) # Sleep for 2 second
+        
         downloaded_file = files[0]
         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))
         
-
         # Get the file size
         file_size = os.path.getsize(downloaded_file)
         
@@ -160,7 +172,7 @@ async def download_file(request: DownloadRequest):
         logger.info(f"Generated download URL: {download_url}")
         
         # Move back to original directory
-       
+        os.chdir(original_dir)
         
         return DownloadResponse(
             success=True,
@@ -169,7 +181,7 @@ async def download_file(request: DownloadRequest):
             download_url=download_url,
             file_size=file_size  # Include the file size in the response
         )
-            
+    
     except subprocess.CalledProcessError as e:
         logger.error(f"Download process failed: stdout={e.stdout}, stderr={e.stderr}")
         raise HTTPException(
@@ -189,6 +201,7 @@ async def download_file(request: DownloadRequest):
 
 
 
+
 @app.get("/test")
 async def test():
     """Test endpoint to verify setup"""