backend/tasks/advanced_unified_evaluation_task.py

"""
ADVANCED UNIFIED EVALUATION CELERY TASK: Background execution with LLM Judge and Multi-turn.

This task provides asynchronous execution of the advanced unified pipeline
through Celery workers for production scalability with all advanced features.
"""

import logging
from celery import Celery
from typing import Dict, Any, List

from workers.celery_worker import celery_app
from services.evaluation_service import EvaluationService
from core.database import get_db

logger = logging.getLogger(__name__)


@celery_app.task(bind=True)
def run_advanced_unified_evaluation_task(
    self, 
    evaluation_id: str, 
    user_id: int,
    enable_multi_turn: bool = True,
    enable_llm_judge: bool = True,
    max_conversation_turns: int = 3,
    judge_evaluation_types: List[str] = None
) -> Dict[str, Any]:
    """
    Celery task for advanced unified evaluation execution.
    
    This task runs the complete advanced unified pipeline in the background:
    - Dataset loading and preparation
    - Learning engine insights
    - Multi-turn conversation attacks (if enabled)
    - LLM Judge evaluation (if enabled)
    - Advanced scoring and audit trails
    - Analytics data flow
    - Enhanced report generation
    
    Args:
        evaluation_id: Evaluation ID to execute
        user_id: User ID requesting execution
        enable_multi_turn: Enable multi-turn conversation attacks
        enable_llm_judge: Enable LLM judge evaluation
        max_conversation_turns: Maximum turns per conversation
        judge_evaluation_types: Types of judge evaluations to perform
        
    Returns:
        Dict with advanced execution results
    """
    logger.info(f"🚀 Starting advanced unified evaluation task: {evaluation_id}")
    
    try:
        # Create async session for database operations
        async def execute_advanced_evaluation():
            async for db in get_db():
                # Initialize evaluation service
                evaluation_service = EvaluationService(db)
                
                # Execute advanced unified evaluation
                result = await evaluation_service.execute_advanced_unified_evaluation(
                    evaluation_id=evaluation_id,
                    user_id=user_id,
                    enable_multi_turn=enable_multi_turn,
                    enable_llm_judge=enable_llm_judge,
                    max_conversation_turns=max_conversation_turns,
                    judge_evaluation_types=judge_evaluation_types
                )
                
                logger.info(f"✅ Advanced unified evaluation task completed: {evaluation_id}")
                return result
        
        # Run async function in sync context
        import asyncio
        return asyncio.run(execute_advanced_evaluation())
        
    except Exception as e:
        logger.error(f"❌ Advanced unified evaluation task failed: {evaluation_id} - {str(e)}")
        
        # Update task status and return error
        self.update_state(
            state='FAILURE',
            meta={'error': str(e), 'evaluation_id': evaluation_id}
        )
        
        raise


@celery_app.task
def cleanup_advanced_evaluation_resources(evaluation_id: str) -> Dict[str, Any]:
    """
    Cleanup task for advanced unified evaluation resources.
    
    This task cleans up temporary resources after advanced evaluation completion:
    - Temporary conversation states
    - Judge evaluation caches
    - Memory caches
    - Connection pools
    
    Args:
        evaluation_id: Evaluation ID to cleanup
        
    Returns:
        Dict with cleanup results
    """
    logger.info(f"🧹 Starting cleanup for advanced unified evaluation: {evaluation_id}")
    
    try:
        cleanup_results = {
            "evaluation_id": evaluation_id,
            "conversation_states_cleaned": 0,
            "judge_cache_cleared": True,
            "memory_cleared": True,
            "connections_closed": True,
            "cleanup_completed_at": None
        }
        
        # Implement cleanup logic here
        # This would integrate with the advanced unified pipeline cleanup methods
        
        logger.info(f"✅ Cleanup completed for advanced unified evaluation: {evaluation_id}")
        return cleanup_results
        
    except Exception as e:
        logger.error(f"❌ Cleanup failed for advanced unified evaluation: {evaluation_id} - {str(e)}")
        raise


# Task chaining for complete advanced workflow
def create_advanced_unified_evaluation_workflow(
    evaluation_id: str, 
    user_id: int,
    enable_multi_turn: bool = True,
    enable_llm_judge: bool = True,
    max_conversation_turns: int = 3,
    judge_evaluation_types: List[str] = None
):
    """
    Create a complete workflow for advanced unified evaluation.
    
    This chains the execution and cleanup tasks for a complete advanced workflow.
    
    Args:
        evaluation_id: Evaluation ID
        user_id: User ID
        enable_multi_turn: Enable multi-turn conversation attacks
        enable_llm_judge: Enable LLM judge evaluation
        max_conversation_turns: Maximum turns per conversation
        judge_evaluation_types: Types of judge evaluations to perform
        
    Returns:
        Celery chain result
    """
    from celery import chain
    
    workflow = chain(
        run_advanced_unified_evaluation_task.s(
            evaluation_id, 
            user_id, 
            enable_multi_turn,
            enable_llm_judge,
            max_conversation_turns,
            judge_evaluation_types
        ),
        cleanup_advanced_evaluation_resources.s(evaluation_id)
    )
    
    return workflow()