PlantAI-rag-new/sdk/app/routers/extraction.py

"""routers/extraction.py —— 抽取任务代理（8000 主 + 7900 回落，API.md §10）。

- 上传仅限 PDF ≤50MB，否则 422 VALIDATION 信封。
- 双上游全不可达 → 503 UPSTREAM_DOWN 降级信封。
- 上游 5xx（如 MinerU 不可用）→ 502 UPSTREAM_ERROR。
错误信封与 mock_server 一致：直接返回 JSONResponse，信封在顶层。
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, AsyncIterator

from fastapi import APIRouter, File, Form, Query, UploadFile
from fastapi.responses import JSONResponse

from ..services import extraction as _extraction
from ..services.extraction import (
    MAX_PDF_BYTES,
    ExtractionUnavailable,
    ExtractionUpstreamError,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/extraction", tags=["extraction"])

_DOWN_MSG = "抽取服务当前不可用，请稍后重试"
_MINERU_MSG = "抽取任务提交失败：上游 MinerU 服务不可用，请检查抽取环境后重试"


def _envelope(code: str, message: str, degraded: bool, status: int) -> JSONResponse:
    return JSONResponse(
        status_code=status,
        content={
            "ok": False,
            "error": {"code": code, "message": message, "degraded": degraded},
        },
    )


def _down(exc: BaseException | None = None) -> JSONResponse:
    logger.warning("extraction: 上游不可用：%s", exc)
    return _envelope("UPSTREAM_DOWN", _DOWN_MSG, True, 503)


def _upstream_error(message: str) -> JSONResponse:
    logger.warning("extraction: 上游错误：%s", message)
    return _envelope("UPSTREAM_ERROR", message, True, 502)


def _validation(message: str) -> JSONResponse:
    return _envelope("VALIDATION", message, False, 422)


def _check_pdf_upload(filename: str, size: int) -> JSONResponse | None:
    """FR-EXTRACT-004：仅 PDF，单文件 ≤50MB。合法返回 None。"""
    if size <= 0:
        return _validation("上传文件为空")
    if size > MAX_PDF_BYTES:
        return _validation("文件超过 50MB 限制")
    if not filename.lower().endswith(".pdf"):
        return _validation("仅支持 PDF 文件")
    return None


async def _stream_upload(file: UploadFile) -> AsyncIterator[bytes]:
    try:
        while chunk := await file.read(1024 * 1024):  # 1MB
            yield chunk
    finally:
        await file.close()


@router.get("/jobs")
async def jobs(limit: int = Query(20, ge=1, le=200)) -> dict[str, Any]:
    try:
        return await _extraction.service.list_jobs(limit=limit)
    except ExtractionUnavailable as exc:
        return _down(exc)


@router.get("/jobs/{job_id}")
async def job_detail(job_id: str) -> dict[str, Any]:
    try:
        return await _extraction.service.get_job(job_id)
    except ExtractionUnavailable as exc:
        return _down(exc)


@router.get("/processing-status")
async def processing_status() -> dict[str, Any]:
    try:
        return await _extraction.service._passthrough("/processing-status")
    except ExtractionUnavailable as exc:
        return _down(exc)


@router.get("/progress")
async def progress() -> dict[str, Any]:
    try:
        return await _extraction.service._passthrough("/progress")
    except ExtractionUnavailable as exc:
        return _down(exc)


@router.get("/summary")
async def summary() -> dict[str, Any]:
    try:
        return await _extraction.service._passthrough("/summary")
    except ExtractionUnavailable as exc:
        return _down(exc)


@router.get("/run-log")
async def run_log() -> dict[str, Any]:
    try:
        return await _extraction.service._passthrough("/run-log")
    except ExtractionUnavailable as exc:
        return _down(exc)


@router.post("/upload")
async def upload(
    file: UploadFile = File(...),
    paper_id: str = Form("uploaded_paper"),
    family: str = Form(""),
    output_name: str = Form(""),
) -> dict[str, Any]:
    filename = file.filename or ""
    size = file.size or 0
    invalid = _check_pdf_upload(filename, size)
    if invalid is not None:
        return invalid
    try:
        return await _extraction.service.upload(
            filename=Path(filename).name,
            stream=_stream_upload(file),
            content_type=file.content_type or "application/pdf",
            paper_id=paper_id,
            family=family,
            output_name=output_name,
        )
    except ExtractionUnavailable as exc:
        return _down(exc)
    except ExtractionUpstreamError as exc:
        return _upstream_error(f"{_MINERU_MSG}（{exc}）")
    except Exception as exc:
        logger.exception("extraction: 上传失败")
        return _upstream_error(f"提交抽取任务失败：{exc}")