app/workers/tasks.py

"""
Celery task definitions.

process_file
    Main file-processing pipeline.  Dispatches to the correct processor class
    (PDF / DOCX / XLSX / Image / Audio / Video), runs extract → summarise →
    chunk → embed → index, and updates the Job row at every step.
    For audio and video files, per-speaker SpeechBrain ECAPA embeddings are
    attached to each chunk's ChromaDB metadata before indexing.
    Retries up to CELERY_MAX_RETRIES times with exponential back-off on
    rate-limit and unknown errors; jumps straight to FAILED_PERMANENT on
    invalid-input (400) errors.  Permanently failed jobs are pushed to a
    Redis dead-letter queue for manual review.

compute_ragas
    Async RAGAS quality evaluation triggered after every successful RAG query.
    Re-embeds the original question, re-retrieves context chunks, and calls
    compute_ragas_scores().  Stores the five metric scores in QueryHistory.

cleanup_old_uploads
    Beat-scheduled daily task that deletes on-disk upload directories for
    jobs that have been in a terminal state (COMPLETED / FAILED_PERMANENT)
    for more than 7 days.
"""

import json
import time
import uuid
from datetime import datetime
from typing import Optional

from sqlmodel import Session

from app.models.db import Job, JobStatus, ErrorType, get_engine
from app.observability.logging import get_logger
from app.workers.celery_app import celery_app

log = get_logger()


# ── State transition helper ────────────────────────────────────────────────────

def update_job_state(
    db: Session,
    job_id,
    status: JobStatus,
    step: Optional[str] = None,
    error_type: Optional[ErrorType] = None,
    error_message: Optional[str] = None,
    chunk_count: Optional[int] = None,
) -> None:
    job = db.get(Job, job_id if isinstance(job_id, uuid.UUID) else uuid.UUID(str(job_id)))
    from_status = job.status.value
    job.status = status
    job.updated_at = datetime.utcnow()
    if step is not None:
        job.step = step
    if error_type is not None:
        job.error_type = error_type
    if error_message is not None:
        job.error_message = error_message
    if chunk_count is not None:
        job.chunk_count = chunk_count
    db.add(job)
    db.commit()
    log.info(
        "job_state_change",
        job_id=str(job_id),
        from_status=from_status,
        to_status=status.value,
        step=step,
        retry_count=job.retry_count,
    )


# ── Error classification ───────────────────────────────────────────────────────

def classify_error(exc: Exception) -> tuple[str, bool]:
    msg = str(exc).lower()
    if "429" in msg or "quota" in msg or "rate" in msg:
        return "RATE_LIMIT", True
    if "400" in msg or "invalid" in msg:
        return "INVALID_INPUT", False
    return "UNKNOWN", True


# ── Main processing task ───────────────────────────────────────────────────────

@celery_app.task(
    bind=True,
    max_retries=3,
    default_retry_delay=60,
)
def process_file(self, job_id: str):
    from app.config import settings
    from app.observability.tracing import tracer

    with tracer.start_as_current_span("process_file") as span:
        span.set_attribute("job_id", job_id)

        with Session(get_engine()) as db:
            job = db.get(Job, uuid.UUID(job_id))
            if not job:
                log.error("process_file_job_not_found", job_id=job_id)
                return

            if job.status == JobStatus.completed:
                log.info("process_file_already_completed", job_id=job_id)
                return

            try:
                span.set_attribute("file_type", job.file_type)
                span.set_attribute("user_id", str(job.user_id))
                update_job_state(db, job_id, JobStatus.processing, step="extracting")

                file_type = job.file_type
                log.info("process_file_start", job_id=job_id, file_type=file_type)

                # ── Dispatch to correct processor ──────────────────────────
                if file_type == "pdf":
                    from app.processors.pdf import PDFProcessor
                    processor = PDFProcessor(job=job, settings=settings)
                elif file_type == "docx":
                    from app.processors.docx_proc import DOCXProcessor
                    processor = DOCXProcessor(job=job, settings=settings)
                elif file_type in ("xlsx", "csv"):
                    from app.processors.xlsx_proc import XLSXProcessor
                    processor = XLSXProcessor(job=job, settings=settings)
                elif file_type == "image":
                    from app.processors.image import ImageProcessor
                    processor = ImageProcessor(job=job, settings=settings)
                elif file_type == "audio":
                    from app.processors.audio import AudioProcessor
                    processor = AudioProcessor(job=job, settings=settings)
                elif file_type == "video":
                    from app.processors.video import VideoProcessor
                    processor = VideoProcessor(job=job, settings=settings)
                else:
                    raise ValueError(f"Unsupported file_type: {file_type}")

                # ── Extract + summarise ────────────────────────────────────
                extracted_text, summary = processor.run(db)

                # ── Chunking ──────────────────────────────────────────────
                update_job_state(db, job_id, JobStatus.processing, step="chunking")

                from app.rag.chunker import chunk_markdown_hierarchical

                chunks = chunk_markdown_hierarchical(
                    extracted_text,
                    job_id=job_id,
                    filename=job.filename,
                    file_type=file_type,
                    parent_size=settings.CHUNK_SIZE,
                    child_size=settings.CHILD_CHUNK_SIZE,
                )

                if not chunks:
                    log.warning("no_chunks_produced", job_id=job_id, file_type=file_type)

                # ── Attach SpeechBrain ECAPA embeddings to audio/video chunks ──
                # speaker_embeddings is a dict {speaker_label: [float, ...]} produced
                # by diarize_audio(return_embeddings=True).  base.py pops it from the
                # summary before DB serialisation to avoid storing 192-float lists in
                # Job.result, then restores it so this task can consume it.
                # Each child chunk's markdown header looks like [Speaker 1 at 00:05],
                # so we extract the speaker label with a regex and store the mean ECAPA
                # embedding for that speaker as JSON in ChromaDB metadata.
                if file_type in ("audio", "video") and chunks:
                    import re as _re
                    import json as _json
                    speaker_embeddings = summary.get("_speaker_embeddings") or {}
                    if speaker_embeddings:
                        _speaker_re = _re.compile(r'\[Speaker (\d+) at')
                        for chunk in chunks:
                            m = _speaker_re.search(chunk["text"])
                            if m:
                                speaker_label = f"Speaker {m.group(1)}"
                                emb = speaker_embeddings.get(speaker_label)
                                if emb:
                                    chunk["metadata"]["speaker_label"] = speaker_label
                                    chunk["metadata"]["speaker_embedding_json"] = _json.dumps(emb)
                        log.info(
                            "speaker_embeddings_tagged",
                            job_id=job_id,
                            speaker_count=len(speaker_embeddings),
                        )

                # ── Embedding ─────────────────────────────────────────────
                update_job_state(db, job_id, JobStatus.processing, step="embedding")

                embeddings = []
                if chunks:
                    from app.rag.embedder import embed_chunks
                    embeddings = embed_chunks(chunks, job.user_id, job.id, settings, db)

                # ── Indexing ──────────────────────────────────────────────
                update_job_state(db, job_id, JobStatus.processing, step="indexing")

                if chunks and embeddings:
                    from app.rag.vectorstore import get_chroma_client, get_or_create_collection, delete_job_chunks, add_chunks
                    from app.rag.bm25_index import invalidate_bm25
                    client = get_chroma_client(settings)
                    collection = get_or_create_collection(client, settings)
                    delete_job_chunks(collection, job_id)
                    add_chunks(collection, chunks, embeddings)
                    invalidate_bm25(settings)  # force BM25 rebuild on next query

                # ── Complete ──────────────────────────────────────────────
                update_job_state(
                    db, job_id, JobStatus.completed,
                    step="completed",
                    chunk_count=len(chunks),
                )
                log.info("process_file_completed", job_id=job_id, chunk_count=len(chunks))

            except Exception as exc:
                error_type_str, retryable = classify_error(exc)
                error_type = ErrorType(error_type_str)

                db.refresh(job)
                job.retry_count += 1
                db.add(job)
                db.commit()

                log.error(
                    "process_file_error",
                    job_id=job_id,
                    error_type=error_type_str,
                    retryable=retryable,
                    retry_count=job.retry_count,
                    error=str(exc),
                )

                if retryable and self.request.retries < self.max_retries:
                    update_job_state(
                        db, job_id, JobStatus.failed,
                        step="failed",
                        error_type=error_type,
                        error_message=str(exc)[:500],
                    )
                    countdown = settings.CELERY_RETRY_BACKOFF * (2 ** self.request.retries)
                    raise self.retry(exc=exc, countdown=countdown)
                else:
                    update_job_state(
                        db, job_id, JobStatus.failed_permanent,
                        step="failed",
                        error_type=error_type,
                        error_message=str(exc)[:500],
                    )
                    # Push to dead letter queue for manual review
                    try:
                        import redis as _redis
                        _rc = _redis.from_url(settings.REDIS_URL)
                        import json as _json
                        _rc.rpush("geminirag:dead_letter", _json.dumps({
                            "job_id": job_id,
                            "error": str(exc)[:500],
                            "error_type": error_type_str,
                        }))
                    except Exception:
                        pass


# ── RAGAS evaluation task ─────────────────────────────────────────────────────

@celery_app.task(bind=True, max_retries=2)
def compute_ragas(self, query_history_id: str):
    from app.config import settings
    from app.models.db import QueryHistory

    with Session(get_engine()) as db:
        qh = db.get(QueryHistory, uuid.UUID(query_history_id))
        if not qh:
            log.error("compute_ragas_not_found", query_history_id=query_history_id)
            return

        try:
            from app.evaluation.ragas_eval import compute_ragas_scores
            from app.rag.embedder import embed_query
            from app.rag.vectorstore import (
                get_chroma_client, get_or_create_collection, search
            )
            import json as _json

            job_ids = _json.loads(qh.job_ids_queried) if qh.job_ids_queried else None

            # Re-embed the question to retrieve context chunks
            q_embedding = embed_query(qh.question, settings)
            client = get_chroma_client(settings)
            collection = get_or_create_collection(client, settings)
            chunks = search(collection, q_embedding, top_k=settings.RAG_TOP_K, job_ids=job_ids or None)
            contexts = [c["text"] for c in chunks]

            scores = compute_ragas_scores(
                question=qh.question,
                answer=qh.answer,
                contexts=contexts,
                ground_truth=None,
                settings=settings,
            )

            qh.ragas_scores = _json.dumps(scores)
            qh.ragas_computed_at = datetime.utcnow()
            db.add(qh)
            db.commit()

            log.info(
                "ragas_computed",
                query_id=query_history_id,
                faithfulness=scores.get("faithfulness"),
                answer_relevancy=scores.get("answer_relevancy"),
            )

        except Exception as exc:
            log.error("compute_ragas_error", query_history_id=query_history_id, error=str(exc))
            if self.request.retries < self.max_retries:
                raise self.retry(exc=exc, countdown=60)


# ── Daily file cleanup task ───────────────────────────────────────────────────

@celery_app.task
def cleanup_old_uploads():
    """Delete upload files for jobs older than 7 days that are COMPLETED or FAILED_PERMANENT."""
    import shutil
    from datetime import timedelta
    from pathlib import Path
    from app.config import settings

    cutoff = datetime.utcnow() - timedelta(days=7)
    terminal_statuses = {JobStatus.completed, JobStatus.failed_permanent}

    deleted = 0
    with Session(get_engine()) as db:
        from sqlmodel import select as _select
        stmt = _select(Job).where(
            Job.updated_at < cutoff,
            Job.status.in_([s.value for s in terminal_statuses]),
        )
        old_jobs = db.exec(stmt).all()
        for job in old_jobs:
            job_dir = Path(settings.UPLOAD_DIR) / str(job.id)
            if job_dir.exists():
                try:
                    shutil.rmtree(job_dir)
                    deleted += 1
                except Exception as exc:
                    log.warning("cleanup_delete_failed", job_id=str(job.id), error=str(exc))

    log.info("cleanup_old_uploads_complete", jobs_cleaned=deleted)