mnote/wolai-backend/app/workers/tasks.py

from __future__ import annotations

import asyncio
import logging
from typing import Dict, Optional

from celery import shared_task

from app.services.lightrag_service import lightrag_service
from app.services.mineru_service import mineru_service
from app.services.storage_fetcher import storage_fetcher
from app.services.supabase_rest import supabase_rest
from app.services.task_tracker import task_tracker
from app.workers.utils import dispatch_task

logger = logging.getLogger(__name__)


@shared_task(name="app.workers.tasks.ocr_pipeline")
def ocr_pipeline(task_id: str, document_id: str, file_url: str, user_id: str) -> Dict[str, str]:
    """
    阶段 0 Celery 任务：调用 MinerU 占位实现并写回 documents/background_tasks。
    阶段 1 会在此处替换为真实的 MinerU CLI / LightRAG 流程。
    """
    supabase_rest.update("background_tasks", {"id": task_id}, {"status": "processing", "progress": 30})
    status = "completed"
    message = "OCR 完成"

    file_path = None
    try:
        file_path = storage_fetcher.download(file_url)
        markdown = mineru_service.extract_markdown_sync(file_path)
        supabase_rest.update(
            "documents",
            {"id": document_id, "user_id": user_id},
            {
                "content": {
                    "blocks": [
                        {
                            "type": "paragraph",
                            "text": markdown,
                        }
                    ]
                },
                "raw_text": markdown,
                "index_status": "pending",
            },
        )
        supabase_rest.update("background_tasks", {"id": task_id}, {"progress": 80})
        index_task = task_tracker.create_task(
            user_id=user_id, document_id=document_id, task_type="index"
        )
        dispatch_task(lightrag_index_pipeline, index_task.task_id, document_id, user_id)
    except Exception as exc:  # pragma: no cover - 网络/IO异常
        status = "failed"
        message = f"OCR 失败：{exc}"
    finally:
        storage_fetcher.cleanup(file_path)

    supabase_rest.update(
        "background_tasks",
        {"id": task_id},
        {
            "status": status,
            "progress": 100,
            "message": message,
        },
    )

    return {
        "task_id": task_id,
        "document_id": document_id,
        "file_url": file_url,
        "status": status,
    }


@shared_task(name="app.workers.tasks.media_ocr_pipeline")
def media_ocr_pipeline(asset_id: str, user_id: str, document_id: Optional[str] = None) -> Dict[str, str]:
    """媒体 OCR：下载资源 → 调用 MinerU → 写回 ocr_text/ocr_payload。"""
    logger.info("媒体 OCR 任务执行 asset_id=%s user_id=%s", asset_id, user_id)
    asset = supabase_rest.select_one("media_assets", {"id": asset_id})
    if not asset:
        logger.warning("媒体 OCR 任务失败：资源不存在 %s", asset_id)
        return {"asset_id": asset_id, "status": "failed"}
    file_url = asset.get("file_url")
    ocr_status = "completed"
    ocr_text: Optional[str] = None
    message = "媒体 OCR 完成"
    file_path = None
    try:
        file_path = storage_fetcher.download(str(file_url))
        ocr_text = mineru_service.extract_markdown_sync(file_path)
    except Exception as exc:  # pragma: no cover - 网络/IO异常
        logger.warning("媒体 OCR 失败 asset=%s err=%s", asset_id, exc)
        ocr_status = "failed"
        message = f"OCR 失败：{exc}"
    finally:
        storage_fetcher.cleanup(file_path)
        supabase_rest.update(
            "media_assets",
            {"id": asset_id},
            {
                "ocr_status": ocr_status,
                "ocr_text": ocr_text,
                "ocr_payload": {"pages": [], "summary": ocr_text} if ocr_text else None,
            },
        )

    return {
        "asset_id": asset_id,
        "document_id": document_id,
        "user_id": user_id,
        "status": ocr_status,
        "message": message,
    }


@shared_task(name="app.workers.tasks.lightrag_index_pipeline")
def lightrag_index_pipeline(task_id: str, document_id: str, user_id: str) -> Dict[str, str]:
    """执行 LightRAG 索引构建，独立于 OCR。"""
    supabase_rest.update(
        "background_tasks",
        {"id": task_id},
        {"status": "processing", "progress": 10, "message": "LightRAG 索引排队中"},
    )
    document = supabase_rest.select_one("documents", {"id": document_id, "user_id": user_id})
    if not document:
        supabase_rest.update(
            "background_tasks",
            {"id": task_id},
            {"status": "failed", "progress": 100, "message": "文档不存在"},
        )
        return {"status": "failed", "document_id": document_id}

    # 先做健康检查，避免长时间阻塞或无谓重试
    try:
        health = asyncio.run(lightrag_service.run_healthcheck())
    except Exception as exc:  # pragma: no cover - 调试辅助
        health = {"ok": False, "error": str(exc)}
    if not health.get("ok"):
        reason = health.get("error", "LightRAG 未就绪")
        supabase_rest.update(
            "background_tasks",
            {"id": task_id},
            {"status": "failed", "progress": 100, "message": reason},
        )
        supabase_rest.update(
            "documents", {"id": document_id}, {"index_status": "failed"}
        )
        return {"status": "failed", "document_id": document_id, "error": reason}

    raw_text = document.get("raw_text") or ""
    text_to_index = str(raw_text).strip()
    if not text_to_index:
        supabase_rest.update(
            "background_tasks",
            {"id": task_id},
            {"status": "failed", "progress": 100, "message": "raw_text 为空，无法索引"},
        )
        return {"status": "failed", "document_id": document_id}
    workspace_id = document.get("workspace_id")
    title = document.get("title")
    supabase_rest.update(
        "documents", {"id": document_id}, {"index_status": "processing"}
    )
    try:
        supabase_rest.update(
            "background_tasks",
            {"id": task_id},
            {"progress": 40, "message": "LightRAG 建立中"},
        )
        lightrag_service.index_document(
            document_id=document_id,
            user_id=user_id,
            workspace_id=str(workspace_id) if workspace_id else None,
            text=text_to_index,
            title=title if isinstance(title, str) else None,
        )
        supabase_rest.update(
            "documents", {"id": document_id}, {"index_status": "completed"}
        )
        supabase_rest.update(
            "background_tasks",
            {"id": task_id},
            {"progress": 100, "status": "completed", "message": "LightRAG 完成"},
        )
        return {"status": "completed", "document_id": document_id}
    except Exception as exc:  # pragma: no cover - 调试错误
        supabase_rest.update(
            "documents", {"id": document_id}, {"index_status": "failed"}
        )
        supabase_rest.update(
            "background_tasks",
            {"id": task_id},
            {
                "progress": 100,
                "status": "failed",
                "message": f"LightRAG 失败：{exc}",
            },
        )
        return {"status": "failed", "document_id": document_id, "error": str(exc)}