app/tasks/ingestion_task.py

"""Author RAG Chatbot SaaS — Document Ingestion Celery Task.

This is the heart of the document pipeline.
Triggered when a document is uploaded and validated.
Pipeline: Parse → Chunk → Embed → Summarize → Mark Ready

RULE: Each stage updates document and book status in the DB.
RULE: SSE updates are broadcast via Redis pub/sub for real-time frontend updates.
RULE: Any failure at any stage marks the document as 'error' — never leaves it stuck.
"""

import asyncio
import os

import structlog

from app.tasks.celery_app import celery_app

logger = structlog.get_logger(__name__)


@celery_app.task(
    bind=True,
    max_retries=3,
    default_retry_delay=30,
    acks_late=True,
    name="app.tasks.ingestion_task.process_document",
)
def process_document(
    self,
    document_id: str,
    author_id: str,
    book_id: str,
    file_path: str,
    file_extension: str,
) -> dict:
    """Process a document through the full ingestion pipeline.

    Runs synchronously within Celery worker — uses asyncio.run() for
    async operations (embedder, summarizer).

    Args:
        document_id: UUID of the Document record.
        author_id: UUID of the author (tenant context).
        book_id: UUID of the associated book.
        file_path: Absolute path to the uploaded file.
        file_extension: Validated extension ('pdf', 'epub', 'docx', 'txt').

    Returns:
        Dict with ingestion results: chunk_count, page_count, collection_name.
    """
    logger.info("Starting document ingestion", doc_id=document_id, extension=file_extension)

    try:
        return asyncio.run(_run_pipeline(
            document_id=document_id,
            author_id=author_id,
            book_id=book_id,
            file_path=file_path,
            file_extension=file_extension,
        ))
    except Exception as exc:
        logger.error("Ingestion pipeline failed", doc_id=document_id, error=str(exc))
        # Update status to error in DB
        asyncio.run(_mark_error(document_id, author_id, book_id, str(exc)))
        raise self.retry(exc=exc)


async def _run_pipeline(
    document_id: str,
    author_id: str,
    book_id: str,
    file_path: str,
    file_extension: str,
) -> dict:
    """Execute ingestion via shared ingestion_service (R-150)."""
    from app.dependencies import _get_session_factory, get_redis
    from app.repositories.book_repo import BookRepository
    from app.repositories.document_repo import DocumentRepository
    from app.services.ingestion_service import ingest_from_file

    redis = await get_redis()
    async with _get_session_factory()() as db:
        book_repo = BookRepository(db)
        doc_repo = DocumentRepository(db)
        book = await book_repo.get_by_id_for_author(book_id, author_id)
        book_title = book.title if book else "Unknown Book"

        await doc_repo.update_status(document_id, author_id, "processing")
        await db.commit()

        await ingest_from_file(
            file_path=file_path,
            book_id=book_id,
            author_id=author_id,
            title=book_title,
            redis=redis,
            db=db,
        )

        await doc_repo.update_status(document_id, author_id, "ready")
        await db.commit()

        refreshed = await book_repo.get_by_id_for_author(book_id, author_id)
        return {
            "chunk_count": getattr(refreshed, "chunk_count", 0) or 0,
            "page_count": getattr(refreshed, "page_count", 0) or 0,
            "collection_name": getattr(refreshed, "chroma_collection_id", "") or "",
        }


async def _mark_error(document_id: str, author_id: str, book_id: str, error: str) -> None:
    """Mark document and book as errored in the database.

    Args:
        document_id: UUID of the document.
        author_id: UUID of the author.
        book_id: UUID of the book.
        error: Error message to store.
    """
    from app.dependencies import _get_session_factory
    from app.repositories.document_repo import DocumentRepository
    from app.repositories.book_repo import BookRepository

    async with _get_session_factory()() as db:
        await DocumentRepository(db).update_status(document_id, author_id, "error", error_message=error)
        await BookRepository(db).update_status(book_id, author_id, "error", error=error)
        await db.commit()
        await _publish_status(author_id, book_id, "error", 0, error=error)


async def _publish_status(
    author_id: str,
    book_id: str,
    stage: str,
    progress: int = 0,
    error: str | None = None,
) -> None:
    """Publish ingestion progress to Redis SSE channel (R-084/R-150)."""
    import json
    from app.dependencies import get_redis

    try:
        redis = await get_redis()
        channel = f"ingestion:{author_id}:{book_id}"
        payload = {"stage": stage, "progress": progress, "detail": error or ""}
        await redis.publish(channel, json.dumps(payload))
    except Exception as e:
        logger.warning("Failed to publish status update", error=str(e))