Update app.py

app.py

@@ -1,99 +1,204 @@
 import os
 import shutil
 import zipfile
-import asyncio
+import json
+import time
 from contextlib import asynccontextmanager
-from typing import List
+from typing import List, Dict, Any
 
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import FileResponse
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, hf_hub_download, HfFileSystem
 
 # --- Configuration ---
 UPLOAD_DIR = "uploaded_files"
+STATE_FILE_PATH = os.path.join(os.getcwd(), "processing_audio_state.json")
 HF_DATASET_REPO = "samfred2/A_Text"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# --- Utility Functions ---
+# Hardcoded variables for automation
+ZIP_UPLOAD_THRESHOLD = 200
+STATE_UPLOAD_THRESHOLD = 100
 
-def get_uploaded_files() -> List[str]:
-    """Returns a list of all files in the upload directory."""
-    if not os.path.exists(UPLOAD_DIR):
-        return []
-    return [os.path.join(UPLOAD_DIR, f) for f in os.listdir(UPLOAD_DIR) if os.path.isfile(os.path.join(UPLOAD_DIR, f))]
+# Global state object (will be loaded from/saved to STATE_FILE_PATH)
+app_state: Dict[str, Any] = {
+    "total_files_uploaded": 0,
+    "current_zip_version": 0,
+    "files_since_last_zip": 0,
+    "files_since_last_state_upload": 0,
+}
+
+# --- Hugging Face Utility Functions ---
+
+def get_hf_api() -> HfApi:
+    """Returns an HfApi instance, raising an error if token is missing."""
+    if not HF_TOKEN:
+        raise ValueError("HF_TOKEN not found in environment variables.")
+    return HfApi(token=HF_TOKEN)
+
+def download_state_file():
+    """Downloads the state file from the Hugging Face dataset."""
+    try:
+        api = get_hf_api()
+        # Use hf_hub_download for simple file download
+        downloaded_path = hf_hub_download(
+            repo_id=HF_DATASET_REPO,
+            filename=os.path.basename(STATE_FILE_PATH),
+            repo_type="dataset",
+            local_dir=os.path.dirname(STATE_FILE_PATH),
+            local_dir_use_symlinks=False
+        )
+        print(f"Successfully downloaded state file to: {downloaded_path}")
+        return True
+    except ValueError:
+        print("HF_TOKEN missing. Cannot download state file.")
+        return False
+    except Exception as e:
+        # File not found is a common case for the first run
+        if "404" in str(e):
+            print("State file not found on Hugging Face. Will start with default state.")
+        else:
+            print(f"Error downloading state file: {e}")
+        return False
 
-def zip_uploaded_files(zip_filename: str = "uploaded_files.zip") -> str:
-    """Zips all files in the upload directory into a single zip file."""
+def upload_file_to_huggingface(local_path: str, path_in_repo: str):
+    """Uploads a single file to the specified Hugging Face dataset."""
+    try:
+        api = get_hf_api()
+        api.upload_file(
+            path_or_fileobj=local_path,
+            path_in_repo=path_in_repo,
+            repo_id=HF_DATASET_REPO,
+            repo_type="dataset",
+            token=HF_TOKEN
+        )
+        print(f"Successfully uploaded {os.path.basename(local_path)} to {HF_DATASET_REPO} as {path_in_repo}")
+        return True
+    except ValueError:
+        print("HF_TOKEN missing. Skipping upload.")
+        return False
+    except Exception as e:
+        print(f"Hugging Face upload failed for {local_path}: {e}")
+        return False
+
+# --- State Management Functions ---
+
+def load_state():
+    """Loads application state from the local state file."""
+    global app_state
+    if os.path.exists(STATE_FILE_PATH):
+        try:
+            with open(STATE_FILE_PATH, "r") as f:
+                app_state.update(json.load(f))
+            print(f"State loaded successfully from {STATE_FILE_PATH}.")
+        except Exception as e:
+            print(f"Error loading state file: {e}. Using default state.")
+    else:
+        print("Local state file not found. Using default state.")
+
+def save_state():
+    """Saves application state to the local state file."""
+    try:
+        with open(STATE_FILE_PATH, "w") as f:
+            json.dump(app_state, f, indent=4)
+        print(f"State saved successfully to {STATE_FILE_PATH}.")
+    except Exception as e:
+        print(f"Error saving state file: {e}")
+
+def upload_state_file_to_hf():
+    """Saves the current state and uploads the state file to Hugging Face."""
+    save_state()
+    upload_file_to_huggingface(STATE_FILE_PATH, os.path.basename(STATE_FILE_PATH))
+    global app_state
+    app_state["files_since_last_state_upload"] = 0
+
+# --- File Processing Functions ---
+
+def zip_uploaded_files_versioned() -> str:
+    """Zips all files in the upload directory (excluding state file) into a versioned zip file."""
     if not os.path.exists(UPLOAD_DIR) or not os.listdir(UPLOAD_DIR):
         print("No files to zip.")
         return None
 
+    # Increment version and create filename
+    global app_state
+    app_state["current_zip_version"] += 1
+    zip_filename = f"uploaded_files_{app_state['current_zip_version']}.zip"
     zip_path = os.path.join(os.getcwd(), zip_filename)
+    
+    exclude_file = os.path.basename(STATE_FILE_PATH)
+
     with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
         for root, _, files in os.walk(UPLOAD_DIR):
             for file in files:
                 file_path = os.path.join(root, file)
-                # Add file to zip, preserving directory structure relative to UPLOAD_DIR
-                zipf.write(file_path, os.path.relpath(file_path, UPLOAD_DIR))
+                if file != exclude_file:
+                    # Add file to zip, preserving directory structure relative to UPLOAD_DIR
+                    zipf.write(file_path, os.path.relpath(file_path, UPLOAD_DIR))
     
-    print(f"Successfully created zip file at: {zip_path}")
+    print(f"Successfully created versioned zip file at: {zip_path}")
     return zip_path
 
-def upload_to_huggingface(zip_path: str):
-    """Uploads the zip file to the specified Hugging Face dataset."""
-    if not HF_TOKEN:
-        print("HF_TOKEN not found in environment variables. Skipping upload.")
-        return
-
+def upload_zip_to_hf(zip_path: str):
+    """Uploads the versioned zip file to the Hugging Face dataset."""
     if not zip_path or not os.path.exists(zip_path):
         print("Zip file not found. Skipping upload.")
         return
 
-    try:
-        api = HfApi()
-        
-        # Upload the zip file to the root of the dataset repository
-        api.upload_file(
-            path_or_fileobj=zip_path,
-            path_in_repo=os.path.basename(zip_path),
-            repo_id=HF_DATASET_REPO,
-            repo_type="dataset",
-            token=HF_TOKEN
-        )
-        print(f"Successfully uploaded {os.path.basename(zip_path)} to {HF_DATASET_REPO}")
-    except Exception as e:
-        print(f"Hugging Face upload failed: {e}")
+    upload_file_to_huggingface(zip_path, os.path.basename(zip_path))
+    
+    # Reset counter after successful upload
+    global app_state
+    app_state["files_since_last_zip"] = 0
 
 def cleanup_upload_dir():
-    """Removes the upload directory and its contents."""
+    """Removes the upload directory and its contents, but keeps the state file."""
     if os.path.exists(UPLOAD_DIR):
-        shutil.rmtree(UPLOAD_DIR)
-        print(f"Cleaned up {UPLOAD_DIR} directory.")
+        for item in os.listdir(UPLOAD_DIR):
+            item_path = os.path.join(UPLOAD_DIR, item)
+            if os.path.isdir(item_path):
+                shutil.rmtree(item_path)
+            elif item != os.path.basename(STATE_FILE_PATH):
+                os.remove(item_path)
+        print(f"Cleaned up files in {UPLOAD_DIR} directory (state file preserved).")
 
 # --- Application Lifespan ---
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup: Ensure upload directory exists
+    # Startup: 1. Ensure upload directory exists
     os.makedirs(UPLOAD_DIR, exist_ok=True)
     print(f"Application starting. Upload directory: {UPLOAD_DIR}")
+    
+    # Startup: 2. Download and load state
+    download_state_file()
+    load_state()
+    print(f"Initial state: {app_state}")
+    
     yield
-    # Shutdown: Zip and upload files
-    print("Application shutting down. Initiating final upload...")
-    zip_path = zip_uploaded_files()
+    
+    # Shutdown: 1. Upload state file
+    print("Application shutting down. Uploading final state...")
+    upload_state_file_to_hf()
+    
+    # Shutdown: 2. Perform final zip and upload of any remaining files
+    print("Performing final zip and upload...")
+    zip_path = zip_uploaded_files_versioned()
     if zip_path:
-        upload_to_huggingface(zip_path)
+        upload_zip_to_hf(zip_path)
         # Clean up the created zip file after upload
         os.remove(zip_path)
+        
+    # Shutdown: 3. Cleanup
     cleanup_upload_dir()
     print("Shutdown complete.")
 
 # --- FastAPI App Initialization ---
 
 app = FastAPI(
-    title="File Uploader and Downloader Service",
-    description="A simple service for file management and Hugging Face dataset synchronization.",
-    version="1.0.0",
+    title="Versioned File Uploader and Dataset Sync Service",
+    description="A service for file management with versioned zipping and state persistence to Hugging Face.",
+    version="2.0.0",
     lifespan=lifespan
 )
 
@@ -101,18 +206,41 @@ app = FastAPI(
 
 @app.post("/upload/")
 async def upload_file(file: UploadFile = File(...)):
-    """Upload a file to the server."""
+    """Upload a file to the server and trigger periodic tasks."""
+    global app_state
+    
+    # 1. Save the file
     try:
         file_path = os.path.join(UPLOAD_DIR, file.filename)
         # Check if file already exists to prevent overwriting without warning
         if os.path.exists(file_path):
             raise HTTPException(status_code=409, detail=f"File '{file.filename}' already exists.")
-            
-        # Write the file content to disk
-        with open(file_path, "wb") as buffer:
-            shutil.copyfileobj(file.file, buffer)
+
+        try:
+            with open(file_path, "wb") as buffer:
+                shutil.copyfileobj(file.file, buffer)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Could not write file to disk: {e}")
         
-        return {"filename": file.filename, "message": "File successfully uploaded"}
+        # 2. Update counters only after successful write
+        app_state["total_files_uploaded"] += 1
+        app_state["files_since_last_zip"] += 1
+        app_state["files_since_last_state_upload"] += 1
+        save_state() # Save state after every upload  
+        # 3. Check for state file upload threshold (100 uploads)
+        if app_state["files_since_last_state_upload"] >= STATE_UPLOAD_THRESHOLD:
+            print(f"State upload threshold ({STATE_UPLOAD_THRESHOLD}) reached. Uploading state file...")
+            upload_state_file_to_hf()
+            
+        # 4. Check for zip upload threshold (200 uploads)
+        if app_state["files_since_last_zip"] >= ZIP_UPLOAD_THRESHOLD:
+            print(f"Zip upload threshold ({ZIP_UPLOAD_THRESHOLD}) reached. Zipping and uploading files...")
+            zip_path = zip_uploaded_files_versioned()
+            if zip_path:
+                upload_zip_to_hf(zip_path)
+                os.remove(zip_path)
+                
+        return {"filename": file.filename, "message": "File successfully uploaded", "state": app_state}
     except HTTPException:
         raise
     except Exception as e:
@@ -136,23 +264,36 @@ async def download_file(filename: str):
 async def sync_dataset():
     """Manually trigger zipping of all uploaded files and uploading to the Hugging Face dataset."""
     print("Manual dataset sync triggered.")
-    zip_path = zip_uploaded_files()
+    
+    # 1. Upload state file
+    upload_state_file_to_hf()
+    
+    # 2. Zip and upload files
+    zip_path = zip_uploaded_files_versioned()
     if not zip_path:
-        return {"message": "No files to sync. Upload directory is empty."}
+        return {"message": "No files to sync. Upload directory is empty.", "state": app_state}
     
-    upload_to_huggingface(zip_path)
+    upload_zip_to_hf(zip_path)
     
     # Clean up the created zip file after upload
     os.remove(zip_path)
     
-    return {"message": "Files zipped and upload to Hugging Face dataset initiated."}
+    return {"message": "Files zipped and upload to Hugging Face dataset initiated.", "state": app_state}
 
 @app.get("/files/")
 async def list_files():
     """List all files currently available for download."""
     if not os.path.exists(UPLOAD_DIR):
         return {"files": []}
-    return {"files": os.listdir(UPLOAD_DIR)}
+    
+    # Exclude the state file from the list of downloadable files
+    files = [f for f in os.listdir(UPLOAD_DIR) if f != os.path.basename(STATE_FILE_PATH)]
+    return {"files": files, "state": app_state}
+
+@app.get("/state/")
+async def get_state():
+    """Get the current application state."""
+    return {"state": app_state}
 
 # --- Main execution block for testing/running ---
 if __name__ == "__main__":
@@ -178,4 +319,4 @@ if __name__ == "__main__":
     finally:
         # Simulate cleanup that happens in the lifespan context manager
         # when running with uvicorn in a real environment.
-        pass
+        pass