"""Async quiz extraction task with step-by-step progress reporting via Redis.""" import json import logging import time from app.tasks import celery_app from app.database import SessionLocal logger = logging.getLogger(__name__) EXPIRE_SECONDS = 3600 # keep progress for 1 hour def _redis(): import redis from app.config import settings return redis.from_url(settings.REDIS_URL, decode_responses=True) def _push_step(r, job_id: str, step: str, message: str): key = f"extraction:steps:{job_id}" entry = json.dumps({"step": step, "message": message, "ts": time.time()}) r.rpush(key, entry) r.expire(key, EXPIRE_SECONDS) @celery_app.task(name="extract_quiz", bind=True) def extract_quiz( self, job_id: str, user_id: int, section_id: int, title: str, mode: str, time_limit_minutes: int | None, model_id: str | None, question_category_id: int | None, ): """Background quiz extraction task with live progress reporting.""" r = _redis() r.set(f"extraction:status:{job_id}", "running", ex=EXPIRE_SECONDS) # Ensure user job index exists (in case set before task ran) r.lpush(f"extraction:user_jobs:{user_id}", job_id) r.expire(f"extraction:user_jobs:{user_id}", 86400) db = SessionLocal() try: from app.models.section import Section from app.models.pdf_document import PDFDocument from app.services import ai_service, vector_service, pdf_service, embedding_service from app.models.quiz import Quiz from app.models.question import Question from app.models.quiz_question_link import QuizQuestionLink from app.config import settings import os _push_step(r, job_id, "start", "Starting extraction…") section = db.query(Section).filter(Section.id == section_id).first() if not section: raise ValueError("Section not found") document = db.query(PDFDocument).filter(PDFDocument.id == section.document_id).first() if not document: raise ValueError("Document not found") total_pages = section.end_page - section.start_page + 1 _push_step(r, job_id, "text", f"Loading text from pages {section.start_page}–{section.end_page} ({total_pages} pages)…") # Determine model first (used for chunk logging) if model_id: from app.models.ai_model_config import AIModelConfig config = db.query(AIModelConfig).filter(AIModelConfig.model_id == model_id).first() api_key = config.api_key if config and config.api_key else None model_name = config.name if config else model_id else: model_id, api_key = ai_service.get_model_for_task(db, "extraction") model_name = model_id # For large page ranges, process in chunks of 50 pages to avoid truncation CHUNK_PAGES = 50 import json as _json if total_pages <= CHUNK_PAGES: chunks = [(section.start_page, section.end_page)] else: chunks = [] p = section.start_page while p <= section.end_page: end = min(p + CHUNK_PAGES - 1, section.end_page) chunks.append((p, end)) p = end + 1 n_chunks = len(chunks) if n_chunks > 1: _push_step(r, job_id, "text", f"Large section: splitting into {n_chunks} chunks of up to {CHUNK_PAGES} pages each.") def _normalize_ocr(text: str) -> str: """Normalize common OCR artifacts in PREP PDFs.""" return (text .replace("Pref erred", "Preferred") .replace("Pre ferred", "Preferred") .replace("Prefer red", "Preferred") .replace("ltem", "Item") .replace("ltcm", "Item")) # --- Detect separate answer key section (e.g. PREP 2013) --- # Scan document in 10-page steps to find where "Preferred Response:" first appears answer_section_start = None scan_step = 10 for scan_p in range(section.start_page, section.end_page, scan_step): scan_end = min(scan_p + scan_step - 1, section.end_page) raw_chunk = vector_service.get_pages_text( document_id=section.document_id, start_page=scan_p, end_page=scan_end, ) if not raw_chunk: continue scan_chunk = _normalize_ocr(raw_chunk) if "Preferred Response:" in scan_chunk or "preferred response:" in scan_chunk.lower(): answer_section_start = max(section.start_page, scan_p - 5) break has_end_answer_key = answer_section_start is not None if has_end_answer_key: _push_step(r, job_id, "ai", f"Detected separate answer key section starting around page {answer_section_start}. Using two-phase extraction.") # Restrict question chunks to pages strictly BEFORE the answer section # Cap the end page so no chunk bleeds into the answer section q_end = answer_section_start - 1 chunks = [(s, min(e, q_end)) for s, e in chunks if s <= q_end] if not chunks: chunks = [(section.start_page, q_end)] all_valid_questions = [] all_skipped = [] if has_end_answer_key: # === PHASE 1: Extract questions (allow null correct_answer) === raw_questions = [] for chunk_idx, (start_p, end_p) in enumerate(chunks, 1): _push_step(r, job_id, "ai", f"Phase 1 – Chunk {chunk_idx}/{n_chunks}: pages {start_p}–{end_p} (questions)…") chunk_content = vector_service.get_pages_text( document_id=section.document_id, start_page=start_p, end_page=end_p, ) if not chunk_content: continue try: chunk_qs = ai_service.extract_questions_no_answers( _normalize_ocr(chunk_content), page_info=f"{start_p}-{end_p}", page_ref=start_p, model_id=model_id, api_key=api_key, ) raw_questions.extend(chunk_qs) _push_step(r, job_id, "ai", f" Pages {start_p}–{end_p}: {len(chunk_qs)} questions found. Total: {len(raw_questions)}.") except Exception as e: _push_step(r, job_id, "ai", f" Pages {start_p}–{end_p} failed: {e}. Continuing…") # === PHASE 2: Extract answer key from the answer section === _push_step(r, job_id, "ai", f"Phase 2 – Extracting answer key from pages {answer_section_start}–{section.end_page}…") # Process answer section in chunks too (it may be large) answer_key: dict = {} for ans_start in range(answer_section_start, section.end_page + 1, CHUNK_PAGES): ans_end = min(ans_start + CHUNK_PAGES - 1, section.end_page) answer_content = vector_service.get_pages_text( document_id=section.document_id, start_page=ans_start, end_page=ans_end, ) if not answer_content: continue chunk_key = ai_service.extract_answer_key( _normalize_ocr(answer_content), page_info=f"{ans_start}-{ans_end}", model_id=model_id, api_key=api_key, ) answer_key.update(chunk_key) _push_step(r, job_id, "ai", f" Answer pages {ans_start}–{ans_end}: +{len(chunk_key)} answers. Total: {len(answer_key)}.") _push_step(r, job_id, "ai", f" Answer key complete: {len(answer_key)} items.") # === PHASE 3: Match questions to answers === _push_step(r, job_id, "ai", "Phase 3 – Matching questions to answers…") for q in raw_questions: item_num = q.get("item_number") if not item_num: # No item number — skip all_skipped.append(q.get("question_text", "")[:120]) continue letter = answer_key.get(item_num) if not letter: all_skipped.append(q.get("question_text", "")[:120]) continue # Resolve letter → full option text options = q.get("options") or [] letter_idx = ord(letter.upper()) - ord("A") if 0 <= letter_idx < len(options): q["correct_answer"] = options[letter_idx] all_valid_questions.append(q) else: all_skipped.append(q.get("question_text", "")[:120]) _push_step(r, job_id, "ai", f"Matching complete: {len(all_valid_questions)} matched, {len(all_skipped)} unmatched.") else: # === STANDARD: inline correct answer format === for chunk_idx, (start_p, end_p) in enumerate(chunks, 1): if n_chunks > 1: _push_step(r, job_id, "ai", f"Chunk {chunk_idx}/{n_chunks}: pages {start_p}–{end_p} → {model_name}…") else: _push_step(r, job_id, "ai", f"Sending pages {start_p}–{end_p} to {model_name}…") chunk_content = vector_service.get_pages_text( document_id=section.document_id, start_page=start_p, end_page=end_p, ) if not chunk_content: _push_step(r, job_id, "ai", f" No text found for pages {start_p}–{end_p}, skipping.") continue try: chunk_data = ai_service.extract_questions( chunk_content, page_info=f"{start_p}-{end_p}", page_ref=start_p, model_id=model_id, api_key=api_key, ) chunk_skipped = [] if chunk_data and chunk_data[0].get("skipped"): chunk_skipped = chunk_data[0].pop("skipped") chunk_valid = [q for q in chunk_data if q.get("correct_answer")] all_valid_questions.extend(chunk_valid) all_skipped.extend(chunk_skipped) _push_step(r, job_id, "ai", f" Pages {start_p}–{end_p}: {len(chunk_valid)} questions extracted{f', {len(chunk_skipped)} skipped' if chunk_skipped else ''}. Total so far: {len(all_valid_questions)}.") except Exception as e: _push_step(r, job_id, "ai", f" Pages {start_p}–{end_p} failed: {e}. Continuing…") valid_questions = all_valid_questions skipped = all_skipped _push_step(r, job_id, "ai", f"Extraction complete: {len(valid_questions)} total valid questions{f', {len(skipped)} skipped' if skipped else ''}.") if not valid_questions: raise ValueError("No valid questions extracted. The AI could not find questions with correct answers in this page range.") # Extract images _push_step(r, job_id, "images", "Extracting question images…") file_path = os.path.join(settings.UPLOAD_DIR, document.filename) page_images = {} if os.path.exists(file_path): try: page_images = pdf_service.extract_all_images( file_path, document.id, section.start_page, section.end_page ) except Exception as e: _push_step(r, job_id, "images", f"Image extraction skipped: {e}") # Create quiz quiz = Quiz( section_id=section_id, user_id=user_id, title=title, questions_count=len(valid_questions), mode=mode, time_limit_minutes=time_limit_minutes, skipped_questions=_json.dumps(skipped) if skipped else None, ) db.add(quiz) db.flush() _push_step(r, job_id, "save", f"Saving {len(valid_questions)} questions and generating embeddings…") for pos, q in enumerate(valid_questions): page_ref = q.get("page_reference") image_path = None if page_ref and page_ref in page_images and page_images[page_ref]: image_path = page_images[page_ref].pop(0) if not page_images[page_ref]: del page_images[page_ref] question = Question( source_quiz_id=quiz.id, question_category_id=question_category_id, question_text=q["question_text"], question_type=q["question_type"], options=q.get("options"), correct_answer=q["correct_answer"], explanation=q.get("explanation", ""), page_reference=page_ref, image_path=image_path, ) db.add(question) db.flush() db.add(QuizQuestionLink(quiz_id=quiz.id, question_id=question.id, position=pos)) try: embedding_service.embed_question(question) except Exception: pass db.commit() db.refresh(quiz) _push_step(r, job_id, "done", f"Quiz ready! {len(valid_questions)} questions extracted and saved.") r.set(f"extraction:status:{job_id}", "completed", ex=EXPIRE_SECONDS) r.set(f"extraction:quiz_id:{job_id}", str(quiz.id), ex=EXPIRE_SECONDS) return quiz.id except Exception as e: logger.exception(f"Quiz extraction failed for job {job_id}") _push_step(r, job_id, "error", f"Extraction failed: {e}") r.set(f"extraction:status:{job_id}", "failed", ex=EXPIRE_SECONDS) r.set(f"extraction:error:{job_id}", str(e)[:500], ex=EXPIRE_SECONDS) raise finally: db.close()