270 lines
7.8 KiB
Python
270 lines
7.8 KiB
Python
"""
|
||
天目 OCR HTTP 客户端。
|
||
|
||
封装对天目微服务的 HTTP 调用，提供同步/异步 OCR 接口。
|
||
后端其他模块通过此客户端调用 OCR 能力，不直接依赖 HTTP 细节。
|
||
|
||
特性：
|
||
- 自动重试（网络故障 / 5xx 错误，指数退避）
|
||
- 连接池复用（httpx.AsyncClient 自动管理）
|
||
- 端点级别超时配置（/healthz 短超时，/v1/ocr 适中，/v1/ocr/vl 长超时）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
|
||
from app.core.logging import logger
|
||
|
||
TIANMU_DEFAULT_BASE = "http://localhost:8101"
|
||
|
||
# 各端点默认超时（秒）
|
||
_TIMEOUTS = {
|
||
"health": 5.0,
|
||
"ocr": 60.0,
|
||
"ocr_vl": 600.0,
|
||
}
|
||
|
||
# 重试配置
|
||
_MAX_RETRIES = 2
|
||
_RETRYABLE_STATUSES = {502, 503, 504}
|
||
_RETRY_BASE_DELAY = 0.5 # 秒
|
||
|
||
|
||
@dataclass
|
||
class OCRBlock:
|
||
"""单行 OCR 识别结果。"""
|
||
|
||
text: str
|
||
confidence: float
|
||
bbox: list[list[float]] = field(default_factory=list)
|
||
|
||
|
||
@dataclass
|
||
class OCRResult:
|
||
"""统一 OCR 响应。"""
|
||
|
||
engine: str
|
||
text: str
|
||
blocks: list[OCRBlock] = field(default_factory=list)
|
||
duration_ms: float = 0.0
|
||
error: str | None = None
|
||
saved_md_path: str | None = None
|
||
|
||
@property
|
||
def success(self) -> bool:
|
||
return self.error is None
|
||
|
||
|
||
class OCRClient:
|
||
"""天目 OCR 服务 HTTP 客户端。
|
||
|
||
用法::
|
||
|
||
client = OCRClient()
|
||
result = await client.ocr(image_bytes)
|
||
await client.close()
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
base_url: str = "",
|
||
http_client: httpx.AsyncClient | None = None,
|
||
timeout: float = 600.0,
|
||
max_retries: int = _MAX_RETRIES,
|
||
) -> None:
|
||
self.base_url = (base_url or TIANMU_DEFAULT_BASE).rstrip("/")
|
||
self._max_retries = max_retries
|
||
self._http = http_client or httpx.AsyncClient(timeout=timeout)
|
||
self._own_http = http_client is None
|
||
|
||
async def close(self) -> None:
|
||
"""释放 HTTP 客户端。"""
|
||
if self._own_http:
|
||
await self._http.aclose()
|
||
|
||
# ── 通用 OCR（PP-OCRv6 优先）─────────────────────────
|
||
|
||
async def ocr(
|
||
self,
|
||
image: bytes | str,
|
||
*,
|
||
engine: str = "auto",
|
||
lang: list[str] | None = None,
|
||
save_md: bool = False,
|
||
) -> OCRResult:
|
||
"""通用文字识别。
|
||
|
||
Args:
|
||
image: 图片二进制数据或 HTTP URL。
|
||
engine: 引擎选择 auto / ppocrv6 / paddleocr-vl。
|
||
lang: 语言列表（默认中英文，仅 ppocrv6 生效）。
|
||
save_md: 是否保存结果为 Markdown 文件。
|
||
|
||
Returns:
|
||
OCRResult: 识别结果。
|
||
"""
|
||
body: dict[str, Any] = {"engine": engine}
|
||
if lang:
|
||
body["lang"] = lang
|
||
if save_md:
|
||
body["save_md"] = True
|
||
|
||
body["image"] = _encode_image(image)
|
||
|
||
data = await self._request_with_retry(
|
||
"POST",
|
||
f"{self.base_url}/v1/ocr",
|
||
json=body,
|
||
timeout=_TIMEOUTS["ocr"],
|
||
)
|
||
if data is None:
|
||
return OCRResult(engine=engine, error="天目服务不可达（重试耗尽）")
|
||
|
||
blocks = [
|
||
OCRBlock(
|
||
text=b["text"],
|
||
confidence=b["confidence"],
|
||
bbox=b.get("bbox", {}).get("points", []),
|
||
)
|
||
for b in data.get("blocks", [])
|
||
]
|
||
|
||
return OCRResult(
|
||
engine=data.get("engine", ""),
|
||
text=data.get("text", ""),
|
||
blocks=blocks,
|
||
duration_ms=data.get("duration_ms", 0),
|
||
error=data.get("error"),
|
||
saved_md_path=data.get("saved_md_path"),
|
||
)
|
||
|
||
# ── 复杂文档解析（PaddleOCR-VL）──────────────────────
|
||
|
||
async def ocr_vl(
|
||
self,
|
||
image: bytes | str,
|
||
*,
|
||
output_format: str = "markdown",
|
||
save_md: bool = False,
|
||
) -> OCRResult:
|
||
"""复杂文档解析（PaddleOCR-VL）。
|
||
|
||
Args:
|
||
image: 图片二进制数据或 HTTP URL。
|
||
output_format: 输出格式 markdown / json / text。
|
||
save_md: 是否保存结果为 Markdown 文件。
|
||
|
||
Returns:
|
||
OCRResult: 结构化解析结果。
|
||
"""
|
||
body: dict[str, Any] = {"output_format": output_format}
|
||
if save_md:
|
||
body["save_md"] = True
|
||
|
||
body["image"] = _encode_image(image)
|
||
|
||
data = await self._request_with_retry(
|
||
"POST",
|
||
f"{self.base_url}/v1/ocr/vl",
|
||
json=body,
|
||
timeout=_TIMEOUTS["ocr_vl"],
|
||
)
|
||
if data is None:
|
||
return OCRResult(engine="paddleocr-vl", error="天目服务不可达（重试耗尽）")
|
||
|
||
return OCRResult(
|
||
engine=data.get("engine", "paddleocr-vl"),
|
||
text=data.get("text", ""),
|
||
duration_ms=data.get("duration_ms", 0),
|
||
error=data.get("error"),
|
||
saved_md_path=data.get("saved_md_path"),
|
||
)
|
||
|
||
# ── 健康检查 ──────────────────────────────────────────
|
||
|
||
async def health(self) -> dict:
|
||
"""检查天目服务健康状态。"""
|
||
resp = await self._http.get(
|
||
f"{self.base_url}/healthz",
|
||
timeout=_TIMEOUTS["health"],
|
||
)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
|
||
# ── 内部：重试逻辑 ────────────────────────────────────
|
||
|
||
async def _request_with_retry(
|
||
self,
|
||
method: str,
|
||
url: str,
|
||
*,
|
||
json: dict[str, Any] | None = None,
|
||
timeout: float = 60.0,
|
||
) -> dict[str, Any] | None:
|
||
"""带指数退避重试的 HTTP 请求。
|
||
|
||
Args:
|
||
method: HTTP 方法。
|
||
url: 请求 URL。
|
||
json: JSON 请求体。
|
||
timeout: 请求超时（秒）。
|
||
|
||
Returns:
|
||
响应 JSON 数据，重试耗尽返回 None。
|
||
"""
|
||
|
||
for attempt in range(self._max_retries + 1):
|
||
try:
|
||
resp = await self._http.request(method, url, json=json, timeout=timeout)
|
||
|
||
# 可重试的服务器错误
|
||
if resp.status_code in _RETRYABLE_STATUSES and attempt < self._max_retries:
|
||
delay = _RETRY_BASE_DELAY * (2 ** attempt)
|
||
logger.warning(
|
||
"天目返回 {}，{}s 后重试 ({}/{})",
|
||
resp.status_code, delay, attempt + 1, self._max_retries,
|
||
)
|
||
await asyncio.sleep(delay)
|
||
continue
|
||
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
|
||
except (httpx.TimeoutException, httpx.ConnectError, httpx.RemoteProtocolError) as e:
|
||
if attempt < self._max_retries:
|
||
delay = _RETRY_BASE_DELAY * (2 ** attempt)
|
||
logger.warning(
|
||
"天目请求失败 ({}，{}s 后重试): {}",
|
||
type(e).__name__, delay, e,
|
||
)
|
||
await asyncio.sleep(delay)
|
||
continue
|
||
|
||
logger.error("天目请求重试耗尽 ({}): {}", type(e).__name__, e)
|
||
|
||
return None
|
||
|
||
|
||
def _encode_image(image: bytes | str | Path) -> str:
|
||
"""将图片编码为 base64 字符串。
|
||
|
||
Args:
|
||
image: 图片二进制数据、文件路径或 HTTP URL。
|
||
|
||
Returns:
|
||
base64 编码字符串（二进制输入）或原 URL。
|
||
"""
|
||
if isinstance(image, Path):
|
||
image = image.read_bytes()
|
||
if isinstance(image, str):
|
||
return image
|
||
|
||
import base64
|
||
return base64.b64encode(image).decode("ascii")
|