mnemosyne/mnemosyne/library/tasks.py

"""
Celery tasks for the embedding pipeline.

All tasks pass UIDs (not model instances) per Red Panda Standards.
Tasks are idempotent, include retry logic, and track progress
via Memcached: library:task:{task_id}:progress.
"""

import logging

from celery import shared_task
from django.core.cache import cache
from neomodel import db

logger = logging.getLogger(__name__)

# Cache key pattern for task progress
PROGRESS_KEY = "library:task:{task_id}:progress"

# MIME type → file extension, for when Daedalus sends content_type as file_type
_MIME_TO_EXT = {
    "text/markdown": "md",
    "text/plain": "txt",
    "text/html": "html",
    "text/csv": "csv",
    "text/xml": "xml",
    "application/pdf": "pdf",
    "application/epub+zip": "epub",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
    "application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx",
    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx",
    "application/json": "json",
    "image/jpeg": "jpg",
    "image/png": "png",
    "image/gif": "gif",
    "image/webp": "webp",
    "image/tiff": "tiff",
}


def _normalize_file_type(raw: str) -> str:
    """Convert a MIME type or extension string to a bare extension."""
    raw = (raw or "").strip().lower()
    if "/" in raw:
        # It's a MIME type — look up or derive from the subtype
        ext = _MIME_TO_EXT.get(raw)
        if ext:
            return ext
        # Fallback: use the part after the slash, strip vendor prefixes
        subtype = raw.split("/", 1)[1]
        subtype = subtype.split("+")[-1]  # e.g. "epub+zip" → "zip"; "vnd.ms-excel" → keep
        return subtype.lstrip(".") or "bin"
    return raw.lstrip(".") or "bin"


def _update_progress(task, percent: int, message: str):
    """
    Update task progress in Memcached and Celery state.

    :param task: Celery task instance (self).
    :param percent: Progress percentage (0-100).
    :param message: Human-readable status message.
    """
    try:
        task.update_state(state="PROGRESS", meta={"percent": percent, "message": message})
        cache.set(
            PROGRESS_KEY.format(task_id=task.request.id),
            {"percent": percent, "message": message},
            timeout=3600,
        )
    except Exception:
        pass


@shared_task(
    name="library.tasks.embed_item",
    bind=True,
    queue="embedding",
    max_retries=3,
    default_retry_delay=60,
    acks_late=True,
)
def embed_item(self, item_uid: str, user_id: int = None):
    """
    Run the full embedding pipeline for a single Item.

    :param item_uid: UID of the Item node to process.
    :param user_id: Optional user ID for usage tracking.
    :returns: Dict with processing results.
    """
    logger.info("Task embed_item starting item_uid=%s task_id=%s", item_uid, self.request.id)

    try:
        from library.services.pipeline import EmbeddingPipeline

        user = _resolve_user(user_id)
        pipeline = EmbeddingPipeline(user=user)

        def progress_cb(percent, message):
            _update_progress(self, percent, message)

        result = pipeline.process_item(item_uid, progress_callback=progress_cb)

        logger.info(
            "Task embed_item completed item_uid=%s chunks=%d images=%d",
            item_uid,
            result.get("chunks_created", 0),
            result.get("images_stored", 0),
        )
        return {"success": True, "item_uid": item_uid, **result}

    except Exception as exc:
        logger.error(
            "Task embed_item failed item_uid=%s: %s",
            item_uid,
            exc,
            exc_info=True,
        )
        # Retry on transient errors
        if self.request.retries < self.max_retries:
            raise self.retry(exc=exc)
        return {"success": False, "item_uid": item_uid, "error": str(exc)}


@shared_task(
    name="library.tasks.reembed_item",
    bind=True,
    queue="embedding",
    max_retries=3,
    default_retry_delay=60,
    acks_late=True,
)
def reembed_item(self, item_uid: str, user_id: int = None):
    """
    Delete existing embeddings and re-process an Item.

    :param item_uid: UID of the Item node to re-embed.
    :param user_id: Optional user ID for usage tracking.
    :returns: Dict with processing results.
    """
    logger.info("Task reembed_item starting item_uid=%s", item_uid)

    try:
        from library.services.pipeline import EmbeddingPipeline

        user = _resolve_user(user_id)
        pipeline = EmbeddingPipeline(user=user)

        def progress_cb(percent, message):
            _update_progress(self, percent, message)

        result = pipeline.reprocess_item(item_uid, progress_callback=progress_cb)

        logger.info("Task reembed_item completed item_uid=%s", item_uid)
        return {"success": True, "item_uid": item_uid, **result}

    except Exception as exc:
        logger.error("Task reembed_item failed item_uid=%s: %s", item_uid, exc, exc_info=True)
        if self.request.retries < self.max_retries:
            raise self.retry(exc=exc)
        return {"success": False, "item_uid": item_uid, "error": str(exc)}


@shared_task(
    name="library.tasks.embed_collection",
    bind=True,
    queue="batch",
    acks_late=True,
)
def embed_collection(self, collection_uid: str, user_id: int = None):
    """
    Embed all items in a collection.

    :param collection_uid: UID of the Collection node.
    :param user_id: Optional user ID for usage tracking.
    :returns: Dict with summary results.
    """
    logger.info("Task embed_collection starting collection_uid=%s", collection_uid)

    try:
        from library.models import Collection

        col = Collection.nodes.get(uid=collection_uid)
        items = col.items.all()

        results = {"total": len(items), "successful": 0, "failed": 0, "skipped": 0}

        for i, item in enumerate(items):
            # Skip already-completed items with unchanged content
            if item.embedding_status == "completed" and item.content_hash:
                results["skipped"] += 1
                logger.debug("Skipping already-embedded item_uid=%s", item.uid)
                continue

            try:
                embed_item.delay(item.uid, user_id)
                results["successful"] += 1
            except Exception as exc:
                results["failed"] += 1
                logger.error(
                    "Failed to queue embed for item_uid=%s: %s", item.uid, exc
                )

            _update_progress(
                self,
                int((i + 1) / len(items) * 100),
                f"Queued {i + 1}/{len(items)} items",
            )

        logger.info(
            "Task embed_collection completed collection_uid=%s queued=%d skipped=%d failed=%d",
            collection_uid,
            results["successful"],
            results["skipped"],
            results["failed"],
        )
        return {"success": True, "collection_uid": collection_uid, **results}

    except Exception as exc:
        logger.error(
            "Task embed_collection failed collection_uid=%s: %s",
            collection_uid,
            exc,
            exc_info=True,
        )
        return {"success": False, "collection_uid": collection_uid, "error": str(exc)}


@shared_task(
    name="library.tasks.embed_library",
    bind=True,
    queue="batch",
    acks_late=True,
)
def embed_library(self, library_uid: str, user_id: int = None):
    """
    Embed all items across all collections in a library.

    :param library_uid: UID of the Library node.
    :param user_id: Optional user ID for usage tracking.
    :returns: Dict with summary results.
    """
    logger.info("Task embed_library starting library_uid=%s", library_uid)

    try:
        from library.models import Library

        lib = Library.nodes.get(uid=library_uid)
        collections = lib.collections.all()

        results = {"total_collections": len(collections), "items_queued": 0}

        for col in collections:
            embed_collection.delay(col.uid, user_id)
            results["items_queued"] += len(col.items.all())

        logger.info(
            "Task embed_library completed library_uid=%s collections=%d items=%d",
            library_uid,
            results["total_collections"],
            results["items_queued"],
        )
        return {"success": True, "library_uid": library_uid, **results}

    except Exception as exc:
        logger.error(
            "Task embed_library failed library_uid=%s: %s",
            library_uid,
            exc,
            exc_info=True,
        )
        return {"success": False, "library_uid": library_uid, "error": str(exc)}


@shared_task(
    name="library.tasks.batch_embed_items",
    bind=True,
    queue="batch",
    acks_late=True,
)
def batch_embed_items(self, item_uids: list[str], user_id: int = None):
    """
    Queue embedding tasks for a specific list of items.

    :param item_uids: List of Item UIDs.
    :param user_id: Optional user ID for usage tracking.
    :returns: Dict with queuing results.
    """
    logger.info("Task batch_embed_items starting count=%d", len(item_uids))

    queued = 0
    for uid in item_uids:
        try:
            embed_item.delay(uid, user_id)
            queued += 1
        except Exception as exc:
            logger.error("Failed to queue item_uid=%s: %s", uid, exc)

    logger.info("Task batch_embed_items completed queued=%d/%d", queued, len(item_uids))
    return {"success": True, "queued": queued, "total": len(item_uids)}


def _resolve_user(user_id: int = None):
    """
    Resolve a user ID to a User instance.

    :param user_id: Optional user ID.
    :returns: User instance, or None.
    """
    if not user_id:
        return None
    try:
        from django.contrib.auth import get_user_model

        User = get_user_model()
        return User.objects.get(pk=user_id)
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Ingest task (Daedalus integration)
# ---------------------------------------------------------------------------


@shared_task(
    name="library.tasks.ingest_from_daedalus",
    bind=True,
    queue="embedding",
    max_retries=3,
    default_retry_delay=60,
    acks_late=True,
)
def ingest_from_daedalus(self, job_id: str):
    """
    Process a single IngestJob: fetch from Daedalus S3 → create Item →
    run embedding pipeline → mark complete.

    Idempotent on (library_uid, source_ref, content_hash) — handled in the
    REST view that creates the IngestJob, so by the time this task runs the
    job either represents new content or a content_hash-changed re-ingest.

    For a content_hash-changed re-ingest, the prior Item with the same
    source_ref is deleted before the new one is processed (ensures no
    stale chunks linger).
    """
    from datetime import datetime, timezone

    from library.models import IngestJob, Item, Library
    from library.services.parsers import UnsupportedFileTypeError
    from library.services.source_s3 import (
        copy_into_mnemosyne,
        fetch_from_source,
    )
    from library.services.pipeline import EmbeddingPipeline

    logger.info(
        "Task ingest_from_daedalus starting job_id=%s task_id=%s",
        job_id, self.request.id,
    )

    try:
        job = IngestJob.objects.get(pk=job_id)
    except IngestJob.DoesNotExist:
        logger.error("IngestJob not found job_id=%s", job_id)
        return {"success": False, "error": "job_not_found"}

    job.status = "processing"
    job.progress = "fetching"
    job.started_at = datetime.now(timezone.utc)
    job.celery_task_id = self.request.id
    job.save(update_fields=["status", "progress", "started_at", "celery_task_id"])

    try:
        # --- 1. Resolve target Library ---
        try:
            lib = Library.nodes.get(uid=job.library_uid)
        except Library.DoesNotExist:
            raise RuntimeError(f"Library not found: {job.library_uid}")

        # --- 2. Supersede prior Item with same source_ref but different hash ---
        prior_item_uid = None
        if job.source_ref:
            rows, _ = db.cypher_query(
                """
                MATCH (l:Library {uid: $library_uid})-[:CONTAINS]->(:Collection)
                -[:CONTAINS]->(i:Item)
                WHERE i.metadata IS NOT NULL
                  AND i.metadata CONTAINS $source_ref_marker
                RETURN i.uid LIMIT 1
                """,
                {
                    "library_uid": lib.uid,
                    "source_ref_marker": f'"source_ref": "{job.source_ref}"',
                },
            )
            if rows:
                prior_item_uid = rows[0][0]
                logger.info(
                    "Superseding prior Item job_id=%s prior_item_uid=%s",
                    job_id, prior_item_uid,
                )
                _delete_item_and_chunks(prior_item_uid)

        # --- 3. Fetch from the source bucket, copy into Mnemosyne bucket ---
        job.progress = "copying"
        job.save(update_fields=["progress"])

        data = fetch_from_source(job.source, job.s3_key)

        # --- 4. Create Item node ---
        ext = _normalize_file_type(job.file_type)
        item = Item(
            title=job.title,
            file_type=ext,
            file_size=len(data),
            content_hash=job.content_hash,
            embedding_status="pending",
            metadata={
                "source": job.source,
                "source_ref": job.source_ref,
            },
        )
        item.save()

        mnemosyne_s3_key = f"items/{item.uid}/original.{ext}"
        copy_into_mnemosyne(data, mnemosyne_s3_key)
        item.s3_key = mnemosyne_s3_key
        item.save()

        # --- 5. Connect to library/collection ---
        col = _resolve_or_create_default_collection(lib, job.collection_uid)
        col.items.connect(item)

        job.item_uid = item.uid
        job.save(update_fields=["item_uid"])

        # --- 6. Run the embedding pipeline ---
        job.progress = "embedding"
        job.save(update_fields=["progress"])

        def progress_cb(percent, message):
            _update_progress(self, percent, message)

        pipeline = EmbeddingPipeline(user=None)
        result = pipeline.process_item(item.uid, progress_callback=progress_cb)

        # --- 7. Mark complete ---
        job.status = "completed"
        job.progress = "done"
        job.chunks_created = result.get("chunks_created", 0)
        job.concepts_extracted = result.get("concepts_extracted", 0)
        job.embedding_model = result.get("embedding_model", "")
        job.completed_at = datetime.now(timezone.utc)
        job.save()

        logger.info(
            "Task ingest_from_daedalus completed job_id=%s item_uid=%s "
            "chunks=%d concepts=%d",
            job_id, item.uid, job.chunks_created, job.concepts_extracted,
        )
        return {
            "success": True,
            "job_id": job_id,
            "item_uid": item.uid,
            **result,
        }

    except UnsupportedFileTypeError as exc:
        # Deterministic, input-driven — re-parsing identical bytes can never
        # succeed. Terminal client-data failure, never retried, and logged at
        # WARNING (not ERROR) because an unparseable input is not a server fault.
        logger.warning(
            "Task ingest_from_daedalus rejected job_id=%s: %s", job_id, exc,
        )
        job.status = "failed"
        job.error = str(exc)
        job.completed_at = datetime.now(timezone.utc)
        job.save(update_fields=["status", "error", "completed_at"])
        return {
            "success": False,
            "job_id": job_id,
            "error": str(exc),
            "reason": "unsupported_file_type",
        }

    except Exception as exc:
        logger.error(
            "Task ingest_from_daedalus failed job_id=%s: %s",
            job_id, exc, exc_info=True,
        )
        if self.request.retries < self.max_retries:
            job.retry_count = self.request.retries + 1
            job.save(update_fields=["retry_count"])
            raise self.retry(exc=exc)

        job.status = "failed"
        job.error = str(exc)
        job.completed_at = datetime.now(timezone.utc)
        job.save(update_fields=["status", "error", "completed_at"])
        return {"success": False, "job_id": job_id, "error": str(exc)}


def _delete_item_and_chunks(item_uid: str):
    """Delete an Item, its chunks, and its images. Concept GC is workspace-delete only."""
    db.cypher_query(
        """
        MATCH (i:Item {uid: $uid})
        OPTIONAL MATCH (i)-[:HAS_CHUNK]->(c:Chunk)
        OPTIONAL MATCH (i)-[:HAS_IMAGE]->(img:Image)
        OPTIONAL MATCH (img)-[:HAS_EMBEDDING]->(emb:ImageEmbedding)
        DETACH DELETE c, img, emb, i
        """,
        {"uid": item_uid},
    )


def _resolve_or_create_default_collection(lib, collection_uid: str = ""):
    """
    Find or create the default Collection for a Library.

    Daedalus integration creates one Collection per Library, named "default".
    Explicit collection_uid is honored if provided.
    """
    from library.models import Collection

    if collection_uid:
        try:
            return Collection.nodes.get(uid=collection_uid)
        except Collection.DoesNotExist:
            pass

    # Look for an existing "default" collection in this library
    rows, _ = db.cypher_query(
        "MATCH (l:Library {uid: $library_uid})-[:CONTAINS]->(c:Collection {name: 'default'}) "
        "RETURN c.uid LIMIT 1",
        {"library_uid": lib.uid},
    )
    if rows:
        return Collection.nodes.get(uid=rows[0][0])

    col = Collection(name="default", description="Default collection")
    col.save()
    lib.collections.connect(col)
    col.library.connect(lib)
    return col