Upload 2 files

api_server.py

@@ -1,193 +1,399 @@
-from fastapi import FastAPI, HTTPException, BackgroundTasks
-from fastapi.responses import JSONResponse
-import asyncio
-import os
-from typing import Optional, Dict, Any
-from pydantic import BaseModel
-import download_channel  # Import the existing downloader
-
-app = FastAPI(title="Telegram Channel Downloader API")
-
-# Track active downloads and their status
-active_downloads: Dict[str, Dict[str, Any]] = {}
-
-@app.on_event("startup")
-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
-    ))
-    print(f"Started initial download task with ID: {task_id}")
-
-class DownloadRequest(BaseModel):
-    channel: Optional[str] = None  # If None, uses default from download_channel.py
-    message_limit: Optional[int] = None
-
-class DownloadStatus(BaseModel):
-    channel: str
-    status: str  # "running", "completed", "failed"
-    message_count: int = 0
-    downloaded: int = 0
-    skipped: int = 0
-    not_rar: int = 0
-    error: Optional[str] = None
-
-async def run_download(channel: Optional[str], message_limit: Optional[int], task_id: str):
-    """Background task to run the download"""
-    try:
-        # Override channel and message limit if provided
-        if channel:
-            download_channel.CHANNEL = channel
-        if message_limit is not None:
-            download_channel.MESSAGE_LIMIT = message_limit
-
-        # Create a status tracker
-        status = {
-            "channel": download_channel.CHANNEL,
-            "status": "running",
-            "message_count": 0,
-            "downloaded": 0,
-            "skipped": 0,
-            "not_rar": 0,
-            "error": None
-        }
-        active_downloads[task_id] = status
-
-        # Patch the download function to update our status
-        original_download = download_channel.download_channel
-        
-        async def wrapped_download():
-            nonlocal status
-            try:
-                # Initialize client and get entity
-                client = download_channel.TelegramClient(
-                    download_channel.SESSION_FILE,
-                    download_channel.API_ID,
-                    download_channel.API_HASH
-                )
-
-                async with client:
-                    try:
-                        entity = await client.get_entity(download_channel.CHANNEL)
-                    except Exception as e:
-                        status["error"] = f"Failed to resolve channel: {str(e)}"
-                        status["status"] = "failed"
-                        return 1
-
-                    try:
-                        async for message in client.iter_messages(entity, limit=download_channel.MESSAGE_LIMIT or None):
-                            status["message_count"] += 1
-                            
-                            if not message.media:
-                                continue
-
-                            # Check if it's a RAR file
-                            is_rar = False
-                            if message.file:
-                                filename = getattr(message.file, 'name', '') or ''
-                                if filename:
-                                    is_rar = filename.lower().endswith('.rar')
-                                else:
-                                    mime_type = getattr(message.file, 'mime_type', '') or ''
-                                    is_rar = 'rar' in mime_type.lower() if mime_type else False
-
-                            if not is_rar:
-                                status["not_rar"] += 1
-                                continue
-
-                            # Use the same filename logic
-                            if filename:
-                                suggested = f"{message.id}_{filename}"
-                            else:
-                                suggested = f"{message.id}.rar"
-
-                            out_path = os.path.join(download_channel.OUTPUT_DIR, suggested)
-
-                            if os.path.exists(out_path):
-                                status["skipped"] += 1
-                                continue
-
-                            try:
-                                await client.download_media(message, file=out_path)
-                                status["downloaded"] += 1
-
-                                # Upload to HF if token available
-                                if download_channel.HF_TOKEN:
-                                    path_in_repo = f"files/{os.path.basename(out_path)}"
-                                    ok = download_channel.upload_file_to_hf(
-                                        out_path, path_in_repo, download_channel.HF_TOKEN
-                                    )
-                                    if not ok:
-                                        print(f"Warning: failed to upload {out_path}")
-
-                                await asyncio.sleep(0.2)  # Be polite
-
-                            except download_channel.errors.FloodWaitError as fw:
-                                wait = int(fw.seconds) if fw.seconds else 60
-                                print(f"Hit FloodWait: sleeping {wait}s")
-                                await asyncio.sleep(wait + 1)
-                            except Exception as e:
-                                print(f"Error downloading {message.id}: {e}")
-
-                    except Exception as e:
-                        status["error"] = str(e)
-                        status["status"] = "failed"
-                        return 1
-
-                    status["status"] = "completed"
-                    return 0
-
-            except Exception as e:
-                status["error"] = str(e)
-                status["status"] = "failed"
-                return 1
-
-        await wrapped_download()
-
-    except Exception as e:
-        active_downloads[task_id] = {
-            "channel": download_channel.CHANNEL,
-            "status": "failed",
-            "message_count": 0,
-            "downloaded": 0,
-            "skipped": 0,
-            "not_rar": 0,
-            "error": str(e)
-        }
-
-@app.post("/download", response_model=Dict[str, str])
-async def start_download(request: DownloadRequest, background_tasks: BackgroundTasks):
-    """Start a new download task"""
-    task_id = f"download_{len(active_downloads) + 1}"
-    
-    # Schedule the download
-    background_tasks.add_task(
-        run_download,
-        channel=request.channel,
-        message_limit=request.message_limit,
-        task_id=task_id
-    )
-    
-    return {"task_id": task_id}
-
-@app.get("/status/{task_id}", response_model=DownloadStatus)
-async def get_status(task_id: str):
-    """Get the status of a download task"""
-    if task_id not in active_downloads:
-        raise HTTPException(status_code=404, detail="Task not found")
-    return active_downloads[task_id]
-
-@app.get("/active", response_model=Dict[str, DownloadStatus])
-async def list_active():
-    """List all active or completed downloads"""
-    return active_downloads
-
-if __name__ == "__main__":
-    import uvicorn
-    # Note: When running directly, this runs on 8000
-    # For production, use: uvicorn api_server:app --host 0.0.0.0 --port 8000
+from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi.responses import JSONResponse
+import asyncio
+import os
+import time
+from typing import Optional, Dict, Any, List
+from enum import Enum
+from pydantic import BaseModel
+from rich.progress import (
+    Progress,
+    SpinnerColumn,
+    TimeElapsedColumn,
+    DownloadColumn,
+    TransferSpeedColumn,
+    BarColumn,
+    TextColumn,
+)
+from rich.console import Console
+from rich.live import Live
+from rich.table import Table
+import download_channel
+
+# Initialize rich console for pretty logging
+console = Console()
+
+app = FastAPI(title="Telegram Channel Downloader API")
+
+# Track active downloads and their status
+active_downloads: Dict[str, Dict[str, Any]] = {}
+
+class FileStatus(str, Enum):
+    PENDING = "pending"
+    DOWNLOADING = "downloading"
+    DOWNLOADED = "downloaded"
+    FAILED = "failed"
+
+class ChannelFile(BaseModel):
+    message_id: int
+    filename: str
+    status: FileStatus
+    size: Optional[int] = None
+    download_time: Optional[float] = None
+    error: Optional[str] = None
+    upload_path: Optional[str] = None
+
+class DownloadState(BaseModel):
+    channel: str
+    last_scanned_id: Optional[int] = None
+    files: List[ChannelFile] = []
+    current_download: Optional[int] = None  # message_id of current download
+    last_updated: float = time.time()
+
+class DownloadRequest(BaseModel):
+    channel: Optional[str] = None
+    message_limit: Optional[int] = None
+
+class DownloadStatus(BaseModel):
+    channel: str
+    status: str
+    message_count: int = 0
+    downloaded: int = 0
+    downloading: Optional[str] = None
+    error: Optional[str] = None
+
+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:
+        return DownloadState(channel=download_channel.CHANNEL)
+    try:
+        # Try to download existing state
+        local_path = download_channel.hf_hub_download(
+            repo_id=download_channel.HF_REPO_ID,
+            filename=download_channel.STATE_FILE,
+            repo_type="dataset",
+            token=token
+        )
+        with open(local_path, "r", encoding="utf-8") as f:
+            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)
+
+async def clean_downloaded_file(file_path: str):
+    """Remove local file after successful upload"""
+    try:
+        os.remove(file_path)
+        console.print(f"[blue]Cleaned up:[/blue] {os.path.basename(file_path)}")
+    except Exception as e:
+        console.print(f"[yellow]Warning:[/yellow] Could not clean up {file_path}: {e}")
+
+async def update_and_upload_state(state: DownloadState, token: str) -> bool:
+    """Update state timestamp and upload to dataset"""
+    state.last_updated = time.time()
+    try:
+        # Save state locally first
+        with open(download_channel.STATE_FILE, "w", encoding="utf-8") as f:
+            json.dump(state.dict(), f, indent=2, ensure_ascii=False)
+        # Upload to dataset
+        return download_channel.upload_file_to_hf(
+            download_channel.STATE_FILE,
+            download_channel.STATE_FILE,
+            token
+        )
+    except Exception as e:
+        console.print(f"[red]Failed to update state:[/red] {e}")
+        return False
+
+async def process_message(message, state: DownloadState, client) -> Optional[str]:
+    """Process a single message, return output path if file downloaded or None"""
+    if not message.media:
+        return None
+
+    # Check if it's a RAR file
+    is_rar = False
+    filename = ""
+    if message.file:
+        filename = getattr(message.file, 'name', '') or ''
+        if filename:
+            is_rar = filename.lower().endswith('.rar')
+        else:
+            mime_type = getattr(message.file, 'mime_type', '') or ''
+            is_rar = 'rar' in mime_type.lower() if mime_type else False
+
+    if not is_rar:
+        return None
+
+    # Use message ID and original filename for saved file
+    if filename:
+        suggested = f"{message.id}_{filename}"
+    else:
+        suggested = f"{message.id}.rar"
+
+    return os.path.join(download_channel.OUTPUT_DIR, suggested)
+
+async def run_download(channel: Optional[str], message_limit: Optional[int], task_id: str):
+    """Background task to run the download with state management"""
+    try:
+        # Override channel if provided
+        if channel:
+            download_channel.CHANNEL = channel
+        if message_limit is not None:
+            download_channel.MESSAGE_LIMIT = message_limit
+
+        # Get or create download state
+        state = download_state_from_hf(download_channel.HF_TOKEN)
+        
+        # Initialize status for API
+        status = {
+            "channel": state.channel,
+            "status": "running",
+            "message_count": len(state.files),
+            "downloaded": len([f for f in state.files if f.status == FileStatus.DOWNLOADED]),
+            "downloading": None,
+            "error": None
+        }
+        active_downloads[task_id] = status
+
+        # Create progress displays
+        progress = Progress(
+            SpinnerColumn(),
+            TextColumn("[bold blue]{task.fields[filename]}", justify="right"),
+            BarColumn(bar_width=40),
+            "[progress.percentage]{task.percentage:>3.1f}%",
+            "•",
+            DownloadColumn(),
+            "•",
+            TransferSpeedColumn(),
+            "•",
+            TimeElapsedColumn(),
+        )
+        
+        overall_progress = Progress(
+            TextColumn("[bold yellow]{task.description}", justify="right"),
+            BarColumn(bar_width=40),
+            "[progress.percentage]{task.percentage:>3.1f}%",
+            "•",
+            TextColumn("[bold green]{task.fields[stats]}")
+        )
+
+        # Initialize client
+        client = download_channel.TelegramClient(
+            download_channel.SESSION_FILE,
+            download_channel.API_ID,
+            download_channel.API_HASH
+        )
+
+        async with client:
+            try:
+                entity = await client.get_entity(download_channel.CHANNEL)
+            except Exception as e:
+                console.print(f"[red]Failed to resolve channel:[/red] {e}")
+                return 1
+
+            console.print(f"[green]Starting download from:[/green] {entity.title if hasattr(entity, 'title') else download_channel.CHANNEL}")
+
+            # First, scan for new messages and update state
+            scan_count = 0
+            last_message_id = state.last_scanned_id
+            
+            try:
+                async for message in client.iter_messages(entity, limit=download_channel.MESSAGE_LIMIT or None):
+                    scan_count += 1
+                    
+                    # Update last scanned ID
+                    if last_message_id is None or message.id > last_message_id:
+                        last_message_id = message.id
+                    
+                    # Skip if we already know about this message
+                    if any(f.message_id == message.id for f in state.files):
+                        continue
+                        
+                    # Check if it's a downloadable file
+                    out_path = await process_message(message, state, client)
+                    if out_path:
+                        # Add to state as pending
+                        file_info = ChannelFile(
+                            message_id=message.id,
+                            filename=os.path.basename(out_path),
+                            status=FileStatus.PENDING,
+                            size=getattr(message.media, 'size', 0) or 0
+                        )
+                        state.files.append(file_info)
+                
+                # Update state with scan results
+                state.last_scanned_id = last_message_id
+                if download_channel.HF_TOKEN:
+                    await update_and_upload_state(state, download_channel.HF_TOKEN)
+                    
+                console.print(f"[green]Channel scan complete:[/green] Found {scan_count} messages")
+                    
+            except Exception as e:
+                console.print(f"[red]Error during channel scan:[/red] {e}")
+                
+            # Now process pending downloads
+            pending_files = [f for f in state.files if f.status == FileStatus.PENDING]
+            total_pending = len(pending_files)
+
+            if total_pending == 0:
+                console.print("[green]No new files to download![/green]")
+                return 0
+
+            console.print(f"[green]Starting downloads:[/green] {total_pending} files pending")
+
+            # Process pending files
+            with Live(progress) as live_progress, Live(overall_progress) as live_overall:
+                overall_task = overall_progress.add_task(
+                    f"Channel: {download_channel.CHANNEL}",
+                    total=total_pending,
+                    stats=f"Pending: {total_pending}"
+                )
+
+                for file_info in pending_files:
+                    try:
+                        # Mark as downloading in state
+                        file_info.status = FileStatus.DOWNLOADING
+                        state.current_download = file_info.message_id
+                        if download_channel.HF_TOKEN:
+                            await update_and_upload_state(state, download_channel.HF_TOKEN)
+
+                        # Update status
+                        status["downloading"] = file_info.filename
+
+                        # Get message and prepare download
+                        message = await client.get_messages(entity, ids=file_info.message_id)
+                        if not message or not message.media:
+                            file_info.status = FileStatus.FAILED
+                            file_info.error = "Message not found or no media"
+                            continue
+
+                        out_path = os.path.join(download_channel.OUTPUT_DIR, file_info.filename)
+                        file_task = progress.add_task(
+                            "download",
+                            total=file_info.size or 100,
+                            filename=file_info.filename
+                        )
+
+                        # Download with progress
+                        start_time = time.time()
+                        try:
+                            async def progress_callback(current, total):
+                                progress.update(file_task, completed=current)
+                                overall_stats = f"Downloaded: {len([f for f in state.files if f.status == FileStatus.DOWNLOADED])}"
+                                overall_progress.update(overall_task, completed=current/total*100, stats=overall_stats)
+
+                            await client.download_media(
+                                message,
+                                file=out_path,
+                                progress_callback=progress_callback
+                            )
+
+                            # Upload to HF
+                            if download_channel.HF_TOKEN:
+                                console.print(f"[yellow]Uploading to HF:[/yellow] {file_info.filename}")
+                                path_in_repo = f"files/{file_info.filename}"
+                                ok = download_channel.upload_file_to_hf(
+                                    out_path,
+                                    path_in_repo,
+                                    download_channel.HF_TOKEN
+                                )
+                                if ok:
+                                    console.print(f"[green]Uploaded:[/green] {file_info.filename}")
+                                    # Clean up local file
+                                    await clean_downloaded_file(out_path)
+                                    file_info.upload_path = path_in_repo
+                                else:
+                                    console.print(f"[red]Upload failed:[/red] {file_info.filename}")
+                                    file_info.error = "Upload to dataset failed"
+                                    file_info.status = FileStatus.FAILED
+                                    continue
+
+                            # Mark as completed in state
+                            file_info.status = FileStatus.DOWNLOADED
+                            file_info.download_time = time.time() - start_time
+                            
+                            # Update state
+                            if download_channel.HF_TOKEN:
+                                await update_and_upload_state(state, download_channel.HF_TOKEN)
+
+                            # Update status
+                            status["downloaded"] += 1
+                            await asyncio.sleep(0.2)  # Be polite
+
+                        except download_channel.errors.FloodWaitError as fw:
+                            wait = int(fw.seconds) if fw.seconds else 60
+                            console.print(f"[yellow]FloodWait:[/yellow] Sleeping {wait}s")
+                            await asyncio.sleep(wait + 1)
+                            # Retry this file
+                            continue
+
+                        except Exception as e:
+                            console.print(f"[red]Error:[/red] {str(e)}")
+                            file_info.status = FileStatus.FAILED
+                            file_info.error = str(e)
+                            if download_channel.HF_TOKEN:
+                                await update_and_upload_state(state, download_channel.HF_TOKEN)
+
+                    except Exception as e:
+                        console.print(f"[red]Fatal error processing {file_info.filename}:[/red] {str(e)}")
+                        continue
+
+            # Clear current download
+            state.current_download = None
+            if download_channel.HF_TOKEN:
+                await update_and_upload_state(state, download_channel.HF_TOKEN)
+
+            console.print("[green]Download session completed![/green]")
+            status["status"] = "completed"
+            status["downloading"] = None
+
+    except Exception as e:
+        console.print(f"[red]Fatal error:[/red] {str(e)}")
+        if "status" in locals():
+            status["status"] = "failed"
+            status["error"] = str(e)
+    
+    return 0
+
+@app.on_event("startup")
+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}")
+
+@app.post("/download", response_model=Dict[str, str])
+async def start_download(request: DownloadRequest, background_tasks: BackgroundTasks):
+    """Start a new download task"""
+    task_id = f"download_{len(active_downloads) + 1}"
+    
+    background_tasks.add_task(
+        run_download,
+        channel=request.channel,
+        message_limit=request.message_limit,
+        task_id=task_id
+    )
+    
+    return {"task_id": task_id}
+
+@app.get("/status/{task_id}", response_model=DownloadStatus)
+async def get_status(task_id: str):
+    """Get the status of a download task"""
+    if task_id not in active_downloads:
+        raise HTTPException(status_code=404, detail="Task not found")
+    return active_downloads[task_id]
+
+@app.get("/active", response_model=Dict[str, DownloadStatus])
+async def list_active():
+    """List all active or completed downloads"""
+    return active_downloads
+
+if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="127.0.0.1", port=8000)

requirements.txt

@@ -3,3 +3,4 @@ huggingface_hub>=0.17.0
 fastapi>=0.104.0
 uvicorn>=0.24.0
 pydantic>=2.4.2
+rich>=13.6.0