81 lines
3.2 KiB
Python
81 lines
3.2 KiB
Python
"""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：直接返回固定内容，保证前端流程贯通。
|
|
阶段 1：调用 MinerU CLI / SDK，从 Supabase Storage 下载文件后解析。
|
|
"""
|
|
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()
|