0.1.0

wolai-backend/app/services/mineru_service.py

@@ -1,7 +1,60 @@
-"""MinerU OCR 占位服务。"""
+"""MinerU OCR 服务封装，支持调用本地 MinerU HTTP 接口。"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional, Tuple
+
+import httpx
+
+
+class MinerUServiceError(RuntimeError):
+    """MinerU 调用失败时抛出的异常。"""
 
 
 class MinerUService:
+    def __init__(self) -> None:
+        self.endpoint = os.getenv("MINERU_ENDPOINT", "http://127.0.0.1:18888").rstrip("/")
+        timeout = float(os.getenv("MINERU_TIMEOUT_SECONDS", "300"))
+        # MinerU 首次加载模型可能耗时较长，这里提高超时时间避免大文件 OCR 直接失败
+        self._client = httpx.Client(timeout=timeout)
+
+    def _request_markdown(self, file_path: str) -> Tuple[Optional[str], dict]:
+        """
+        调用 MinerU /file_parse，返回 (markdown, raw_response)。
+        若接口未返回内容则返回 (None, response_dict)。
+        """
+        if not self.endpoint:
+            raise MinerUServiceError("未配置 MINERU_ENDPOINT，无法调用 MinerU")
+
+        url = f"{self.endpoint}/file_parse"
+        file_name = Path(file_path).name
+        backend = os.getenv("MINERU_DEFAULT_BACKEND", "vlm-transformers")
+        model_path = os.getenv("MINERU_MODEL_PATH")
+        with open(file_path, "rb") as fp:
+            response = self._client.post(
+                url,
+                files={"files": (file_name, fp, "application/octet-stream")},
+                data={
+                    "return_md": "true",
+                    "return_content_list": "false",
+                    "return_middle_json": "false",
+                    "response_format_zip": "false",
+                    "backend": backend,
+                    # fast_api main 会从环境注入 model_path；这里双保险随请求传递
+                    **({"model_path": model_path} if model_path else {}),
+                },
+            )
+        response.raise_for_status()
+        payload = response.json()
+        results = payload.get("results") or {}
+        if not isinstance(results, dict) or not results:
+            return None, payload
+        first_key = next(iter(results.keys()))
+        md_content = results.get(first_key, {}).get("md_content")
+        return md_content, payload
+
     async def extract_markdown(self, file_url: str) -> str:
         """
         阶段 0：直接返回固定内容，保证前端流程贯通。
@@ -9,5 +62,19 @@ class MinerUService:
         """
         return f"# OCR Placeholder\n\n源文件：{file_url}"
 
+    def extract_markdown_sync(self, file_path: str) -> str:
+        """
+        Celery 任务使用的同步封装。
+        - 若 MinerU 服务可用：调用 HTTP 接口返回 markdown
+        - 若失败：抛出 MinerUServiceError 让上层标记失败
+        """
+        try:
+            markdown, raw = self._request_markdown(file_path)
+            if markdown:
+                return markdown
+            raise MinerUServiceError(f"MinerU 未返回 md_content，响应片段：{str(raw)[:300]}")
+        except Exception as exc:  # pragma: no cover - IO/网络异常
+            raise MinerUServiceError(f"MinerU 调用失败：{exc}") from exc
+
 
 mineru_service = MinerUService()