0.1.0

wolai-backend/app/workers/celery_app.py

@@ -1,15 +1,53 @@
+import os
+
 from celery import Celery
+from kombu import Queue
 
 from app.config import settings
 
 
+def _env_flag(name: str, default: bool = False) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
 def create_celery_app() -> Celery:
     app = Celery("wolai-backend")
     app.conf.broker_url = settings.redis_url
     app.conf.result_backend = settings.redis_url
-    app.conf.task_routes = {"app.workers.tasks.*": {"queue": "wolai-tasks"}}
+    app.conf.task_default_queue = "wolai-index"
+    app.conf.task_queues = (
+        Queue("wolai-ocr", routing_key="tasks.ocr", max_priority=10),
+        Queue("wolai-index", routing_key="tasks.index", max_priority=10),
+    )
+    app.conf.task_routes = {
+        "app.workers.tasks.ocr_pipeline": {
+            "queue": "wolai-ocr",
+            "routing_key": "tasks.ocr",
+            "priority": 0,
+        },
+        "app.workers.tasks.media_ocr_pipeline": {
+            "queue": "wolai-ocr",
+            "routing_key": "tasks.ocr",
+            "priority": 1,
+        },
+        "app.workers.tasks.lightrag_index_pipeline": {
+            "queue": "wolai-index",
+            "routing_key": "tasks.index",
+            "priority": 5,
+        },
+    }
+    app.conf.broker_transport_options = {"priority_steps": list(range(10))}
+    app.conf.worker_prefetch_multiplier = 1
+    app.conf.task_acks_late = True
+    app.conf.worker_concurrency = 3
+    app.conf.task_always_eager = _env_flag("CELERY_TASK_ALWAYS_EAGER", False)
+    app.conf.task_eager_propagates = _env_flag("CELERY_TASK_EAGER_PROPAGATES", True)
     app.autodiscover_tasks(["app.workers"])
     return app
 
 
 celery_app = create_celery_app()
+celery_app.set_default()

wolai-backend/app/workers/tasks.py

@@ -1,45 +1,68 @@
 from __future__ import annotations
 
-from typing import Dict
+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 任务：模拟 OCR，向 documents+background_tasks 回写占位结果。
-    阶段 1 将在此处串联 MinerU OCR 与 LightRAG 索引。
+    阶段 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 完成"
 
-    markdown = f"# OCR 结果占位\\n\\n文件地址：{file_url}\\n\\n> 阶段 1 将替换为 MinerU 输出。"
-    supabase_rest.update(
-        "documents",
-        {"id": document_id, "user_id": user_id},
-        {
-            "content": {
-                "blocks": [
-                    {
-                        "type": "paragraph",
-                        "text": markdown,
-                    }
-                ]
+    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",
             },
-            "raw_text": markdown,
-            "index_status": "completed",
-        },
-    )
+        )
+        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": "completed",
+            "status": status,
             "progress": 100,
-            "message": "OCR 模拟完成",
+            "message": message,
         },
     )
 
@@ -47,5 +70,132 @@ def ocr_pipeline(task_id: str, document_id: str, file_url: str, user_id: str) ->
         "task_id": task_id,
         "document_id": document_id,
         "file_url": file_url,
-        "status": "completed",
+        "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)}

wolai-backend/app/workers/utils.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Callable
+
+logger = logging.getLogger(__name__)
+
+
+def _env_flag(name: str) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return False
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def dispatch_task(task: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+    """
+    根据环境变量决定直接执行任务还是通过 Celery delay。
+    本地开发默认走 eager，同步执行便于调试。
+    """
+    eager = _env_flag("CELERY_TASK_ALWAYS_EAGER")
+    logger.info("dispatch_task task=%s eager=%s", getattr(task, "name", task.__name__), eager)
+    if _env_flag("CELERY_TASK_ALWAYS_EAGER"):
+        return task.apply(args=args, kwargs=kwargs)
+    return task.delay(*args, **kwargs)