mineru-rocm/docker/scripts/mineru_webui.py

#!/usr/bin/env python3
"""MinerU Gradio 客户端 — 纯前端，所有解析任务通过 HTTP 转发到 Router/Worker。

与 mineru-gradio 的区别：
- 不启动内嵌 FastAPI server
- 不需要 GPU / vLLM
- 通过 HTTP 调用 mineru-router → mineru-worker 处理任务
"""
import os
import time
import asyncio
import tempfile
import zipfile
import json
from pathlib import Path

import gradio as gr
import httpx

API_BASE = os.environ.get("MINERU_API_BASE", "http://mineru-nginx:8000")
DEFAULT_BACKEND = os.environ.get("MINERU_BACKEND", "hybrid-engine")
DEFAULT_LANG = os.environ.get("MINERU_LANG", "ch")
POLL_INTERVAL = float(os.environ.get("MINERU_POLL_INTERVAL", "1.0"))
MAX_WAIT = float(os.environ.get("MINERU_MAX_WAIT", "600"))

CUSTOM_CSS = """
.result-panel { min-height: 400px; }
.status-panel { min-height: 200px; font-size: 13px; }
"""


async def discover_api(client: httpx.AsyncClient) -> str:
    """探测实际可用的 API 前缀: /file_parse 端点"""
    candidates = [
        f"{API_BASE}/openapi.json",
        f"{API_BASE}/docs",
    ]
    for url in candidates:
        try:
            r = await client.get(url, timeout=5)
            if r.status_code == 200:
                return ""  # worker uses root-level endpoints like /file_parse
        except Exception:
            continue
    return ""


async def submit_task(client: httpx.AsyncClient, api_prefix: str,
                      file_path: str, file_name: str,
                      backend: str, lang: str) -> dict:
    """提交解析任务到 worker 的 /file_parse 端点"""
    submit_url = f"{API_BASE}/file_parse"
    with open(file_path, "rb") as f:
        files = {"files": (file_name, f, "application/pdf")}
        data = {
            "backend": backend,
            "lang_list": lang,
            "parse_method": "auto",
        }
        r = await client.post(submit_url, files=files, data=data, timeout=30)
        r.raise_for_status()
        return r.json()


async def get_task_status(client: httpx.AsyncClient, api_prefix: str,
                          task_id: str) -> dict:
    """查询任务状态"""
    url = f"{API_BASE}/tasks/{task_id}"
    r = await client.get(url, timeout=10)
    r.raise_for_status()
    return r.json()


async def download_result(client: httpx.AsyncClient, api_prefix: str,
                          task_id: str, output_dir: str) -> str:
    """下载任务结果"""
    url = f"{API_BASE}/tasks/{task_id}/result"
    r = await client.get(url, timeout=60, follow_redirects=True)
    r.raise_for_status()
    zip_path = os.path.join(output_dir, f"{task_id}.zip")
    with open(zip_path, "wb") as f:
        f.write(r.content)
    return zip_path


def extract_readme(zip_path: str, output_dir: str) -> str:
    """解压结果并返回 markdown 内容"""
    extract_dir = os.path.join(output_dir, "extracted")
    os.makedirs(extract_dir, exist_ok=True)
    with zipfile.ZipFile(zip_path, "r") as zf:
        zf.extractall(extract_dir)
    # 查找 markdown 文件
    md_files = list(Path(extract_dir).rglob("*.md"))
    if md_files:
        return md_files[0].read_text(encoding="utf-8")
    return "No markdown output found in result."


# ---------------------------------------------------------------------------
# Gradio 处理函数
# ---------------------------------------------------------------------------
async def process_pdf(file_obj, backend, lang, progress=gr.Progress()):
    """Gradio 事件处理：上传 → 提交 → 轮询 → 返回结果"""
    if file_obj is None:
        return "Please upload a PDF file.", "", ""

    tmp_dir = tempfile.mkdtemp(prefix="mineru-client-")

    # Gradio 6.0 binary mode：file_obj 是 bytes
    if isinstance(file_obj, bytes):
        file_path = os.path.join(tmp_dir, "input.pdf")
        file_name = "input.pdf"
        with open(file_path, "wb") as f:
            f.write(file_obj)
    elif isinstance(file_obj, str):
        file_path = file_obj
        file_name = os.path.basename(file_path)
    else:
        return "Unsupported file object type", "", ""

    async with httpx.AsyncClient(timeout=httpx.Timeout(120)) as client:
        # 1. 探测 API
        progress(0.05, desc="Connecting to API...")
        api_prefix = await discover_api(client)
        progress(0.1, desc=f"API prefix: {api_prefix}")

        # 2. 提交任务
        progress(0.15, desc="Submitting task...")
        try:
            submit_resp = await submit_task(client, api_prefix, file_path, file_name, backend, lang)
        except Exception as e:
            return f"Task submission failed: {e}", "", ""

        task_id = submit_resp.get("task_id") or submit_resp.get("data", {}).get("task_id")
        if not task_id:
            return f"Unexpected submit response:\n{json.dumps(submit_resp, indent=2)}", "", ""

        progress(0.2, desc=f"Task ID: {task_id}")

        # 3. 轮询状态
        status_text = ""
        start_time = time.time()
        last_idx = -1
        tickers = ["⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"]

        while True:
            elapsed = time.time() - start_time
            if elapsed > MAX_WAIT:
                return f"Task timed out after {MAX_WAIT}s\nLast status:\n{status_text}", "", ""

            t = (int(elapsed * 2)) % len(tickers)
            progress(min(0.2 + 0.6 * (elapsed / 60), 0.8),
                     desc=f"{tickers[t]} Processing... ({int(elapsed)}s)")

            try:
                status_resp = await get_task_status(client, api_prefix, task_id)
            except Exception as e:
                status_text = f"Polling error: {e}"
                await asyncio.sleep(POLL_INTERVAL)
                continue

            status = status_resp.get("status", "unknown")
            progress_pct = status_resp.get("progress", 0)
            err_msg = status_resp.get("error", "")
            status_text = json.dumps(status_resp, indent=2, ensure_ascii=False)

            if status == "completed" or status == "success":
                progress(0.85, desc="Downloading result...")
                break
            elif status == "failed":
                return f"## Task Failed\n\nError: {err_msg}\n\n```json\n{status_text}\n```", "", ""
            elif status == "processing":
                pass
            else:
                pass

            await asyncio.sleep(POLL_INTERVAL)

        # 4. 下载结果
        progress(0.9, desc="Downloading result...")
        try:
            zip_path = await download_result(client, api_prefix, task_id, tmp_dir)
        except Exception as e:
            return f"Result download failed: {e}\n\nStatus:\n{status_text}", "", ""

        # 5. 提取 markdown
        progress(0.95, desc="Extracting results...")
        try:
            md_content = extract_readme(zip_path, tmp_dir)
        except Exception as e:
            md_content = f"Extraction error: {e}"

        progress(1.0, desc="Done!")
        return md_content, status_text, zip_path


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
def create_ui():
    with gr.Blocks(title="MinerU ROCm — Document Parser") as demo:
        gr.Markdown("""
        # MinerU ROCm — Document Parser

        Upload a PDF file. Processing is handled by the **mineru-router → mineru-worker** backend.
        """)

        with gr.Row():
            with gr.Column(scale=1):
                file_input = gr.UploadButton(
                    label="Upload PDF",
                    file_types=[".pdf"],
                    type="binary",
                )
                backend_dd = gr.Dropdown(
                    label="Backend",
                    choices=["hybrid-engine", "hybrid-http-client",
                             "vlm-http-client", "pipeline", "auto"],
                    value=DEFAULT_BACKEND,
                )
                lang_dd = gr.Dropdown(
                    label="Language",
                    choices=["ch", "en", "japan", "korean"],
                    value=DEFAULT_LANG,
                )
                submit_btn = gr.Button("Parse Document", variant="primary", size="lg")
                status_display = gr.Textbox(
                    label="Status",
                    lines=6,
                    max_lines=12,
                    elem_classes=["status-panel"],
                )
                zip_display = gr.File(
                    label="Download Result ZIP",
                    type="filepath",
                    visible=True,
                )

            with gr.Column(scale=2):
                result_display = gr.Markdown(
                    value="*Upload a PDF to start parsing...*",
                    elem_classes=["result-panel"],
                )

        file_input.upload(
            fn=process_pdf,
            inputs=[file_input, backend_dd, lang_dd],
            outputs=[result_display, status_display, zip_display],
        )
        submit_btn.click(
            fn=lambda f,b,l: "Please use the Upload button above.",
            inputs=[file_input, backend_dd, lang_dd],
            outputs=[result_display, status_display, zip_display],
        )

    return demo


if __name__ == "__main__":
    server_name = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0")
    server_port = int(os.environ.get("GRADIO_SERVER_PORT", "7860"))
    demo = create_ui()
    demo.queue(default_concurrency_limit=3, max_size=10).launch(
        server_name=server_name,
        server_port=server_port,
        share=False,
        css=CUSTOM_CSS,
    )