pragativeer-v2/backend/worker.py

import os
from datetime import datetime, timedelta
from celery import Celery
from celery.schedules import crontab
from sqlmodel import Session, select
import json
import asyncio

from database import engine
import models
import storage
from pdf_parser import parse_pdf_to_questions

REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")

# Initialize Celery connected to our local Redis broker
celery_app = Celery(
    "pragativeer_tasks",
    broker=REDIS_URL,
    backend=REDIS_URL
)

# Configure timezone to IST for accurate 5:00 AM scheduling
celery_app.conf.update(
    timezone='Asia/Kolkata',
    enable_utc=True,
)

@celery_app.task(name="generate_daily_study_plans", bind=True, default_retry_delay=60, max_retries=3)
def generate_daily_study_plans(self):
    """
    Runs daily to evaluate students' weak areas and generate 
    a personalized AI study plan for the day.
    """
    try:
        with Session(engine) as session:
            # Fetch all active users on the AI Coach tier
            ai_students = session.exec(
                select(models.User).where(models.User.tier == models.TierEnum.AI_COACH)
            ).all()

            today = datetime.utcnow().date()
            plans_created = 0

            for student in ai_students:
                # 1. Fetch pending spaced repetition mistakes for today
                pending_mistakes = session.exec(
                    select(models.Mistake)
                    .where(models.Mistake.user_id == student.id)
                    .where(models.Mistake.is_mastered == False)
                    .where(models.Mistake.next_review <= datetime.utcnow())
                ).all()

                # 2. Extract weak topics (for revision recommendation)
                topics = [m.topic for m in pending_mistakes]
                most_weak_topic = max(set(topics), key=topics.count) if topics else "Mixed Syllabus"

                # 3. Construct the time-budgeted dynamic task list
                tasks = [
                    {"type": "Revision", "description": f"Revise {most_weak_topic}", "duration_minutes": 15, "status": "Pending"},
                    {"type": "Mistake_Review", "description": "Daily Mistake Review Session", "duration_minutes": 10, "status": "Pending"},
                    {"type": "DPS", "description": f"Attempt {most_weak_topic} Practice Sheet", "duration_minutes": 20, "status": "Pending"}
                ]

                # 4. Save to Database
                new_plan = models.StudyPlan(
                    user_id=student.id,
                    date=today,
                    tasks=tasks,
                    is_completed=False
                )
                session.add(new_plan)
                plans_created += 1
                
            session.commit()
            return f"Successfully generated {plans_created} daily study plans."
    except Exception as exc:
        print(f"Database error during study plan generation: {exc}")
        raise self.retry(exc=exc)

@celery_app.task(name="process_pdf_test", bind=True, max_retries=1)
def process_pdf_test(self, test_id_str: str, object_name: str, admin_id_str: str, exam_pattern: str):
    """
    Background task to download PDF from MinIO, extract text, 
    use LLM to generate questions, and attach them to a Test record.
    """
    try:
        print(f"Starting PDF processing for test {test_id_str}")
        
        # 1. Download PDF from MinIO
        pdf_bytes = storage.get_file(object_name)
        if not pdf_bytes:
            print("Failed to download PDF from MinIO")
            return
            
        # 2. Extract Questions using LLM (Async)
        questions_data = asyncio.run(parse_pdf_to_questions(pdf_bytes, exam_pattern))
        if not questions_data:
            print("LLM parsing returned empty data")
            return
            
        import uuid
        test_id = uuid.UUID(test_id_str)
        admin_id = uuid.UUID(admin_id_str)
        
        with Session(engine) as session:
            test_obj = session.get(models.Test, test_id)
            if not test_obj:
                print("Test record not found in DB")
                return

            question_ids = []
            for q_data in questions_data:
                # Create a Question DB object
                new_q = models.Question(
                    subject=q_data.get("subject", "Physics"),
                    topic=q_data.get("topic", "General"),
                    subtopic=q_data.get("subtopic", "General"),
                    difficulty=q_data.get("difficulty", "Medium"),
                    type=q_data.get("type", models.QuestionTypeEnum.MCQ_SINGLE),
                    question_text=q_data.get("question_text", ""),
                    options=q_data.get("options", {}),
                    correct_answer=q_data.get("correct_answer", {}),
                    solution_text=q_data.get("solution_text", ""),
                    language=models.MediumEnum.ENGLISH,
                    exam=models.ExamEnum(exam_pattern),
                    created_by=admin_id
                )
                session.add(new_q)
                session.flush() # To get the ID
                question_ids.append(str(new_q.id))
                
            # Update Test record with the new sections
            test_obj.sections = {
                "Section 1": {
                    "questions": question_ids,
                    "time_limit": test_obj.total_time
                }
            }
            # Automatically activate it when done processing
            test_obj.is_active = True
            
            session.commit()
            print(f"Successfully processed PDF and created {len(question_ids)} questions for test {test_id_str}")
            return f"Processed {len(question_ids)} questions"
    except Exception as exc:
        print(f"Error processing PDF test: {exc}")
        # Mark test as failed if we had a status field, for now just log it
        raise self.retry(exc=exc)

# Schedule the tasks using Celery Beat
celery_app.conf.beat_schedule = {
    'generate-study-plans-every-morning': {
        'task': 'generate_daily_study_plans',
        'schedule': crontab(hour=5, minute=0), # 5:00 AM IST
    },
}