Update api_server.py

api_server.py

@@ -1,181 +1,193 @@
-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]] = {}
-
-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
+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
     uvicorn.run(app, host="127.0.0.1", port=8000)