services/ocr/client.py

from __future__ import annotations

import base64
from pathlib import Path
from typing import Any

import httpx


class OCRClientError(RuntimeError):
    def __init__(self, code: str, message: str, details: Any = None) -> None:
        super().__init__(f"{code}: {message}")
        self.code = code
        self.message = message
        self.details = details


class OCRClient:
    def __init__(
        self, base_url: str = "http://127.0.0.1:5020", timeout: float = 120.0
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    @staticmethod
    def path_image(path: str | Path) -> dict[str, str]:
        return {"path": str(Path(path).resolve())}

    @staticmethod
    def base64_image(content: bytes) -> dict[str, str]:
        return {"base64": base64.b64encode(content).decode("ascii")}

    def invoke(self, operation: str, payload: dict[str, Any] | None = None) -> Any:
        response = httpx.post(
            f"{self.base_url}/api/invoke",
            json={"operation": operation, "payload": payload or {}},
            timeout=self.timeout,
        )
        try:
            body = response.json()
        except ValueError as exc:
            raise OCRClientError(
                "INVALID_RESPONSE",
                f"OCR 服务返回非 JSON 响应: HTTP {response.status_code}",
                response.text[:500],
            ) from exc
        if not body.get("ok"):
            error = body.get("error") or {}
            raise OCRClientError(
                error.get("code", "UNKNOWN_ERROR"),
                error.get("message", "OCR 服务调用失败"),
                error.get("details"),
            )
        response.raise_for_status()
        return body["result"]

    def status(self) -> dict[str, Any]:
        return self.invoke("providers.status")

    def warmup(
        self,
        providers: list[str],
        languages: dict[str, list[str]] | None = None,
    ) -> dict[str, Any]:
        return self.invoke(
            "providers.warmup",
            {"providers": providers, "languages": languages or {}},
        )

    def recognize(
        self,
        image: dict[str, str],
        *,
        provider: str = "rapidocr",
        languages: list[str] | None = None,
        min_confidence: float = 0.0,
        options: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        return self.invoke(
            "ocr.recognize",
            {
                "image": image,
                "provider": provider,
                "languages": languages,
                "min_confidence": min_confidence,
                "options": options or {},
            },
        )

    def compare(
        self,
        image: dict[str, str],
        *,
        providers: list[str],
        languages: list[str] | None = None,
        min_confidence: float = 0.0,
        options: dict[str, dict[str, Any]] | None = None,
    ) -> dict[str, Any]:
        return self.invoke(
            "ocr.compare",
            {
                "image": image,
                "providers": providers,
                "languages": languages,
                "min_confidence": min_confidence,
                "options": options or {},
            },
        )