Update api_server.py

api_server.py

@@ -3,7 +3,6 @@ from fastapi.responses import JSONResponse
 import asyncio
 import os
 import time
-import json
 from typing import Optional, Dict, Any, List
 from enum import Enum
 from pydantic import BaseModel
@@ -63,6 +62,41 @@ class DownloadStatus(BaseModel):
     downloading: Optional[str] = None
     error: Optional[str] = None
 
+def create_hf_dataset(token: str) -> bool:
+    """Create the Hugging Face dataset if it doesn't exist."""
+    try:
+        from huggingface_hub import create_repo, RepoNotFoundError
+        try:
+            # Try to create the dataset repository
+            create_repo(
+                repo_id=download_channel.HF_REPO_ID,
+                token=token,
+                repo_type="dataset",
+                exist_ok=True
+            )
+            console.print(f"[green]Created or verified dataset:[/green] {download_channel.HF_REPO_ID}")
+            
+            # Create initial state file
+            initial_state = DownloadState(channel=download_channel.CHANNEL)
+            with open(download_channel.STATE_FILE, "w", encoding="utf-8") as f:
+                json.dump(initial_state.dict(), f, indent=2, ensure_ascii=False)
+            
+            # Upload initial state
+            if download_channel.upload_file_to_hf(
+                download_channel.STATE_FILE,
+                download_channel.STATE_FILE,
+                token
+            ):
+                console.print("[green]Initialized dataset with empty state file[/green]")
+                return True
+        except Exception as e:
+            console.print(f"[red]Failed to create dataset:[/red] {str(e)}")
+            return False
+    except ImportError:
+        console.print("[red]huggingface_hub not properly installed[/red]")
+        return False
+    return True
+
 def download_state_from_hf(token: str) -> DownloadState:
     """Try to download the state file from the HF dataset. Returns state dict or creates new."""
     if not token:
@@ -79,8 +113,13 @@ def download_state_from_hf(token: str) -> DownloadState:
             data = json.load(f)
             return DownloadState(**data)
     except Exception as e:
-        console.print(f"[yellow]No existing state found, creating new:[/yellow] {str(e)}")
-        return DownloadState(channel=download_channel.CHANNEL)
+        console.print(f"[yellow]No existing state found, creating new dataset:[/yellow] {str(e)}")
+        if create_hf_dataset(token):
+            console.print("[green]Dataset created successfully![/green]")
+            return DownloadState(channel=download_channel.CHANNEL)
+        else:
+            console.print("[red]Failed to create dataset, using local state only[/red]")
+            return DownloadState(channel=download_channel.CHANNEL)
 
 async def clean_downloaded_file(file_path: str):
     """Remove local file after successful upload"""
@@ -361,13 +400,31 @@ async def run_download(channel: Optional[str], message_limit: Optional[int], tas
 async def start_initial_download():
     """Start the download process automatically when the server starts"""
     task_id = "initial_download"
-    # Start the download process with default settings
-    asyncio.create_task(run_download(
-        channel=None,  # Use default from download_channel.py
-        message_limit=None,  # Use default
-        task_id=task_id
-    ))
-    console.print(f"[green]Started initial download task:[/green] {task_id}")
+    
+    # Verify HF token is set
+    if not download_channel.HF_TOKEN:
+        console.print("[red]ERROR: HF_TOKEN not set. Please set your Hugging Face token.[/red]")
+        return
+        
+    # Create dataset structure if needed
+    console.print("[yellow]Checking Hugging Face dataset...[/yellow]")
+    try:
+        state = download_state_from_hf(download_channel.HF_TOKEN)
+        console.print(f"[green]Using channel:[/green] {state.channel}")
+        
+        # Create files directory in dataset if it doesn't exist
+        os.makedirs(download_channel.OUTPUT_DIR, exist_ok=True)
+        
+        # Start the download process with default settings
+        asyncio.create_task(run_download(
+            channel=None,  # Use default from download_channel.py
+            message_limit=None,  # Use default
+            task_id=task_id
+        ))
+        console.print(f"[green]Started initial download task:[/green] {task_id}")
+        
+    except Exception as e:
+        console.print(f"[red]Failed to initialize:[/red] {str(e)}")
 
 @app.post("/download", response_model=Dict[str, str])
 async def start_download(request: DownloadRequest, background_tasks: BackgroundTasks):