iomgaa
86b19d1e07
- app/ports.py: 新增 OCRProvider Protocol（runtime_checkable，与 EmbeddingProvider 同级），定义 async transcribe_frames 端口 - adapters/ocr.py: 从 TRM4 core/tree/ocr.py 保真迁移 MonkeyOCRClient - assert → ValueError（P5 防御性校验） - 公开方法改 async（asyncio.to_thread 包装同步 HTTP） - 内部逻辑不变：多端点轮询、线程安全 Session、单帧降级、行去重 - tests/unit/test_ocr_adapter.py: 17 个测试覆盖 Protocol 合规、 构造校验、健康检查、转录、降级、去重、轮询 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
129 lines
4.6 KiB
Python
129 lines
4.6 KiB
Python
"""MonkeyOCR HTTP 客户端 — 帧文字转录的异构硬证据源。
|
||
|
||
服务由用户在 LAN 部署（双端点轮询）；请求必须绕过代理（trust_env=False）。
|
||
实现 OCRProvider Protocol（app/ports.py）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import itertools
|
||
import threading
|
||
from pathlib import Path # noqa: TC003 — 运行时需要（方法签名 + open()）
|
||
|
||
import requests
|
||
from loguru import logger
|
||
|
||
_TIMEOUT_S = 15
|
||
|
||
|
||
class MonkeyOCRClient:
|
||
"""MonkeyOCR 服务客户端：多端点轮询、单帧失败降级为跳过。
|
||
|
||
关键实现细节：实例可被多线程共享——端点轮询加锁、Session 线程局部
|
||
（A/B 评测会以 4 线程并发调用同一实例）。
|
||
|
||
参数:
|
||
urls: 服务端点列表（如 ["http://10.77.0.20:7866", ...]），非空。
|
||
|
||
异常:
|
||
ValueError: urls 为空时抛出。
|
||
"""
|
||
|
||
def __init__(self, urls: list[str]) -> None:
|
||
if not urls:
|
||
raise ValueError("MonkeyOCR 端点列表不能为空")
|
||
self._urls = [u.rstrip("/") for u in urls]
|
||
self._rr = itertools.cycle(self._urls)
|
||
self._rr_lock = threading.Lock()
|
||
self._local = threading.local()
|
||
|
||
def _get_session(self) -> requests.Session:
|
||
"""返回当前线程专属的 Session（惰性创建并复用，trust_env=False 绕代理）。"""
|
||
session = getattr(self._local, "session", None)
|
||
if session is None:
|
||
session = requests.Session()
|
||
session.trust_env = False # LAN 直连，绕过代理
|
||
self._local.session = session
|
||
return session
|
||
|
||
def _check_health_sync(self) -> None:
|
||
"""同步预检所有端点，任一不可达即抛错（供 asyncio.to_thread 调用）。
|
||
|
||
异常:
|
||
RuntimeError: 端点不可达或 /health 非 2xx。
|
||
"""
|
||
for url in self._urls:
|
||
try:
|
||
resp = self._get_session().get(f"{url}/health", timeout=5)
|
||
except requests.RequestException as e:
|
||
raise RuntimeError(f"MonkeyOCR 端点不可达: {url}: {e}") from e
|
||
if not resp.ok:
|
||
raise RuntimeError(f"MonkeyOCR 健康检查失败: {url}: {resp.status_code}")
|
||
|
||
async def check_health(self) -> None:
|
||
"""异步预检所有端点，任一不可达即抛错（A/B qtr_ocr 臂启动门）。
|
||
|
||
异常:
|
||
RuntimeError: 端点不可达或 /health 非 2xx。
|
||
"""
|
||
await asyncio.to_thread(self._check_health_sync)
|
||
|
||
def _transcribe_frames_sync(self, frame_paths: list[Path]) -> str:
|
||
"""同步逐帧转录并拼接（供 asyncio.to_thread 调用）。
|
||
|
||
参数:
|
||
frame_paths: 帧文件路径列表。
|
||
|
||
返回:
|
||
"帧1: <行1> | <行2>\\n帧2: ..." 格式文本；无任何有效结果时空串。
|
||
"""
|
||
parts: list[str] = []
|
||
for i, path in enumerate(frame_paths, 1):
|
||
lines = self._transcribe_one(path)
|
||
if lines:
|
||
parts.append(f"帧{i}: " + " | ".join(lines))
|
||
return "\n".join(parts)
|
||
|
||
async def transcribe_frames(self, frame_paths: list[Path]) -> str:
|
||
"""异步逐帧转录并拼接为注入文本；单帧失败跳过，全失败返回空串。
|
||
|
||
参数:
|
||
frame_paths: 帧文件路径列表。
|
||
|
||
返回:
|
||
"帧1: <行1> | <行2>\\n帧2: ..." 格式文本；无任何有效结果时空串。
|
||
"""
|
||
return await asyncio.to_thread(self._transcribe_frames_sync, frame_paths)
|
||
|
||
def _transcribe_one(self, path: Path) -> list[str]:
|
||
"""单帧转录：空结果/单字符行过滤 + 帧内行级去重。
|
||
|
||
参数:
|
||
path: 帧文件路径。
|
||
|
||
返回:
|
||
过滤去重后的文本行列表；请求失败或无有效行时空列表。
|
||
"""
|
||
with self._rr_lock:
|
||
url = next(self._rr)
|
||
try:
|
||
with open(path, "rb") as f:
|
||
resp = self._get_session().post(
|
||
f"{url}/ocr/text", files={"file": f}, timeout=_TIMEOUT_S
|
||
)
|
||
resp.raise_for_status()
|
||
content = resp.json().get("content", "")
|
||
except (requests.RequestException, ValueError) as e:
|
||
logger.warning("MonkeyOCR 单帧转录失败，跳过 {}: {}", path.name, e)
|
||
return []
|
||
seen: set[str] = set()
|
||
lines: list[str] = []
|
||
for ln in content.splitlines():
|
||
ln = ln.strip()
|
||
if len(ln) <= 1 or ln in seen:
|
||
continue
|
||
seen.add(ln)
|
||
lines.append(ln)
|
||
return lines
|