fix(avatar): OCR image-only knowledge PDFs

2026-09-08 16:44:18 +08:00
parent 1bcdcead8d
commit 848657219e
6 changed files with 302 additions and 3 deletions
@@ -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,