Update app.py

app.py

@@ -12,12 +12,14 @@ 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")
+STATE_FILE_NAME = "processing_audio_state.json"
+STATE_FILE_PATH = os.path.join(os.getcwd(), STATE_FILE_NAME)
 HF_DATASET_REPO = "samfred2/A_Text"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Hardcoded variables for automation
-ZIP_UPLOAD_THRESHOLD = 200
+# Both thresholds are now 100 as requested
+ZIP_UPLOAD_THRESHOLD = 100
 STATE_UPLOAD_THRESHOLD = 100
 
 # Global state object (will be loaded from/saved to STATE_FILE_PATH)
@@ -36,14 +38,14 @@ def get_hf_api() -> HfApi:
         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."""
+def download_state_file_from_hf():
+    """Downloads the state file from the Hugging Face dataset to the local path."""
     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),
+            filename=STATE_FILE_NAME,
             repo_type="dataset",
             local_dir=os.path.dirname(STATE_FILE_PATH),
             local_dir_use_symlinks=False
@@ -70,7 +72,8 @@ def upload_file_to_huggingface(local_path: str, path_in_repo: str):
             path_in_repo=path_in_repo,
             repo_id=HF_DATASET_REPO,
             repo_type="dataset",
-            token=HF_TOKEN
+            token=HF_TOKEN,
+            commit_message=f"Sync: {path_in_repo}"
         )
         print(f"Successfully uploaded {os.path.basename(local_path)} to {HF_DATASET_REPO} as {path_in_repo}")
         return True
@@ -89,26 +92,41 @@ def load_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}.")
+                # Load the state file content, which contains the file_states and next_download_index
+                state_content = json.load(f)
+                
+                # The app_state itself only tracks counters, not the file_states
+                # We only update the counters here if they exist in the state file
+                if "total_files_uploaded" in state_content:
+                    app_state["total_files_uploaded"] = state_content["total_files_uploaded"]
+                if "current_zip_version" in state_content:
+                    app_state["current_zip_version"] = state_content["current_zip_version"]
+                if "files_since_last_zip" in state_content:
+                    app_state["files_since_last_zip"] = state_content["files_since_last_zip"]
+                if "files_since_last_state_upload" in state_content:
+                    app_state["files_since_last_state_upload"] = state_content["files_since_last_state_upload"]
+                    
+            print(f"State counters 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."""
+    """Saves application state (counters) to the local state file."""
+    # This function is now only for saving the internal app_state counters, 
+    # not the file_states, which are managed by the client.
     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}.")
+        print(f"State counters 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."""
+    """Saves the current app_state counters and uploads the state file to Hugging Face."""
     save_state()
-    upload_file_to_huggingface(STATE_FILE_PATH, os.path.basename(STATE_FILE_PATH))
+    upload_file_to_huggingface(STATE_FILE_PATH, STATE_FILE_NAME)
     global app_state
     app_state["files_since_last_state_upload"] = 0
 
@@ -126,12 +144,13 @@ def zip_uploaded_files_versioned() -> str:
     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)
+    exclude_file = STATE_FILE_NAME
 
     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)
+                # Only zip files that are NOT the state file
                 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))
@@ -158,7 +177,8 @@ def cleanup_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):
+            # Only remove files that are NOT the state file
+            elif item != STATE_FILE_NAME:
                 os.remove(item_path)
         print(f"Cleaned up files in {UPLOAD_DIR} directory (state file preserved).")
 
@@ -170,15 +190,16 @@ async def lifespan(app: FastAPI):
     os.makedirs(UPLOAD_DIR, exist_ok=True)
     print(f"Application starting. Upload directory: {UPLOAD_DIR}")
     
-    # Startup: 2. Download and load state
-    download_state_file()
+    # Startup: 2. Download and load state (only internal counters are loaded here)
+    # We rely on the /state/ endpoint to download the file_states from HF when needed.
+    download_state_file_from_hf()
     load_state()
-    print(f"Initial state: {app_state}")
+    print(f"Initial state counters: {app_state}")
     
     yield
     
-    # Shutdown: 1. Upload state file
-    print("Application shutting down. Uploading final state...")
+    # Shutdown: 1. Upload state file (only internal counters are uploaded here)
+    print("Application shutting down. Uploading final state counters...")
     upload_state_file_to_hf()
     
     # Shutdown: 2. Perform final zip and upload of any remaining files
@@ -212,27 +233,42 @@ async def upload_file(file: UploadFile = File(...)):
     # 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.")
+            # Special handling for the state file: allow overwrite as it's a sync mechanism
+            if file.filename != STATE_FILE_NAME:
+                raise HTTPException(status_code=409, detail=f"File '{file.filename}' already exists.")
 
         try:
+            # Write the uploaded file to the local UPLOAD_DIR
             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}")
         
-        # 2. Update counters only after successful write
+        # --- LOGIC FOR STATE FILE HANDLING ---
+        if file.filename == STATE_FILE_NAME:
+            print(f"Intercepted state file upload: {STATE_FILE_NAME}. Saving locally only.")
+            # The client's lock/unlock is now persisted locally on the server.
+            # It will only be uploaded to HF when the STATE_UPLOAD_THRESHOLD is met.
+            # We return success immediately.
+            return {"filename": file.filename, "message": "State file successfully saved locally", "state": app_state}
+        # --- END LOGIC FOR STATE FILE HANDLING ---
+        
+        # 2. Update counters only after successful write of a non-state file (transcription JSON)
         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)
+        save_state() # Save internal counters locally
+        
+        # 3. Check for state file upload threshold (100 non-state file uploads)
         if app_state["files_since_last_state_upload"] >= STATE_UPLOAD_THRESHOLD:
             print(f"State upload threshold ({STATE_UPLOAD_THRESHOLD}) reached. Uploading state file...")
+            # This uploads the local state file (which includes the latest file_states) to HF
             upload_state_file_to_hf()
             
-        # 4. Check for zip upload threshold (200 uploads)
+        # 4. Check for zip upload threshold (100 non-state file 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()
@@ -265,7 +301,7 @@ async def sync_dataset():
     """Manually trigger zipping of all uploaded files and uploading to the Hugging Face dataset."""
     print("Manual dataset sync triggered.")
     
-    # 1. Upload state file
+    # 1. Upload state file (internal counters)
     upload_state_file_to_hf()
     
     # 2. Zip and upload files
@@ -287,13 +323,44 @@ async def list_files():
         return {"files": []}
     
     # 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)]
+    files = [f for f in os.listdir(UPLOAD_DIR) if f != STATE_FILE_NAME]
     return {"files": files, "state": app_state}
 
 @app.get("/state/")
 async def get_state():
-    """Get the current application state."""
-    return {"state": app_state}
+    """
+    Get the current application state (including file processing status).
+    This endpoint now ensures it fetches the latest state from Hugging Face.
+    """
+    # 1. Load the file_states from the local file, which is updated by client uploads.
+    # We assume the local file exists because it's either downloaded on startup
+    # or created by a client upload.
+    if not os.path.exists(STATE_FILE_PATH):
+        # If the file is missing, attempt a download from HF as a fallback
+        download_state_file_from_hf()
+        if not os.path.exists(STATE_FILE_PATH):
+        # If still not found (e.g., first run and 404 on HF), return default
+        return {"state": {"next_download_index": 0, "file_states": {}}}
+        
+    try:
+        with open(STATE_FILE_PATH, "r") as f:
+            state_content = json.load(f)
+            
+            # Ensure the required keys for the client are present
+            if "file_states" not in state_content:
+                state_content["file_states"] = {}
+            if "next_download_index" not in state_content:
+                state_content["next_download_index"] = 0
+                
+            # Merge the internal app_state counters (which are not critical for locking)
+            # into the returned state for completeness, but the core file_states come from the file.
+            state_content.update(app_state)
+            
+            return {"state": state_content}
+            
+    except Exception as e:
+        print(f"Error reading local state file for /state/ endpoint: {e}")
+        raise HTTPException(status_code=500, detail="Error processing state file.")
 
 # --- Main execution block for testing/running ---
 if __name__ == "__main__":
@@ -319,4 +386,4 @@ if __name__ == "__main__":
     finally:
         # Simulate cleanup that happens in the lifespan context manager
         # when running with uvicorn in a real environment.
-        pass
+        pass