qwen_agent/task_queue/manager.py

#!/usr/bin/env python3
"""
Queue manager for handling file processing queues.
"""

import os
import json
import time
import logging
from typing import Dict, List, Optional, Any
from huey import Huey
from huey.api import Task
from datetime import datetime, timedelta

# Configure logging
logger = logging.getLogger('app')

from .config import huey
from .tasks import process_file_async, process_multiple_files_async, process_zip_file_async, cleanup_processed_files


class QueueManager:
    """Queue manager for file processing tasks."""

    def __init__(self):
        self.huey = huey
        logger.info("Queue manager initialized with Redis backend")

    def enqueue_file(
        self,
        project_id: str,
        file_path: str,
        original_filename: str = None,
        delay: int = 0
    ) -> str:
        """
        Add a file to the processing queue.

        Args:
            project_id: Project ID
            file_path: File path
            original_filename: Original filename
            delay: Delay before execution in seconds

        Returns:
            Task ID
        """
        if delay > 0:
            task = process_file_async.schedule(
                args=(project_id, file_path, original_filename),
                delay=timedelta(seconds=delay)
            )
        else:
            task = process_file_async(project_id, file_path, original_filename)

        logger.info(f"File queued for processing: {file_path}, task ID: {task.id}")
        return task.id

    def enqueue_multiple_files(
        self,
        project_id: str,
        file_paths: List[str],
        original_filenames: List[str] = None,
        delay: int = 0
    ) -> List[str]:
        """
        Add multiple files to the processing queue.

        Args:
            project_id: Project ID
            file_paths: List of file paths
            original_filenames: List of original filenames
            delay: Delay before execution in seconds

        Returns:
            List of task IDs
        """
        if delay > 0:
            task = process_multiple_files_async.schedule(
                args=(project_id, file_paths, original_filenames),
                delay=timedelta(seconds=delay)
            )
        else:
            task = process_multiple_files_async(project_id, file_paths, original_filenames)

        logger.info(f"Batch files queued for processing: {len(file_paths)} files, task ID: {task.id}")
        return [task.id]

    def enqueue_zip_file(
        self,
        project_id: str,
        zip_path: str,
        extract_to: str = None,
        delay: int = 0
    ) -> str:
        """
        Add a zip file to the processing queue.

        Args:
            project_id: Project ID
            zip_path: Path to the zip file
            extract_to: Extraction target directory
            delay: Delay before execution in seconds

        Returns:
            Task ID
        """
        if delay > 0:
            task = process_zip_file_async.schedule(
                args=(project_id, zip_path, extract_to),
                delay=timedelta(seconds=delay)
            )
        else:
            task = process_zip_file_async(project_id, zip_path, extract_to)

        logger.info(f"Zip file queued for processing: {zip_path}, task ID: {task.id}")
        return task.id

    def get_task_status(self, task_id: str) -> Dict[str, Any]:
        """
        Get task status.

        Args:
            task_id: Task ID

        Returns:
            Task status information
        """
        try:
            # Try getting the task result from result storage
            try:
                # Use Huey's built-in result lookup when available
                if hasattr(self.huey, 'result') and self.huey.result:
                    result = self.huey.result(task_id)
                    if result is not None:
                        return {
                            "task_id": task_id,
                            "status": "complete",
                            "result": result
                        }
            except Exception:
                pass

            # Check whether the task is in the pending queue
            try:
                pending_tasks = list(self.huey.pending())
                for task in pending_tasks:
                    if hasattr(task, 'id') and task.id == task_id:
                        return {
                            "task_id": task_id,
                            "status": "pending"
                        }
            except Exception:
                pass

            # Check whether the task is in the scheduled queue
            try:
                scheduled_tasks = list(self.huey.scheduled())
                for task in scheduled_tasks:
                    if hasattr(task, 'id') and task.id == task_id:
                        return {
                            "task_id": task_id,
                            "status": "scheduled"
                        }
            except Exception:
                pass

            # If not found anywhere, it may not exist or may have completed with cleaned results
            return {
                "task_id": task_id,
                "status": "unknown",
                "message": "Task status is unknown; it may already be complete or may not exist"
            }

        except Exception as e:
            return {
                "task_id": task_id,
                "status": "error",
                "message": f"Failed to get task status: {str(e)}"
            }

    def get_queue_stats(self) -> Dict[str, Any]:
        """
        Get queue statistics.

        Returns:
            Queue statistics information
        """
        try:
            # Use a simplified approach for queue statistics
            stats = {
                "total_tasks": 0,
                "pending_tasks": 0,
                "running_tasks": 0,
                "completed_tasks": 0,
                "error_tasks": 0,
                "scheduled_tasks": 0,
                "recent_tasks": [],
                "queue_backend": "redis"
            }

            # Try to get the number of pending tasks
            try:
                pending_tasks = list(self.huey.pending())
                stats["pending_tasks"] = len(pending_tasks)
                stats["total_tasks"] += len(pending_tasks)
            except Exception as e:
                logger.error(f"Failed to get pending tasks: {e}")

            # Try to get the number of scheduled tasks
            try:
                scheduled_tasks = list(self.huey.scheduled())
                stats["scheduled_tasks"] = len(scheduled_tasks)
                stats["total_tasks"] += len(scheduled_tasks)
            except Exception as e:
                logger.error(f"Failed to get scheduled tasks: {e}")

            return stats

        except Exception as e:
            return {
                "error": str(e),
                "queue_backend": "redis"
            }


# Global singleton instance
queue_manager = QueueManager()