diff --git a/api/ocr/client.py b/api/ocr/client.py index 4c91487..7c662e3 100644 --- a/api/ocr/client.py +++ b/api/ocr/client.py @@ -25,6 +25,16 @@ class OcrClient: resp.raise_for_status() return OcrResult.model_validate(resp.json()) + def recognize_batch(self, urls: list[dict[str, str]], country: str = "US") -> list[OcrResult]: + """v2 批量识别:urls=[{"url":"...", "filename":"orderNo"}], 最多 200 项/批。 + + 返回 OcrResult 列表,file 字段 = filename(用于回匹配 orderNo)。 + """ + resp = self._http.post("/api/v2/batch", json={"urls": urls, "country": country}, timeout=300.0) + resp.raise_for_status() + data = resp.json() + return [OcrResult.model_validate(r) for r in data.get("results", [])] + def health(self) -> bool: """健康检查。""" try: diff --git a/api/riin/models.py b/api/riin/models.py index 71f72ba..0a4d857 100644 --- a/api/riin/models.py +++ b/api/riin/models.py @@ -122,6 +122,7 @@ class Order(_RiinModel): express_fee: float | None = None express_source: int | None = None courier_number: str | None = None + waybill: str | None = None consignee_name: str | None = None phone: str | None = None email: str | None = None diff --git a/flows/label_ocr.py b/flows/label_ocr.py new file mode 100644 index 0000000..b53a056 --- /dev/null +++ b/flows/label_ocr.py @@ -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 * * * *") diff --git a/flows/sds/label_ocr.py b/flows/sds/label_ocr.py deleted file mode 100644 index 55ccc24..0000000 --- a/flows/sds/label_ocr.py +++ /dev/null @@ -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 * * *") diff --git a/flows/shared/db.py b/flows/shared/db.py index 9fb2db1..7bdeab4 100644 --- a/flows/shared/db.py +++ b/flows/shared/db.py @@ -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: diff --git a/flows/shared/mappers.py b/flows/shared/mappers.py index 7e37678..04e9b3e 100644 --- a/flows/shared/mappers.py +++ b/flows/shared/mappers.py @@ -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, ) diff --git a/prefect.yaml b/prefect.yaml index 4f6a7a6..228910b 100644 --- a/prefect.yaml +++ b/prefect.yaml @@ -80,10 +80,10 @@ deployments: timezone: Asia/Shanghai work_pool: name: inkreach-pool -- name: sds-label-ocr - entrypoint: flows/sds/label_ocr.py:sds_label_ocr +- name: label-ocr + entrypoint: flows/label_ocr.py:label_ocr schedules: - - cron: 0 */4 * * * + - cron: 0 * * * * timezone: Asia/Shanghai work_pool: name: inkreach-pool