"""Small async client for the Baidu OCR endpoints used by the application.""" import asyncio import base64 import time from typing import Any from urllib.parse import urlparse import httpx from app.core.config import get_settings from app.core.exceptions import AppException class BaiduOcrClient: token_url = "https://aip.baidubce.com/oauth/2.0/token" def __init__(self, transport: httpx.AsyncBaseTransport | None = None) -> None: settings = get_settings() self.api_key = settings.baidu_ocr_api_key.get_secret_value() if settings.baidu_ocr_api_key else "" self.secret_key = settings.baidu_ocr_secret_key.get_secret_value() if settings.baidu_ocr_secret_key else "" self.image_url = settings.baidu_ocr_image_url.strip() self.document_url = settings.baidu_ocr_document_url.strip().rstrip("/") self.timeout_seconds = settings.baidu_ocr_timeout_seconds self.poll_interval = settings.baidu_ocr_poll_interval_seconds self.transport = transport self._access_token = "" self._token_expires_at = 0.0 @property def configured(self) -> bool: return bool(self.api_key and self.secret_key) async def recognize_image(self, content: bytes) -> tuple[str, float | None, int]: token = await self._get_access_token() payload = { "image": base64.b64encode(content).decode("ascii"), "language_type": "CHN_ENG", "detect_direction": "true", "probability": "true", } body = await self._post_form(self.image_url, token, payload) words = body.get("words_result") or [] lines = [str(item.get("words") or "").strip() for item in words if isinstance(item, dict)] confidences = [ float(item["probability"]["average"]) for item in words if isinstance(item, dict) and isinstance(item.get("probability"), dict) and isinstance(item["probability"].get("average"), (int, float)) ] text = "\n".join(line for line in lines if line) if not text: raise AppException("图片中未识别到文字，请重新拍摄或手动填写", code=422, status_code=422) confidence = sum(confidences) / len(confidences) if confidences else None return text, confidence, 1 async def recognize_document(self, content: bytes, file_name: str) -> tuple[str, float | None, int]: token = await self._get_access_token() body = await self._post_form( self.document_url, token, { "file_data": base64.b64encode(content).decode("ascii"), "file_name": file_name, "merge_tables": "true", "return_span_boxes": "true", "angle_adjust": "true", }, ) task_id = str((body.get("result") or {}).get("task_id") or "") if not task_id: raise AppException("百度文档识别未返回任务编号", code=502, status_code=502) deadline = time.monotonic() + self.timeout_seconds query_url = f"{self.document_url}/query" result: dict[str, Any] = {} while time.monotonic() < deadline: await asyncio.sleep(self.poll_interval) query_body = await self._post_form(query_url, token, {"task_id": task_id}) result = query_body.get("result") or {} status = result.get("status") if status == "success": break if status == "failed": detail = str(result.get("task_error") or "文档解析失败") raise AppException(f"百度文档识别失败：{detail}", code=502, status_code=502) else: raise AppException("文档识别超时，请稍后重试", code=504, status_code=504) parse_url = str(result.get("parse_result_url") or "") markdown_url = str(result.get("markdown_url") or "") if parse_url: parsed = await self._download_result(parse_url, expect_json=True) text, pages = self._text_from_parse_result(parsed) elif markdown_url: text = str(await self._download_result(markdown_url, expect_json=False)) pages = max(1, text.count("\n---\n") + 1) else: raise AppException("百度文档识别未返回解析结果", code=502, status_code=502) if not text.strip(): raise AppException("文档中未识别到文字，请重新上传或手动填写", code=422, status_code=422) return text[:100_000], None, pages async def _get_access_token(self) -> str: if not self.configured: raise AppException("未配置百度 OCR API Key 或 Secret Key", code=503, status_code=503) if self._access_token and time.monotonic() < self._token_expires_at: return self._access_token timeout = httpx.Timeout(15, connect=10) async with httpx.AsyncClient(timeout=timeout, transport=self.transport) as client: try: response = await client.post( self.token_url, params={ "grant_type": "client_credentials", "client_id": self.api_key, "client_secret": self.secret_key, }, ) response.raise_for_status() body = response.json() except (httpx.HTTPError, ValueError) as exc: raise AppException("无法连接百度 OCR 鉴权服务", code=502, status_code=502) from exc token = str(body.get("access_token") or "") if not token: raise AppException("百度 OCR 密钥无效或服务尚未开通", code=502, status_code=502) expires_in = int(body.get("expires_in") or 2_592_000) self._access_token = token self._token_expires_at = time.monotonic() + max(60, expires_in - 300) return token async def _post_form(self, url: str, token: str, data: dict[str, str]) -> dict[str, Any]: timeout = httpx.Timeout(self.timeout_seconds, connect=15) async with httpx.AsyncClient(timeout=timeout, transport=self.transport) as client: try: response = await client.post(url, params={"access_token": token}, data=data) response.raise_for_status() body = response.json() except httpx.TimeoutException as exc: raise AppException("百度 OCR 请求超时，请重试", code=504, status_code=504) from exc except (httpx.HTTPError, ValueError) as exc: raise AppException("百度 OCR 服务暂时不可用", code=502, status_code=502) from exc if body.get("error_code") not in (None, 0, "0"): message = str(body.get("error_msg") or "识别失败") raise AppException(f"百度 OCR 返回错误：{message}", code=502, status_code=502) return body async def _download_result(self, url: str, *, expect_json: bool) -> dict[str, Any] | str: parsed = urlparse(url) allowed = parsed.scheme == "https" and ( parsed.hostname == "baidubce.com" or bool(parsed.hostname and parsed.hostname.endswith((".baidubce.com", ".bcebos.com"))) ) if not allowed: raise AppException("百度 OCR 返回了不受信任的结果地址", code=502, status_code=502) timeout = httpx.Timeout(self.timeout_seconds, connect=15) async with httpx.AsyncClient(timeout=timeout, transport=self.transport, follow_redirects=True) as client: try: response = await client.get(url) response.raise_for_status() return response.json() if expect_json else response.text except (httpx.HTTPError, ValueError) as exc: raise AppException("无法下载百度 OCR 解析结果", code=502, status_code=502) from exc @staticmethod def _text_from_parse_result(body: dict[str, Any]) -> tuple[str, int]: pages = body.get("pages") or [] chunks: list[str] = [] for page in pages: if not isinstance(page, dict): continue page_text = str(page.get("text") or "").strip() if page_text: chunks.append(page_text) for table in page.get("tables") or []: if isinstance(table, dict) and table.get("markdown"): chunks.append(str(table["markdown"]).strip()) return "\n\n".join(chunks), max(1, len(pages))