Spaces:
Runtime error
Runtime error
| """ | |
| Progress Tracking System | |
| Manages progress updates for long-running enhancement tasks | |
| Supports Server-Sent Events (SSE) for real-time frontend updates | |
| """ | |
| import asyncio | |
| import logging | |
| from typing import Dict, Optional | |
| from datetime import datetime | |
| logger = logging.getLogger(__name__) | |
| class ProgressTracker: | |
| """ | |
| Tracks progress for enhancement operations | |
| Allows multiple clients to subscribe to progress updates | |
| """ | |
| def __init__(self): | |
| self.tasks: Dict[str, Dict] = {} | |
| self.subscribers: Dict[str, list] = {} | |
| def create_task(self, task_id: str, total_steps: int = 100): | |
| """Initialize a new task for progress tracking""" | |
| self.tasks[task_id] = { | |
| "id": task_id, | |
| "status": "initializing", | |
| "progress": 0, | |
| "total_steps": total_steps, | |
| "current_step": 0, | |
| "message": "Starting enhancement...", | |
| "started_at": datetime.now().isoformat(), | |
| "error": None | |
| } | |
| self.subscribers[task_id] = [] | |
| logger.info(f"Created progress tracker for task {task_id}") | |
| def update_progress( | |
| self, | |
| task_id: str, | |
| current_step: int, | |
| message: str, | |
| status: str = "processing" | |
| ): | |
| """Update progress for a task (synchronous)""" | |
| if task_id not in self.tasks: | |
| logger.warning(f"Task {task_id} not found for progress update") | |
| return | |
| task = self.tasks[task_id] | |
| task["current_step"] = current_step | |
| task["progress"] = int((current_step / task["total_steps"]) * 100) | |
| task["message"] = message | |
| task["status"] = status | |
| logger.info(f"Task {task_id}: {task['progress']}% - {message}") | |
| def complete_task(self, task_id: str, result_url: Optional[str] = None): | |
| """Mark a task as completed""" | |
| if task_id not in self.tasks: | |
| return | |
| self.tasks[task_id].update({ | |
| "status": "completed", | |
| "progress": 100, | |
| "message": "Enhancement completed!", | |
| "result_url": result_url, | |
| "completed_at": datetime.now().isoformat() | |
| }) | |
| # Notify subscribers if event loop is running | |
| try: | |
| asyncio.create_task(self._notify_subscribers(task_id)) | |
| except RuntimeError: | |
| # No event loop running (called from thread) | |
| pass | |
| logger.info(f"Task {task_id} completed") | |
| def fail_task(self, task_id: str, error: str): | |
| """Mark a task as failed""" | |
| if task_id not in self.tasks: | |
| return | |
| self.tasks[task_id].update({ | |
| "status": "failed", | |
| "message": f"Error: {error}", | |
| "error": error, | |
| "failed_at": datetime.now().isoformat() | |
| }) | |
| # Notify subscribers if event loop is running | |
| try: | |
| asyncio.create_task(self._notify_subscribers(task_id)) | |
| except RuntimeError: | |
| # No event loop running (called from thread) | |
| pass | |
| logger.error(f"Task {task_id} failed: {error}") | |
| def get_task_status(self, task_id: str) -> Optional[Dict]: | |
| """Get current status of a task""" | |
| return self.tasks.get(task_id) | |
| def subscribe(self, task_id: str, queue: asyncio.Queue): | |
| """Subscribe to progress updates for a task""" | |
| if task_id not in self.subscribers: | |
| self.subscribers[task_id] = [] | |
| self.subscribers[task_id].append(queue) | |
| logger.debug(f"New subscriber for task {task_id}") | |
| def unsubscribe(self, task_id: str, queue: asyncio.Queue): | |
| """Unsubscribe from progress updates""" | |
| if task_id in self.subscribers and queue in self.subscribers[task_id]: | |
| self.subscribers[task_id].remove(queue) | |
| async def _notify_subscribers(self, task_id: str): | |
| """Send current task status to all subscribers""" | |
| if task_id not in self.subscribers: | |
| return | |
| task_data = self.tasks.get(task_id) | |
| if not task_data: | |
| return | |
| # Send to all subscribers | |
| for queue in self.subscribers[task_id]: | |
| try: | |
| await queue.put(task_data.copy()) | |
| except Exception as e: | |
| logger.error(f"Error notifying subscriber: {e}") | |
| def cleanup_task(self, task_id: str): | |
| """Remove task data after completion""" | |
| if task_id in self.tasks: | |
| del self.tasks[task_id] | |
| if task_id in self.subscribers: | |
| del self.subscribers[task_id] | |
| logger.debug(f"Cleaned up task {task_id}") | |
| # Global progress tracker instance | |
| progress_tracker = ProgressTracker() | |