plugins/ocr/mcp.py

"""OCR 插件 MCP 工具入口。"""
import base64
import tempfile
from pathlib import Path

from app.mcp.decorators import mcp_tool
from pydantic import BaseModel, Field

from .core import extract_text


class OCROutput(BaseModel):
    success: bool = Field(description="是否成功")
    text: str = Field(description="识别出的文本")
    lines: list[str] = Field(default_factory=list, description="按行识别结果")
    error: str | None = Field(default=None, description="错误信息")


def _decode_to_tempfile(image_b64: str, suffix: str) -> str | None:
    """将 Base64 写入临时文件，返回路径；解码失败返回 None。"""
    if suffix and not suffix.startswith("."):
        suffix = f".{suffix}"
    content = base64.b64decode(image_b64)
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(content)
        return tmp.name


@mcp_tool(
    name="ocr-image",
    title="OCR 图片识别",
    description="从 Base64 图片中识别文字",
    annotations={"readOnlyHint": True, "destructiveHint": False},
)
async def ocr_image_base64(image_base64: str, suffix: str = ".jpg") -> OCROutput:
    try:
        tmp_path = _decode_to_tempfile(image_base64, suffix)
    except Exception as e:
        return OCROutput(success=False, text="", lines=[], error=f"Base64 解码失败: {e}")
    try:
        result = extract_text(tmp_path)
        return OCROutput(**{k: result[k] for k in ("success", "text", "lines", "error")})
    finally:
        Path(tmp_path).unlink(missing_ok=True)


@mcp_tool(
    name="ocr-file",
    title="OCR 文件识别",
    description="从本地图片文件路径中识别文字",
    annotations={"readOnlyHint": True, "destructiveHint": False},
)
async def ocr_file_path(image_path: str) -> OCROutput:
    result = extract_text(image_path)
    return OCROutput(**{k: result[k] for k in ("success", "text", "lines", "error")})