feat: OCR v2 batch + RIIN waybill + 双系统 label-ocr flow(hourly)

2026-07-23 23:44:24 +00:00
parent 51dbde4538
commit 70178e9295
7 changed files with 175 additions and 114 deletions
@@ -0,0 +1,117 @@
+"""面单 OCR enrichment:对所有系统(RIIN+SDS)有待处理面单的订单,用 v2 batch API 批量识别。
+
+对应 label-ocr。每小时跑一次,跟上增量同步节奏。
+按国家分组,每组 200 单/批调 /api/v2/batch(服务端并发下载+重试)。
+"""
+
+from __future__ import annotations
+
+from prefect import flow, task
+
+from api.ocr import OcrClient, country_to_ocr_code
+from flows.shared import SyncDB, build_stats_rows, emit_sync_summary
+
+BATCH_SIZE = 200
+
+_OCR_FIELD_MAP = [
+    ("layout", "ocr_layout"), ("courier", "ocr_courier"), ("tracking", "ocr_tracking"),
+    ("marker", "ocr_marker"), ("s_name", "ocr_s_name"), ("s_street", "ocr_s_street"),
+    ("s_city", "ocr_s_city"), ("s_state", "ocr_s_state"), ("s_zip", "ocr_s_zip"),
+    ("r_name", "ocr_r_name"), ("r_street", "ocr_r_street"), ("r_city", "ocr_r_city"),
+    ("r_state", "ocr_r_state"), ("r_zip", "ocr_r_zip"),
+]
+
+
+def _result_to_fields(r) -> dict[str, str | None]:
+    if not r.ok:
+        return {"ocr_layout": "ERROR"}
+    return {col: (getattr(r, attr) or "") for attr, col in _OCR_FIELD_MAP}
+
+
+@task
+def process_all(pending: list) -> dict[str, int]:
+    db = SyncDB()
+    ocr = OcrClient()
+    stats = {"success": 0, "error": 0, "unsupported": 0, "batches": 0}
+    try:
+        # 1. 按国家分组;不支持的国家直接标 UNSUPPORTED
+        by_country: dict[str, list[tuple[str, object, str]]] = {}
+        unsupp: list[tuple[str, object]] = []
+        for order_no, created_at, pdf_url, dest_country in pending:
+            code = country_to_ocr_code(dest_country)
+            if code:
+                by_country.setdefault(code, []).append((order_no, created_at, pdf_url))
+            else:
+                unsupp.append((order_no, created_at))
+
+        if unsupp:
+            db.update_ocr_batch([(no, ca, {"ocr_layout": "UNSUPPORTED"}) for no, ca in unsupp])
+            stats["unsupported"] = len(unsupp)
+            print(f"[OCR] 不支持国家: {len(unsupp)} 单 → UNSUPPORTED")
+
+        # 2. 每个国家分批调 v2 batch API
+        for country_code, orders in by_country.items():
+            total = len(orders)
+            for i in range(0, total, BATCH_SIZE):
+                batch = orders[i : i + BATCH_SIZE]
+                created_map = {no: ca for no, ca, _ in batch}
+                urls = [{"url": url, "filename": no} for no, _, url in batch]
+                batch_no = i // BATCH_SIZE + 1
+                total_batches = -(-total // BATCH_SIZE)
+                try:
+                    results = ocr.recognize_batch(urls, country_code)
+                except Exception as err:  # noqa: BLE001
+                    print(f"[OCR] {country_code} batch {batch_no}/{total_batches} 请求失败: {str(err)[:100]}")
+                    db.update_ocr_batch([(no, created_map[no], {"ocr_layout": "ERROR"}) for no in created_map])
+                    stats["error"] += len(batch)
+                    continue
+
+                updates = []
+                for r in results:
+                    order_no = r.file
+                    created_at = created_map.get(order_no)
+                    if created_at is None:
+                        continue
+                    fields = _result_to_fields(r)
+                    updates.append((order_no, created_at, fields))
+                    stats["success" if r.ok else "error"] += 1
+                db.update_ocr_batch(updates)
+                stats["batches"] += 1
+                ok = sum(1 for r in results if r.ok)
+                print(f"[OCR] {country_code} batch {batch_no}/{total_batches}: "
+                      f"{ok}/{len(results)} 成功")
+        return stats
+    finally:
+        ocr.close()
+
+
+@flow(name="label-ocr", log_prints=True)
+def label_ocr() -> dict:
+    db = SyncDB()
+    before = db.count_ocr_status()
+    print(f"OCR 状态(开始): {before}")
+
+    pending = db.get_ocr_pending()
+    print(f"待 OCR: {len(pending)} 单 (RIIN+SDS)")
+    if not pending:
+        emit_sync_summary("面单 OCR", build_stats_rows("ALL", {}, [
+            ("待处理", 0), ("累计已完成", before.get("done", 0)),
+        ]))
+        return {"processed": 0, **before}
+
+    stats = process_all(pending)
+    after = db.count_ocr_status()
+    total_processed = stats["success"] + stats["error"] + stats["unsupported"]
+    print(f"完成: 处理 {total_processed} 单 / {stats['batches']} 批 | "
+          f"成功 {stats['success']} 失败 {stats['error']} 不支持 {stats['unsupported']}")
+    print(f"累计: {after}")
+    emit_sync_summary("面单 OCR", build_stats_rows("ALL", {}, [
+        ("本批处理", total_processed), ("成功", stats["success"]),
+        ("失败", stats["error"]), ("不支持", stats["unsupported"]),
+        ("累计已完成", after.get("done", 0)), ("累计待处理", after.get("pending", 0)),
+    ]))
+    return {"processed": total_processed, **stats, **after}
+
+
+if __name__ == "__main__":
+    label_ocr.serve(name="label-ocr", cron="0 * * * *")
@@ -1,98 +0,0 @@
-"""SDS 面单 OCR enrichment:对已发货的 SDS 订单,用面单 PDF URL 调 OCR 服务回填标签字段。
-
-对应 sds-label-ocr。独立于同步 flow,每 4 小时跑一批。
-"""
-
-from __future__ import annotations
-
-from concurrent.futures import ThreadPoolExecutor, as_completed
-
-from prefect import flow, task
-
-from api.ocr import OcrClient, country_to_ocr_code
-from flows.shared import SyncDB, build_stats_rows, emit_sync_summary
-
-OCR_CONCURRENCY = 5
-SHIPPED_STATUSES = ["已完成", "部分发货", "备货中", "待确认"]
-
-_OCR_FIELD_MAP = [
-    ("layout", "ocr_layout"), ("courier", "ocr_courier"), ("tracking", "ocr_tracking"),
-    ("marker", "ocr_marker"), ("s_name", "ocr_s_name"), ("s_street", "ocr_s_street"),
-    ("s_city", "ocr_s_city"), ("s_state", "ocr_s_state"), ("s_zip", "ocr_s_zip"),
-    ("r_name", "ocr_r_name"), ("r_street", "ocr_r_street"), ("r_city", "ocr_r_city"),
-    ("r_state", "ocr_r_state"), ("r_zip", "ocr_r_zip"),
-]
-
-
-def _process_one(db: SyncDB, ocr: OcrClient, order_no: str, created_at, pdf_url: str, dest_country: str | None) -> str:
-    country_code = country_to_ocr_code(dest_country)
-    if not country_code:
-        db.update_ocr_fields(order_no, created_at, {"ocr_layout": "UNSUPPORTED"})
-        return "unsupported"
-    try:
-        r = ocr.recognize_url(pdf_url, country_code)
-    except Exception as err:  # noqa: BLE001
-        print(f"[OCR] {order_no} 请求失败: {str(err)[:100]}")
-        db.update_ocr_fields(order_no, created_at, {"ocr_layout": "ERROR"})
-        return "error"
-    if not r.ok:
-        print(f"[OCR] {order_no} 识别失败: {(r.error or '')[:80]}")
-        db.update_ocr_fields(order_no, created_at, {"ocr_layout": "ERROR"})
-        return "error"
-    fields = {col: getattr(r, attr) or "" for attr, col in _OCR_FIELD_MAP}
-    db.update_ocr_fields(order_no, created_at, fields)
-    return "success"
-
-
-@task
-def process_batch(pending: list) -> dict[str, int]:
-    db = SyncDB()
-    ocr = OcrClient()
-    stats = {"success": 0, "error": 0, "unsupported": 0}
-    try:
-        done = 0
-        total = len(pending)
-        with ThreadPoolExecutor(max_workers=min(OCR_CONCURRENCY, total)) as pool:
-            futures = {
-                pool.submit(_process_one, db, ocr, no, ca, url, dc): no
-                for no, ca, url, dc in pending
-            }
-            for fut in as_completed(futures):
-                result = fut.result()
-                stats[result] += 1
-                done += 1
-                if done % 50 == 0 or done == total:
-                    print(f"[OCR] 进度: {done}/{total}  成功 {stats['success']} / 失败 {stats['error']} / 不支持 {stats['unsupported']}")
-        stats["processed"] = total
-        return stats
-    finally:
-        ocr.close()
-
-
-@flow(name="sds-label-ocr", log_prints=True)
-def sds_label_ocr(batch_size: int = 200) -> dict:
-    db = SyncDB()
-    before = db.count_ocr_status()
-    print(f"OCR 状态(开始): {before}")
-
-    pending = db.get_ocr_pending(batch_size)
-    print(f"待 OCR: {len(pending)} 单 (batch_size={batch_size})")
-    if not pending:
-        emit_sync_summary("SDS 面单 OCR", build_stats_rows("SDS", {}, [
-            ("待处理", 0), ("累计已完成", before.get("done", 0)), ("累计失败", before.get("error", 0)),
-        ]))
-        return {"processed": 0, **before}
-
-    stats = process_batch(pending)
-    after = db.count_ocr_status()
-    print(f"完成: {stats}  累计: {after}")
-    emit_sync_summary("SDS 面单 OCR", build_stats_rows("SDS", {}, [
-        ("本批处理", stats["processed"]), ("本批成功", stats["success"]),
-        ("本批失败", stats["error"]), ("本批不支持", stats["unsupported"]),
-        ("累计已完成", after.get("done", 0)), ("累计待处理", after.get("pending", 0)),
-    ]))
-    return {**stats, **after}
-
-
-if __name__ == "__main__":
-    sds_label_ocr.serve(name="sds-label-ocr", cron="0 */4 * * *")
@@ -100,22 +100,25 @@ class SyncDB:

    # ---------- OCR 面单 enrichment ----------

-    def get_ocr_pending(self, limit: int = 200) -> list[tuple[str, datetime, str, str | None]]:
-        """查待 OCR 的 SDS 订单:(orderNo, createdAt, labelPdfUrl, destinationCountry)。
+    def get_ocr_pending(self, limit: int = 0) -> list[tuple[str, datetime, str, str | None]]:
+        """查待 OCR 订单(所有系统):(orderNo, createdAt, labelPdfUrl, destinationCountry)。

-        条件:SDS、有 labelPdfUrl、ocr_layout 未处理(Null)、已发货/已完成。
+        条件:有 labelPdfUrl、ocr_layout 未处理(Null)。
+        limit=0 表示不限。
        """
+        sql = (
+            'SELECT "orderNo","createdAt","labelPdfUrl","destinationCountry" '
+            'FROM "Orders" '
+            'WHERE "labelPdfUrl" IS NOT NULL AND "ocr_layout" IS NULL '
+            'ORDER BY "orderedAt" DESC'
+        )
+        params: list = []
+        if limit > 0:
+            sql += " LIMIT %s"
+            params.append(limit)
        with self._conn() as conn:
            with conn.cursor() as cur:
-                cur.execute(
-                    '''SELECT "orderNo","createdAt","labelPdfUrl","destinationCountry"
-                       FROM "Orders"
-                       WHERE "systemName"=%s AND "labelPdfUrl" IS NOT NULL
-                         AND "ocr_layout" IS NULL
-                         AND "systemStatus"=ANY(%s)
-                       ORDER BY "orderedAt" DESC LIMIT %s''',
-                    ("SDS", ["已完成", "部分发货", "备货中", "待确认"], limit),
-                )
+                cur.execute(sql, params)
                return [(r[0], r[1], r[2], r[3]) for r in cur.fetchall()]

    def update_ocr_fields(
@@ -135,6 +138,34 @@ class SyncDB:
                )
            conn.commit()

+    def update_ocr_batch(self, updates: list[tuple[str, datetime, dict[str, str | None]]]) -> int:
+        """批量更新 ocr 字段(单连接单事务)。updates: [(orderNo, createdAt, {field: value})]。
+
+        返回更新行数。
+        """
+        if not updates:
+            return 0
+        rows = 0
+        with self._conn() as conn:
+            try:
+                with conn.cursor() as cur:
+                    for order_no, created_at, fields in updates:
+                        if not fields:
+                            continue
+                        assignments = ", ".join(f'"{k}"=%s' for k in fields)
+                        params = list(fields.values()) + [order_no, created_at]
+                        cur.execute(
+                            f'''UPDATE "Orders" SET {assignments}
+                                WHERE "orderNo"=%s AND "createdAt"=%s''',
+                            params,
+                        )
+                        rows += cur.rowcount
+                conn.commit()
+            except Exception:
+                conn.rollback()
+                raise
+        return rows
+
    def count_ocr_status(self) -> dict[str, int]:
        """统计 OCR 处理状态(done/error/unsupported/pending)。"""
        with self._conn() as conn:
@@ -195,7 +195,7 @@ def map_riin_order(raw: Any) -> OrderUpsertInput:
        addressLine3=raw.receiver_town,
        postalCode=raw.post_code,
        storeName=raw.platform_shop_name,
-        labelPdfUrl=None,
+        labelPdfUrl=raw.waybill,
    )