fix(avatar): OCR image-only knowledge PDFs

digital-avatar-app/backend/tests/test_knowledge_storage.py

@@ -238,6 +238,53 @@ def test_background_vectorizer_keeps_failure_reason_for_retry(
         db.close()
 
 
+def test_background_vectorizer_uses_ocr_for_image_only_pdf(
+    tmp_path: Path,
+    authorization_context,
+):
+    context = authorization_context
+    with (
+        patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)),
+        patch("routers.knowledge.knowledge_vectorizer.enqueue"),
+    ):
+        response = client.post(
+            f"/api/avatar/{context['avatar'].id}/knowledge/docs",
+            headers=context["owner_headers"],
+            files={"file": ("scanned.pdf", b"image-only-pdf", "application/pdf")},
+        )
+
+    payload = response.json()["data"]
+    progress = []
+    with (
+        patch("services.knowledge_vectorizer.UPLOAD_DIR", str(tmp_path)),
+        patch("services.knowledge_vectorizer.embeddings.extract_text", return_value=""),
+        patch(
+            "services.knowledge_vectorizer.extract_scanned_pdf_text",
+            side_effect=lambda _db, _avatar, _path, on_progress: (
+                on_progress(1, 2), on_progress(2, 2), "扫描页文字"
+            )[-1],
+        ) as ocr,
+        patch("services.knowledge_vectorizer.embeddings.embed", return_value=[[1.0, 0.0]]),
+        patch.object(knowledge_vectorizer, "_set_progress", wraps=knowledge_vectorizer._set_progress) as set_progress,
+    ):
+        knowledge_vectorizer.vectorize_document(payload["id"])
+        progress = [(call.args[2], call.args[3]) for call in set_progress.call_args_list]
+
+    db = SessionLocal()
+    try:
+        stored = db.query(KnowledgeDoc).filter(KnowledgeDoc.id == payload["id"]).one()
+        assert stored.status == "ready"
+        assert stored.chunk_count == 1
+        assert ("ocr", 18) in progress
+        assert ("ocr", 28) in progress
+        ocr.assert_called_once()
+        db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).delete()
+        db.delete(stored)
+        db.commit()
+    finally:
+        db.close()
+
+
 def test_retry_queues_a_failed_document_again(
     tmp_path: Path,
     authorization_context,

digital-avatar-app/backend/tests/test_pdf_ocr_service.py

@@ -0,0 +1,104 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from services.pdf_ocr_service import extract_scanned_pdf_text
+
+
+class FakePixmap:
+    def tobytes(self, *_args, **_kwargs):
+        return b"jpeg-page"
+
+
+class FakePage:
+    def get_pixmap(self, **_kwargs):
+        return FakePixmap()
+
+
+class FakeDocument:
+    page_count = 2
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_args):
+        return None
+
+    def load_page(self, _index):
+        return FakePage()
+
+
+def test_scanned_pdf_ocr_preserves_page_order_and_reports_progress(monkeypatch):
+    fake_pymupdf = SimpleNamespace(
+        open=lambda _path: FakeDocument(),
+        Matrix=lambda x, y: (x, y),
+        csRGB="rgb",
+    )
+    monkeypatch.setitem(__import__("sys").modules, "pymupdf", fake_pymupdf)
+    progress = []
+    reservation = SimpleNamespace()
+    config = SimpleNamespace(
+        api_key="configured",
+        ocr_model="qwen-vl-ocr",
+        vision_model="vision",
+        vision_max_tokens=2048,
+    )
+
+    with (
+        patch("services.pdf_ocr_service.get_chat_model_config", return_value=config),
+        patch("services.pdf_ocr_service.prepare_image", return_value=SimpleNamespace()),
+        patch(
+            "services.pdf_ocr_service.call_vision_model",
+            side_effect=[
+                {"content": "第一页文字", "usage": {"total_tokens": 10}},
+                {"content": "第二页文字", "usage": {"total_tokens": 12}},
+            ],
+        ),
+        patch("services.pdf_ocr_service.reserve_avatar_tokens", return_value=reservation) as reserve,
+        patch("services.pdf_ocr_service.settle_reservation") as settle,
+    ):
+        text = extract_scanned_pdf_text(
+            MagicMock(),
+            SimpleNamespace(id="avatar-1"),
+            "/tmp/scanned.pdf",
+            on_progress=lambda done, total: progress.append((done, total)),
+        )
+
+    assert text == "[第 1 页]\n第一页文字\n\n[第 2 页]\n第二页文字"
+    assert progress == [(1, 2), (2, 2)]
+    assert reserve.call_count == 2
+    assert settle.call_count == 2
+
+
+def test_scanned_pdf_ocr_releases_tokens_after_retries_fail(monkeypatch):
+    fake_document = FakeDocument()
+    fake_document.page_count = 1
+    fake_pymupdf = SimpleNamespace(
+        open=lambda _path: fake_document,
+        Matrix=lambda x, y: (x, y),
+        csRGB="rgb",
+    )
+    monkeypatch.setitem(__import__("sys").modules, "pymupdf", fake_pymupdf)
+    monkeypatch.setenv("KNOWLEDGE_PDF_OCR_ATTEMPTS", "2")
+    reservation = SimpleNamespace()
+    config = SimpleNamespace(
+        api_key="configured",
+        ocr_model="qwen-vl-ocr",
+        vision_model="vision",
+        vision_max_tokens=2048,
+    )
+
+    with (
+        patch("services.pdf_ocr_service.get_chat_model_config", return_value=config),
+        patch("services.pdf_ocr_service.prepare_image", return_value=SimpleNamespace()),
+        patch("services.pdf_ocr_service.call_vision_model", side_effect=RuntimeError("timeout")) as call,
+        patch("services.pdf_ocr_service.reserve_avatar_tokens", return_value=reservation),
+        patch("services.pdf_ocr_service.release_reservation") as release,
+        patch("services.pdf_ocr_service.time.sleep"),
+    ):
+        with pytest.raises(RuntimeError, match="第 1/1 页识别失败"):
+            extract_scanned_pdf_text(MagicMock(), SimpleNamespace(id="avatar-1"), "/tmp/scanned.pdf")
+
+    assert call.call_count == 2
+    release.assert_called_once()