"""OCR 插件 HTTP 接口。 所有识别请求走 FastAPI BackgroundTasks 真正异步执行:上传/解码完成后立即返回 run_id,OCR 在后台任务中完成并通过 run/log 服务写事件,前端轮询日志面板即可。 """ import base64 import logging import tempfile from pathlib import Path from fastapi import APIRouter, BackgroundTasks, HTTPException, UploadFile from pydantic import BaseModel from .core import extract_text, get_model_info_with_providers from app.plugins.run_log import get_run_log_service, RunStatus logger = logging.getLogger(__name__) router = APIRouter() plugin = None def set_plugin_instance(plugin_instance): global plugin plugin = plugin_instance @router.get("/status") async def get_status(): """三态状态:插件启用 / 引擎就绪 / 模型 ID。""" if plugin is None: return {"name": "ocr", "enabled": False, "engine_ready": False} return plugin.get_status() @router.get("/model-info") async def model_info(): return get_model_info_with_providers() class OCRBase64Request(BaseModel): image_base64: str suffix: str = ".jpg" @router.post("/upload") async def ocr_upload(file: UploadFile, background: BackgroundTasks): """上传图片并异步识别,立即返回 run_id。""" run = _create_run(f"图片已上传: {file.filename}") content = await file.read() background.add_task( _run_ocr_task, content, _suffix_from_filename(file.filename), run.run_id, ) return {"run_id": run.run_id, "status": "running"} @router.post("/image-base64") async def ocr_image_base64(request: OCRBase64Request, background: BackgroundTasks): """Base64 图片异步识别。""" try: content = base64.b64decode(request.image_base64) except Exception as e: raise HTTPException(status_code=400, detail=f"Base64 解码失败: {e}") run = _create_run("Base64 图片已接收") background.add_task( _run_ocr_task, content, _normalize_suffix(request.suffix), run.run_id, ) return {"run_id": run.run_id, "status": "running"} # ---- 后台任务 ---- def _create_run(upload_msg: str): """创建 run 并写上传事件;前置校验 plugin 是否启用。""" if plugin is None or not plugin.enabled: raise HTTPException(status_code=400, detail="插件未启用") run_service = get_run_log_service() run = run_service.create_run("ocr") run_service.add_event(run_id=run.run_id, stage="upload", message=upload_msg) return run def _run_ocr_task(content: bytes, suffix: str, run_id: str) -> None: """后台识别任务:落临时文件 → 调引擎 → 写 run 结果 → 清理临时文件。 整段 try/finally 保证临时文件最终被清理;run_service.finish_run 一定被调用。 """ run_service = get_run_log_service() tmp_path: str | None = None try: with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(content) tmp_path = tmp.name result = extract_text(tmp_path, run_id=run_id) if result["success"]: run_service.finish_run( run_id=run_id, status=RunStatus.SUCCEEDED, result={ "text": result["text"], "lines": result["lines"], "boxes": result["boxes"], "model_info": result["model_info"], "timings": result["timings"], }, ) else: run_service.finish_run( run_id=run_id, status=RunStatus.FAILED, error=result["error"], ) except Exception as e: logger.exception("后台 OCR 异常") run_service.add_event(run_id, "ocr", f"OCR 异常: {e}", level="error") run_service.finish_run(run_id=run_id, status=RunStatus.FAILED, error=str(e)) finally: if tmp_path: Path(tmp_path).unlink(missing_ok=True) # ---- 工具 ---- def _suffix_from_filename(filename: str) -> str: return "." + filename.rsplit(".", 1)[-1] if "." in filename else ".jpg" def _normalize_suffix(suffix: str) -> str: return suffix if suffix.startswith(".") else f".{suffix}"