diff --git a/src/polygateway/transports/monkey_ocr.py b/src/polygateway/transports/monkey_ocr.py index 9e0edc9..254cc55 100644 --- a/src/polygateway/transports/monkey_ocr.py +++ b/src/polygateway/transports/monkey_ocr.py @@ -14,7 +14,7 @@ import math import zipfile from contextlib import contextmanager from io import BytesIO -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import httpx @@ -42,6 +42,50 @@ if TYPE_CHECKING: _HEALTH_TIMEOUT_S = 5.0 # 探测常量(同 jitter 系数先例,不进配置) +_MediaKind = Literal["pdf", "image", "unknown"] + +# 已知图片魔数。起始命中即判图片,且**优先于** PDF 判定: 含 XMP 元数据的图片 +# 可能在前 1KB 内带 `%PDF-` 字样,不能因此被当成 PDF。 +_IMAGE_MAGICS: tuple[bytes, ...] = ( + b"\xff\xd8\xff", # JPEG + b"\x89PNG\r\n\x1a\n", # PNG + b"II*\x00", # TIFF 小端 + b"MM\x00*", # TIFF 大端 + b"BM", # BMP + b"GIF87a", + b"GIF89a", +) +_PDF_HEADER = b"%PDF-" +# Adobe 阅读器容忍头部前有垃圾字节(惯例而非 ISO 32000 正文);服务端实测可解析 +# 前缀 200 字节垃圾的 PDF(#34 设计 E3) +_PDF_HEADER_WINDOW = 1024 + + +def _sniff_media(data: bytes) -> _MediaKind: + """按内容判定上传类型: 图片魔数 → 前 1KB 内 `%PDF-` → 未知。 + + 服务端按**文件名扩展名**分派(`.pdf` 走 PDF 数据集,其余走图片数据集),而标签 + 若与内容不符会得到 HTTP 500(`is no PDF`)或——在 `/ocr/text` 上——`success:true` + 的幻觉文本(#34 设计 E2/E4)。所以标签只能由库按内容定,不开放给调用方声明。 + """ + if data.startswith(_IMAGE_MAGICS) or (data[:4] == b"RIFF" and data[8:12] == b"WEBP"): + return "image" + if _PDF_HEADER in data[:_PDF_HEADER_WINDOW]: + return "pdf" + return "unknown" + + +def _upload_part(data: bytes) -> tuple[str, bytes, str]: + """multipart 的 (filename, content, content_type)。 + + 只有 PDF 改标 `doc.pdf`;**其余(含未知)一律沿用历史的 `image.jpg`**——服务对 + PNG/TIFF/BMP 字节标 `image.jpg` 照样按内容解码(2026-10-21 实测),任何"更精确" + 的改标都只会把现在能成功的路径变成失败(webp/gif 没有对应的服务端扩展名)。 + """ + if _sniff_media(data) == "pdf": + return ("doc.pdf", data, "application/pdf") + return ("image.jpg", data, "image/jpeg") + def _default_client_factory(source: SourceConfig) -> httpx.AsyncClient: return httpx.AsyncClient( @@ -205,9 +249,10 @@ class MonkeyOcrTransport: return client @staticmethod - def _files(image: bytes) -> dict[str, tuple[str, bytes, str]]: - # 固定文件名/类型(CHS invokers.py:496 原样;服务按内容处理,35 样本实测) - return {"file": ("image.jpg", image, "image/jpeg")} + def _files(data: bytes) -> dict[str, tuple[str, bytes, str]]: + # 标签按内容判定(#34): 非 PDF 沿用 CHS invokers.py:496 的 image.jpg, + # PDF 改标 doc.pdf——服务按扩展名分派,固定名对 PDF 是错的 + return {"file": _upload_part(data)} def _response_json(self, resp: httpx.Response, source: SourceConfig, operation: str) -> dict: """响应体 JSON 化;坏 JSON/非 dict → Transient(设计 §4: 服务侧瞬时异常)。""" diff --git a/tests/unit/test_monkey_ocr.py b/tests/unit/test_monkey_ocr.py index f0c100e..21ae776 100644 --- a/tests/unit/test_monkey_ocr.py +++ b/tests/unit/test_monkey_ocr.py @@ -23,6 +23,8 @@ from polygateway.transports.monkey_ocr import ( MonkeyOcrTransport, _classify_status, _parse_middle_json, + _sniff_media, + _upload_part, ) from polygateway.types import SourceConfig @@ -167,6 +169,86 @@ class TestRecognizeText: assert b'filename="image.jpg"' in seen["body"] # CHS invokers.py:496 固定名 +_PDF_BODY = b"%PDF-1.4\n%synthetic body" +_JUNK = b"x" * 200 + + +class TestSniffMedia: + """上传类型按内容判定(#34 设计 §3 Q2): 图片魔数优先 → 前 1KB 内 `%PDF-` → 未知。""" + + @pytest.mark.parametrize( + ("data", "expected"), + [ + pytest.param(_PDF_BODY, "pdf", id="pdf"), + pytest.param(_JUNK + _PDF_BODY, "pdf", id="pdf-with-junk-prefix"), + pytest.param(b"\xff\xd8\xff\xe0" + b"0" * 20, "image", id="jpeg"), + pytest.param(b"\x89PNG\r\n\x1a\n" + b"0" * 20, "image", id="png"), + pytest.param(b"II*\x00" + b"0" * 20, "image", id="tiff-le"), + pytest.param(b"MM\x00*" + b"0" * 20, "image", id="tiff-be"), + pytest.param(b"BM" + b"0" * 20, "image", id="bmp"), + pytest.param(b"GIF89a" + b"0" * 20, "image", id="gif89a"), + pytest.param(b"GIF87a" + b"0" * 20, "image", id="gif87a"), + pytest.param(b"RIFF\x00\x00\x00\x00WEBPVP8 ", "image", id="webp"), + # 图片魔数优先: 含 XMP 元数据的图片可能在 1KB 内带 `%PDF-` 字样 + pytest.param(b"\xff\xd8\xff\xe0" + b"%PDF-", "image", id="jpeg-mentions-pdf"), + pytest.param(b"paywall", "unknown", id="html"), + pytest.param(b"GARBAGE-NOT-AN-IMAGE", "unknown", id="garbage"), + pytest.param(b"", "unknown", id="empty"), + # 1KB 窗口之外的 `%PDF-` 不算 + pytest.param(b"x" * 1100 + b"%PDF-", "unknown", id="header-beyond-window"), + ], + ) + def test_sniff(self, data, expected): + assert _sniff_media(data) == expected + + def test_upload_part_labels_pdf_only(self): + assert _upload_part(_PDF_BODY) == ("doc.pdf", _PDF_BODY, "application/pdf") + assert _upload_part(_JUNK + _PDF_BODY)[0] == "doc.pdf" + + @pytest.mark.parametrize("data", [b"\xff\xd8\xff\xe0abc", b"\x89PNG\r\n\x1a\nabc", b"junk", b""]) + def test_everything_else_keeps_the_legacy_label(self, data): + """非 PDF 路径字节级不变: 服务对图片按内容解码(2026-10-21 实测 PNG/TIFF/BMP 标 image.jpg 成功)。""" + assert _upload_part(data) == ("image.jpg", data, "image/jpeg") + + +class TestUploadLabelPerMethod: + """三个调用点共用同一个标签点: 漏掉任何一个,E4 的静默错误结果就在那里复活。""" + + @staticmethod + def _capturing_transport(seen): + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/ocr/text": + seen["body"] = request.read() + return httpx.Response(200, json=_text_body()) + if request.url.path == "/parse": + seen["body"] = request.read() + return httpx.Response(200, json=_parse_body()) + return httpx.Response(200, content=_zip_bytes()) + + return _transport_for(handler) + + async def test_recognize_text_labels_pdf_as_pdf(self): + seen = {} + t = self._capturing_transport(seen) + await t.recognize_text(image=_PDF_BODY, source=_source(), call_id="c1") + assert b'filename="doc.pdf"' in seen["body"] + assert b"Content-Type: application/pdf" in seen["body"] + + async def test_parse_layout_labels_pdf_as_pdf(self): + seen = {} + t = self._capturing_transport(seen) + await t.parse_layout(image=_PDF_BODY, source=_source(), call_id="c1") + assert b'filename="doc.pdf"' in seen["body"] + assert b"Content-Type: application/pdf" in seen["body"] + + async def test_parse_layout_keeps_image_label_for_images(self): + seen = {} + t = self._capturing_transport(seen) + await t.parse_layout(image=b"\xff\xd8\xff\xe0abc", source=_source(), call_id="c1") + assert b'filename="image.jpg"' in seen["body"] + assert b"Content-Type: image/jpeg" in seen["body"] + + class TestParseLayout: async def test_full_chain_relative_url(self): t = _transport_for(