app/workers/tasks.py

"""arq tasks. Heavy ML imports live in app.workers.ocr.* (worker image only)."""

import asyncio
import logging
from typing import Any
from uuid import uuid4

from sqlalchemy import select

from app.core.config import settings
from app.db.session import async_session_maker
from app.documents.chunks_model import DocumentChunk
from app.documents.models import Document, DocumentStatus
from app.storage.minio_client import get_minio_client

log = logging.getLogger("mmap.worker")


async def process_document_ocr(ctx: dict[str, Any], document_id: str) -> str:
    # OCR → chunk → embed → Qdrant → graph → summary.
    from app.embeddings import embed_texts
    from app.workers.chunking import chunk_text, is_meaningful
    from app.workers.ocr.pipeline import extract_text_from_bytes

    async with async_session_maker() as db:
        doc = await db.get(Document, document_id)
        if doc is None:
            log.warning("process_document_ocr: doc %s not found", document_id)
            return "missing"

        doc.status = DocumentStatus.PROCESSING
        await db.commit()
        await db.refresh(doc)

        try:
            data = await asyncio.to_thread(_download_bytes, doc.storage_key)
            text = await extract_text_from_bytes(
                data,
                doc.content_type,
                filename=doc.filename,
            )
            doc.extracted_text = text

            raw_chunks = chunk_text(text) if text.strip() else []
            # Drop OCR/PDF noise that would otherwise pollute retrieval. Only
            # filter when the doc produced more than one chunk — for short
            # transcripts (audio, voice notes) we'd rather index a small chunk
            # than nothing.
            if len(raw_chunks) > 1:
                chunks = [c for c in raw_chunks if is_meaningful(c.text)]
            else:
                chunks = raw_chunks
            log.info(
                "doc=%s ocr=%d chars chunks=%d (filtered %d noise)",
                document_id,
                len(text),
                len(chunks),
                len(raw_chunks) - len(chunks),
            )

            chunk_rows: list[DocumentChunk] = []
            for c in chunks:
                row = DocumentChunk(
                    id=uuid4(),
                    document_id=doc.id,
                    chunk_index=c.index,
                    text=c.text,
                    char_start=c.char_start,
                    char_end=c.char_end,
                )
                chunk_rows.append(row)
                db.add(row)
            await db.flush()

            if chunk_rows:
                vectors = await asyncio.to_thread(embed_texts, [r.text for r in chunk_rows])
                await asyncio.to_thread(
                    _upsert_qdrant_points,
                    chunk_rows=chunk_rows,
                    vectors=vectors,
                    user_id=str(doc.user_id),
                    document_id=str(doc.id),
                )

            doc.status = DocumentStatus.PROCESSED
            await db.commit()

            # Graph + summary are best-effort. Vector RAG still works without
            # them, so failures here must not fail the document.
            if text.strip():
                await _ingest_graph(
                    chunks=[r.text for r in chunk_rows] or [text],
                    user_id=str(doc.user_id),
                    document_id=str(doc.id),
                )
                await _summarize_and_store(text=text, document_id=str(doc.id))
            return "processed"

        except (Exception, asyncio.CancelledError) as exc:
            # asyncio.CancelledError is a BaseException in 3.8+, so plain
            # `except Exception` missed it — arq's job-timeout cancel would
            # bypass this handler entirely and leave the doc row wedged at
            # `processing` forever. Catch both and re-raise CancelledError
            # after our cleanup so arq still records the timeout correctly.
            timed_out = isinstance(exc, asyncio.CancelledError)
            log.exception("ocr pipeline failed for doc=%s", document_id)
            # Discard pending chunk inserts so we don't persist orphan rows
            # whose vectors never reached Qdrant.
            try:
                await db.rollback()
                failed_doc = await db.get(Document, document_id)
                if failed_doc is not None:
                    failed_doc.status = DocumentStatus.FAILED
                    failed_doc.error_message = (
                        "OCR timed out — try a smaller or text-based PDF."
                        if timed_out
                        else str(exc)
                    )
                    await db.commit()
            except Exception:
                log.exception("failed to record failure status for doc=%s", document_id)
            if timed_out:
                # Let arq see the cancel so its own bookkeeping stays right.
                raise
            return "failed"


def _download_bytes(storage_key: str) -> bytes:
    client = get_minio_client()
    response = client.get_object(settings.minio_bucket, storage_key)
    try:
        return response.read()
    finally:
        response.close()
        response.release_conn()


def _upsert_qdrant_points(
    *,
    chunk_rows: list[DocumentChunk],
    vectors: list[list[float]],
    user_id: str,
    document_id: str,
) -> None:
    from qdrant_client.http import models as qmodels

    from app.rag.sparse import SPARSE_VECTOR_NAME, encode_passages
    from app.storage.qdrant_client import (
        COLLECTION_NAME,
        DENSE_VECTOR_NAME,
        ensure_collection,
        get_qdrant_client,
    )

    ensure_collection()
    client = get_qdrant_client()

    sparse_vectors = encode_passages([row.text for row in chunk_rows])

    points = [
        qmodels.PointStruct(
            id=str(row.id),
            vector={
                DENSE_VECTOR_NAME: vec,
                SPARSE_VECTOR_NAME: qmodels.SparseVector(indices=s_idx, values=s_val),
            },
            payload={
                "chunk_id": str(row.id),
                "document_id": document_id,
                "user_id": user_id,
                "chunk_index": row.chunk_index,
                "text": row.text,
            },
        )
        for row, vec, (s_idx, s_val) in zip(chunk_rows, vectors, sparse_vectors, strict=True)
    ]
    client.upsert(collection_name=COLLECTION_NAME, points=points, wait=True)


GRAPH_EXTRACT_MAX_ATTEMPTS = 3
GRAPH_EXTRACT_BACKOFF_SECONDS = (15, 30)  # sleep between attempt 1→2 and 2→3


async def _ingest_graph(*, chunks: list[str], user_id: str, document_id: str) -> None:
    # Best-effort: swallow errors so the worker stays non-blocking.
    try:
        from app.embeddings import embed_texts
        from app.graph.alignment import (
            AlignedEntity,
            AlignedRelationship,
            Candidate,
            align_batch,
        )
        from app.graph.extraction import safe_extract_entities
        from app.graph.neo4j_client import (
            ensure_indexes,
            list_entity_candidates,
            list_entity_semantic_candidates,
            upsert_entity,
            upsert_relationship,
        )
        from app.graph.semantic_alignment import (
            SemanticCandidate,
            find_semantic_alias,
            format_entity_text,
        )

        # Groq's json_object validator flakes intermittently — the same
        # extraction that returned 0 entities on upload will often return
        # 20+ on a manual reindex 60-120s later. Retry transient failures
        # in-line so the user's graph populates without a manual click.
        outcome = await safe_extract_entities(chunks)
        for attempt in range(2, GRAPH_EXTRACT_MAX_ATTEMPTS + 1):
            if outcome.result.entities or not outcome.transient_failure:
                break
            sleep_for = GRAPH_EXTRACT_BACKOFF_SECONDS[attempt - 2]
            log.info(
                "graph: transient extraction failure on doc=%s (attempt %d/%d) — sleeping %ds",
                document_id,
                attempt - 1,
                GRAPH_EXTRACT_MAX_ATTEMPTS,
                sleep_for,
            )
            await asyncio.sleep(sleep_for)
            outcome = await safe_extract_entities(chunks)

        if not outcome.result.entities:
            log.info(
                "graph: no entities extracted for doc=%s (transient=%s)",
                document_id,
                outcome.transient_failure,
            )
            return
        result = outcome.result

        await ensure_indexes()

        # Alignment (#43a): pull the user's existing entities as fuzzy-match
        # candidates so this doc's writes converge on the shared KG instead
        # of creating parallel islands ("GWU" vs "gwu.", "Kamalasankaris" vs
        # "Kamalasankari" etc.). One Neo4j fetch per doc; the aligner runs
        # in-memory from there.
        existing_pairs = await list_entity_candidates(user_id)
        existing = [Candidate(name_lower=nl, type=t) for nl, t in existing_pairs]
        aligned_entities, aligned_rels = align_batch(
            result.entities,
            result.relationships,
            existing,
        )
        pre_l3_ent_count = len(aligned_entities)
        pre_l3_rel_count = len(aligned_rels)

        # Semantic alignment (#43c): L3 layer on top of L1/L2. Embed each
        # aligned entity's `name: description`, compare cosine similarity
        # against existing same-type entities' stored embeddings. If a
        # match clears the threshold, rewrite the entity's name_lower to
        # the existing canonical form and cascade the change through the
        # relationships. Catches abbreviation/expansion pairs (SFT ↔
        # Supervised Fine-Tuning) that L1/L2 string-fuzzy can't.
        entity_embeddings: dict[str, list[float]] = {}
        if settings.graph_semantic_align and aligned_entities:
            existing_semantic = await list_entity_semantic_candidates(user_id)
            running_candidates = [
                SemanticCandidate(
                    name_lower=nl,
                    type=t,
                    embedding=tuple(emb),
                )
                for nl, t, emb in existing_semantic
            ]
            texts = [format_entity_text(e.name, e.description) for e in aligned_entities]
            embeddings = await asyncio.to_thread(embed_texts, texts)
            alias_map: dict[str, str] = {}  # old name_lower → canonical
            for e, emb in zip(aligned_entities, embeddings, strict=True):
                match = find_semantic_alias(
                    emb,
                    e.type,
                    running_candidates,
                    settings.graph_semantic_threshold_same,
                    settings.graph_semantic_threshold_cross,
                )
                if match and match != e.name_lower:
                    alias_map[e.name_lower] = match
                else:
                    # Not aliased — this entity becomes a candidate that
                    # later same-batch entities (or future docs) can align to.
                    running_candidates.append(
                        SemanticCandidate(
                            name_lower=e.name_lower,
                            type=e.type,
                            embedding=tuple(emb),
                        )
                    )
                    entity_embeddings[e.name_lower] = emb

            if alias_map:
                seen: set[str] = set()
                collapsed: list[AlignedEntity] = []
                for e in aligned_entities:
                    canonical = alias_map.get(e.name_lower, e.name_lower)
                    if canonical in seen:
                        continue
                    seen.add(canonical)
                    if canonical == e.name_lower:
                        collapsed.append(e)
                    else:
                        # Aliased into a canonical form — keep the aliased
                        # display name; ON CREATE on the existing node is
                        # a no-op so display stays as first-write.
                        collapsed.append(
                            AlignedEntity(
                                name=e.name,
                                name_lower=canonical,
                                type=e.type,
                                description=e.description,
                            )
                        )
                aligned_entities = collapsed

                rewritten_rels: list[AlignedRelationship] = []
                for r in aligned_rels:
                    src = alias_map.get(r.source_lower, r.source_lower)
                    tgt = alias_map.get(r.target_lower, r.target_lower)
                    if src == tgt:
                        continue  # self-loop after L3
                    rewritten_rels.append(
                        AlignedRelationship(
                            source_lower=src,
                            target_lower=tgt,
                            relation=r.relation,
                        )
                    )
                aligned_rels = rewritten_rels

        for e in aligned_entities:
            await upsert_entity(
                user_id=user_id,
                document_id=document_id,
                name=e.name,
                name_lower=e.name_lower,
                entity_type=e.type,
                description=e.description,
                embedding=entity_embeddings.get(e.name_lower),
            )
        for r in aligned_rels:
            await upsert_relationship(
                user_id=user_id,
                document_id=document_id,
                source_lower=r.source_lower,
                target_lower=r.target_lower,
                relation=r.relation,
            )
        log.info(
            "graph: doc=%s chunks=%d entities=%d rels=%d "
            "(pre-l3 %d ents, %d rels; raw %d ents, %d rels)",
            document_id,
            len(chunks),
            len(aligned_entities),
            len(aligned_rels),
            pre_l3_ent_count,
            pre_l3_rel_count,
            len(result.entities),
            len(result.relationships),
        )
    except Exception as exc:  # noqa: BLE001
        log.warning("graph ingest failed (non-blocking): %s", exc)


async def _summarize_and_store(*, text: str, document_id: str) -> None:
    # Best-effort: failure leaves the existing summary columns untouched.
    try:
        from app.agents.summarization import summarize_document

        result = await summarize_document(text)
        if result.is_empty():
            log.info("summary: no content produced for doc=%s", document_id)
            return

        async with async_session_maker() as db:
            doc = await db.get(Document, document_id)
            if doc is None:
                return
            doc.summary_tldr = result.tldr or None
            doc.summary_key_points = list(result.key_points) or None
            doc.summary_topics = list(result.topics) or None
            await db.commit()
        log.info(
            "summary: doc=%s tldr=%d chars points=%d topics=%d",
            document_id,
            len(result.tldr),
            len(result.key_points),
            len(result.topics),
        )
    except Exception as exc:  # noqa: BLE001
        log.warning("summarization failed (non-blocking): %s", exc)


async def resummarize_document(ctx: dict[str, Any], document_id: str) -> str:
    # Re-run summarization on stored chunks. No re-OCR or re-embed.
    async with async_session_maker() as db:
        doc = await db.get(Document, document_id)
        if doc is None:
            return "missing"

        if str(doc.status) != DocumentStatus.PROCESSED.value:
            log.info(
                "resummarize: doc %s status=%s — skipping (expected 'processed')",
                document_id,
                doc.status,
            )
            return "skipped"

        stmt = (
            select(DocumentChunk)
            .where(DocumentChunk.document_id == doc.id)
            .order_by(DocumentChunk.chunk_index.asc())
        )
        rows = (await db.execute(stmt)).scalars().all()

    text = "\n\n".join(r.text for r in rows) if rows else (doc.extracted_text or "")
    if not text.strip():
        log.info("resummarize: doc %s has no text — nothing to summarize", document_id)
        return "no-text"

    await _summarize_and_store(text=text, document_id=document_id)
    return "resummarized"


async def reindex_graph_for_document(ctx: dict[str, Any], document_id: str) -> str:
    # Re-run entity extraction on stored chunks. No re-OCR or re-embed.
    async with async_session_maker() as db:
        doc = await db.get(Document, document_id)
        if doc is None:
            log.warning("reindex_graph: doc %s not found", document_id)
            return "missing"

        if str(doc.status) != DocumentStatus.PROCESSED.value:
            log.info(
                "reindex_graph: doc %s status=%s — skipping (expected 'processed')",
                document_id,
                doc.status,
            )
            return "skipped"

        stmt = (
            select(DocumentChunk)
            .where(DocumentChunk.document_id == doc.id)
            .order_by(DocumentChunk.chunk_index.asc())
        )
        rows = (await db.execute(stmt)).scalars().all()

    # Prefer stored chunk rows; fall back to extracted_text as a single
    # chunk for legacy docs indexed before chunking existed.
    if rows:
        chunk_texts = [r.text for r in rows]
    elif doc.extracted_text and doc.extracted_text.strip():
        chunk_texts = [doc.extracted_text]
    else:
        chunk_texts = []

    if not any(t.strip() for t in chunk_texts):
        log.info("reindex_graph: doc %s has no text — nothing to extract", document_id)
        return "no-text"

    await _ingest_graph(
        chunks=chunk_texts,
        user_id=str(doc.user_id),
        document_id=str(doc.id),
    )
    return "reindexed"


async def fetch_document(document_id: str) -> Document | None:
    # Test helper; not part of arq's surface.
    async with async_session_maker() as db:
        result = await db.execute(select(Document).where(Document.id == document_id))
        return result.scalar_one_or_none()