diff --git a/docs/implementation/media/22-after.png b/docs/implementation/media/22-after.png new file mode 100644 index 0000000..e447e1b Binary files /dev/null and b/docs/implementation/media/22-after.png differ diff --git a/docs/implementation/media/22-before.png b/docs/implementation/media/22-before.png new file mode 100644 index 0000000..b307d55 Binary files /dev/null and b/docs/implementation/media/22-before.png differ diff --git a/frontend/src/pages/ImportDetail.vue b/frontend/src/pages/ImportDetail.vue index 73ace06..42952a8 100644 --- a/frontend/src/pages/ImportDetail.vue +++ b/frontend/src/pages/ImportDetail.vue @@ -127,6 +127,36 @@ function confirmReclassify() { }); } +// A killed or timed-out job leaves the import stopped with no job behind it. +const canResume = useCall({ + url: "/api/v2/method/wikify.api.imports.can_resume_import", + method: "GET", + immediate: false, +}); +const mayBeResumable = computed(() => ["Parsing", "Remediating", "Failed"].includes(status.value)); +const isStuck = computed(() => mayBeResumable.value && canResume.data === true); +watch( + status, + () => { + if (mayBeResumable.value) canResume.submit({ import_name: props.name }); + }, + { immediate: true } +); +const resume = useCall({ + url: "/api/v2/method/wikify.api.imports.resume_import", + method: "POST", + immediate: false, +}); +async function resumeImport() { + await resume.submit({ import_name: props.name }); + if (resume.error) { + toast.error(resume.error?.messages?.[0] || "Could not resume the import"); + await canResume.submit({ import_name: props.name }); + } else { + imp.reload(); + } +} + // Attach the document (with its project) as the agent's default context. Keyed on the // source_document id so reloads (progress ticks) don't reset a page/section selection. watch( @@ -309,8 +339,13 @@ const levelColor = { info: "text-ink-gray-7", warn: "text-ink-amber-6", error: " :route="{ name: 'ImportGraph', params: { name: props.name } }" /> None: @@ -27,14 +29,25 @@ def _create_import(pdf_file_url: str, title: str, project: str) -> str: imp.status = "Queued" imp.insert() + enqueue_parse(imp.name) + return imp.name + + +def enqueue_parse(import_name: str) -> None: frappe.enqueue( "wikify.jobs.parse.run", queue="long", timeout=IMPORT_JOB_TIMEOUT, - job_id=import_job_id(imp.name), - import_name=imp.name, + job_id=import_job_id(import_name), + deduplicate=True, + import_name=import_name, ) - return imp.name + + +def is_stuck(import_doc) -> bool: + if import_doc.status not in RESUMABLE_STATUSES or is_job_enqueued(import_job_id(import_doc.name)): + return False + return store.needs_resume(import_doc.source_document) @frappe.whitelist(methods=["POST"]) @@ -87,6 +100,28 @@ def trigger_remediation(import_name: str, scope: str = "flagged") -> str: return import_name +@frappe.whitelist() +def can_resume_import(import_name: str) -> bool: + import_doc = frappe.get_doc("Wikify Import", import_name) + return import_doc.has_permission("write") and is_stuck(import_doc) + + +@frappe.whitelist(methods=["POST"]) +def resume_import(import_name: str) -> str: + import_doc = frappe.get_doc("Wikify Import", import_name) + import_doc.check_permission("write") + if not is_stuck(import_doc): + frappe.throw( + _("Only a stopped parse or remediation can be resumed (current status: {0}).").format( + import_doc.status + ) + ) + + import_doc.db_set({"status": "Queued", "error": None}) + enqueue_parse(import_name) + return import_name + + @frappe.whitelist(methods=["POST"]) def reclassify(import_name: str) -> str: imp = frappe.get_doc("Wikify Import", import_name) diff --git a/wikify/engine/remediate.py b/wikify/engine/remediate.py index 9059306..9c943ac 100644 --- a/wikify/engine/remediate.py +++ b/wikify/engine/remediate.py @@ -4,6 +4,7 @@ from collections.abc import Callable import fitz +import frappe from wikify.engine import diagrams, llm, pdf_utils, regions, settings, store from wikify.engine.loader.cleanup_llm import clean_markdown @@ -16,6 +17,7 @@ ADOPTION_COMPOSITE_RATIO = 0.9 MIN_CANONICAL_CHARS = 40 +SAVE_BATCH_PAGES = 10 _IMAGE_EMBED_RE = re.compile(r"(!\[[^\]]*\]\()([^)]*)(\))") _MARKUP_RE = re.compile(r"[\s#*_>|`~\-]+") @@ -74,6 +76,7 @@ def remediate_pdf( # noqa: C901 progress_cb: Callable[[int, int], None] | None = None, page_cb: Callable[..., None] | None = None, stage_cb: Callable[[str], None] | None = None, + resume: bool = False, ) -> dict: if not llm.has_openrouter(): raise RuntimeError("OpenRouter key not set — remediation needs cloud models.") @@ -90,12 +93,21 @@ def remediate_pdf( # noqa: C901 canon_md = {p["page_no"]: p["baseline_markdown"] or "" for p in pages} canon_comp = {p["page_no"]: p["composite"] for p in pages} canon_src = {p["page_no"]: "baseline" for p in pages} + saved_page_nos = {p["page_no"] for p in targets if resume and p["canonical_source"]} + for p in targets: + if p["page_no"] in saved_page_nos: + canon_md[p["page_no"]] = p["canonical_markdown"] or "" + canon_comp[p["page_no"]] = p["canonical_composite"] + canon_src[p["page_no"]] = p["canonical_source"] + unsaved_pages: list[dict] = [] with fitz.open(pdf_path) as doc: furniture = det.find_furniture_lines([doc[p["page_no"] - 1].get_text("text") for p in pages]) doc_cost = 0.0 for i, p in enumerate(targets): + if p["page_no"] in saved_page_nos: + continue page = doc[p["page_no"] - 1] gt = page.get_text("text") kind = p["kind"] @@ -193,6 +205,16 @@ def remediate_pdf( # noqa: C901 if progress_cb: progress_cb(i + 1, total) + unsaved_pages.append(p) + if len(unsaved_pages) == SAVE_BATCH_PAGES: + with events.suspended_indexing(): + for saved in unsaved_pages: + pno = saved["page_no"] + store.set_canonical(saved["name"], canon_md[pno], canon_comp[pno], canon_src[pno]) + # nosemgrep + frappe.db.commit() + unsaved_pages = [] + stitched = dict(stitch_cross_page_tables([(p["page_no"], canon_md[p["page_no"]]) for p in pages])) with events.suspended_indexing(): for p in pages: diff --git a/wikify/engine/store.py b/wikify/engine/store.py index 36b28c2..400d766 100644 --- a/wikify/engine/store.py +++ b/wikify/engine/store.py @@ -105,7 +105,18 @@ def get_pages(source_document: str) -> list[dict]: return frappe.get_all( "Source Page", filters={"source_document": source_document}, - fields=["name", "page_no", "kind", "baseline_markdown", "verdict", "composite", "image"], + fields=[ + "name", + "page_no", + "kind", + "baseline_markdown", + "verdict", + "composite", + "image", + "canonical_markdown", + "canonical_composite", + "canonical_source", + ], order_by="page_no asc", ) @@ -140,6 +151,14 @@ def set_canonical(page_name: str, markdown: str, composite: float | None, source invalidate_page(page_name) +def needs_resume(source_document: str | None) -> bool: + return not source_document or bool( + frappe.db.exists( + "Source Page", {"source_document": source_document, "canonical_source": ("is", "not set")} + ) + ) + + def set_canonical_mean(source_document: str, mean: float | None) -> None: frappe.db.set_value("Source Document", source_document, "canonical_mean", mean) diff --git a/wikify/jobs/parse.py b/wikify/jobs/parse.py index 5f3db21..f5ea5b6 100644 --- a/wikify/jobs/parse.py +++ b/wikify/jobs/parse.py @@ -51,7 +51,10 @@ def page_cb(page_no, total, kind, score, metrics) -> None: # Parse + score only; the tree is built after remediation (below) so it's built once, # over the cleaned canonical markdown rather than the raw baseline. - source_document = parse_pdf( + resuming = bool(imp.source_document) + if not resuming: + delete_partial_documents(import_name) + source_document = imp.source_document or parse_pdf( pdf_path, title=imp.import_title, import_name=import_name, @@ -101,6 +104,7 @@ def rem_page_cb(page_no, total, method, adopted, base_c, new_c, metrics) -> None progress_cb=rem_progress_cb, page_cb=rem_page_cb, stage_cb=stage_cb, + resume=resuming, ) if result.get("cost"): log( @@ -139,3 +143,8 @@ def rem_page_cb(page_no, total, method, adopted, base_c, new_c, metrics) -> None frappe.db.commit() log(import_name, "error", "parse", "Parse failed — see error on the import") raise + + +def delete_partial_documents(import_name: str) -> None: + for name in frappe.get_all("Source Document", filters={"import": import_name}, pluck="name"): + frappe.delete_doc("Source Document", name, ignore_permissions=True, force=True) diff --git a/wikify/jobs/stalled.py b/wikify/jobs/stalled.py index 5037e41..2197379 100644 --- a/wikify/jobs/stalled.py +++ b/wikify/jobs/stalled.py @@ -4,6 +4,7 @@ from frappe.utils import add_to_date, now_datetime from frappe.utils.background_jobs import is_job_enqueued +from wikify.engine import store from wikify.jobs._util import import_job_id, log, publish_progress RUNNING_STATUSES = ("Queued", "Parsing", "Remediating") @@ -21,7 +22,11 @@ def fail_stalled_imports() -> None: for row in imports: if is_job_enqueued(import_job_id(row.name)): continue - retry_hint = "Click Remediate to retry." if row.source_document else "Import the PDF again to retry." + retry_hint = ( + "Click Resume import to continue." + if store.needs_resume(row.source_document) + else "Click Remediate to retry." + ) error = ( "The background job stopped before finishing (it timed out or the worker restarted). " + retry_hint diff --git a/wikify/tests/test_imports_api.py b/wikify/tests/test_imports_api.py index 6043258..6872961 100644 --- a/wikify/tests/test_imports_api.py +++ b/wikify/tests/test_imports_api.py @@ -8,6 +8,8 @@ from frappe.tests.utils import FrappeTestCase from wikify.api import imports as imports_api +from wikify.engine import store +from wikify.jobs import parse as parse_job from wikify.jobs._util import import_job_id from wikify.seed import seed_uncategorized_project @@ -226,3 +228,129 @@ def test_remediation_refuses_while_the_import_job_is_running(self): enqueue.assert_not_called() self.assertEqual(frappe.db.get_value("Wikify Import", name, "status"), "Failed") + + +class TestResumeImport(FrappeTestCase): + def make_import(self, status: str, canonical_sources: tuple[str, ...] | None = None): + imp = frappe.get_doc( + { + "doctype": "Wikify Import", + "import_title": "Resume Test", + "pdf": "/private/files/resume-test.pdf", + "project": seed_uncategorized_project(), + "status": status, + } + ).insert(ignore_permissions=True) + if canonical_sources is not None: + source_document = store.create_document("Resume Test", import_name=imp.name) + for page_no, canonical_source in enumerate(canonical_sources, start=1): + frappe.get_doc( + { + "doctype": "Source Page", + "source_document": source_document, + "page_no": page_no, + "kind": "text", + "canonical_source": canonical_source, + } + ).insert(ignore_permissions=True) + imp.db_set("source_document", source_document) + return imp + + def test_a_stuck_import_requeues_its_parse_job(self): + for status in ("Parsing", "Remediating"): + imp = self.make_import(status) + with ( + self.subTest(status=status), + patch.object(imports_api, "is_job_enqueued", return_value=False), + patch.object(frappe, "enqueue") as enqueue, + ): + self.assertTrue(imports_api.can_resume_import(imp.name)) + imports_api.resume_import(imp.name) + + enqueue.assert_called_once() + self.assertEqual(enqueue.call_args.args[0], "wikify.jobs.parse.run") + self.assertEqual(enqueue.call_args.kwargs["queue"], "long") + self.assertEqual(enqueue.call_args.kwargs["import_name"], imp.name) + self.assertEqual(enqueue.call_args.kwargs["job_id"], import_job_id(imp.name)) + + def test_refuses_while_the_import_job_is_queued_or_running(self): + imp = self.make_import("Remediating") + with ( + patch.object( + imports_api, "is_job_enqueued", side_effect=lambda job_id: job_id == import_job_id(imp.name) + ), + patch.object(frappe, "enqueue") as enqueue, + ): + self.assertFalse(imports_api.can_resume_import(imp.name)) + with self.assertRaises(frappe.ValidationError): + imports_api.resume_import(imp.name) + enqueue.assert_not_called() + + def test_refuses_an_import_that_is_not_parsing_or_remediating(self): + imp = self.make_import("Review") + with patch.object(frappe, "enqueue") as enqueue: + self.assertFalse(imports_api.can_resume_import(imp.name)) + with self.assertRaises(frappe.ValidationError): + imports_api.resume_import(imp.name) + enqueue.assert_not_called() + + def test_a_user_without_write_permission_cannot_resume(self): + imp = self.make_import("Remediating") + outsider = frappe.get_doc( + { + "doctype": "User", + "email": f"resume-outsider-{frappe.generate_hash(length=6)}@example.com", + "first_name": "Outsider", + } + ).insert(ignore_permissions=True) + + frappe.set_user(outsider.name) + try: + with patch.object(frappe, "enqueue") as enqueue: + self.assertFalse(imports_api.can_resume_import(imp.name)) + with self.assertRaises(frappe.PermissionError): + imports_api.resume_import(imp.name) + enqueue.assert_not_called() + finally: + frappe.set_user("Administrator") + + def test_a_failed_import_with_unremediated_pages_can_resume(self): + imp = self.make_import("Failed", canonical_sources=("vlm", "baseline", "")) + imp.db_set("error", "The background job stopped before finishing.") + with ( + patch.object(imports_api, "is_job_enqueued", return_value=False), + patch.object(frappe, "enqueue") as enqueue, + ): + self.assertTrue(imports_api.can_resume_import(imp.name)) + imports_api.resume_import(imp.name) + + enqueue.assert_called_once() + status, error = frappe.db.get_value("Wikify Import", imp.name, ["status", "error"]) + self.assertEqual(status, "Queued") + self.assertFalse(error) + + def test_an_import_with_every_page_remediated_cannot_resume(self): + for status in ("Remediating", "Failed"): + imp = self.make_import(status, canonical_sources=("vlm", "baseline", "cleanup")) + with ( + self.subTest(status=status), + patch.object(imports_api, "is_job_enqueued", return_value=False), + patch.object(frappe, "enqueue") as enqueue, + ): + self.assertFalse(imports_api.can_resume_import(imp.name)) + with self.assertRaises(frappe.ValidationError): + imports_api.resume_import(imp.name) + enqueue.assert_not_called() + + def test_a_failed_import_with_no_document_can_resume(self): + imp = self.make_import("Failed") + with patch.object(imports_api, "is_job_enqueued", return_value=False): + self.assertTrue(imports_api.can_resume_import(imp.name)) + + def test_reparse_deletes_the_half_parsed_document(self): + imp = self.make_import("Parsing") + half_parsed = store.create_document("Resume Test", import_name=imp.name) + + parse_job.delete_partial_documents(imp.name) + + self.assertFalse(frappe.db.exists("Source Document", half_parsed)) diff --git a/wikify/tests/test_remediate_pipeline.py b/wikify/tests/test_remediate_pipeline.py index 071e53b..1645eda 100644 --- a/wikify/tests/test_remediate_pipeline.py +++ b/wikify/tests/test_remediate_pipeline.py @@ -6,6 +6,7 @@ from pathlib import Path from unittest.mock import patch +import fitz import frappe from frappe.tests.utils import FrappeTestCase @@ -146,6 +147,59 @@ def test_every_page_gets_vlm_attempt_even_with_perfect_recall(self): total = frappe.db.count("Source Page", {"source_document": sd}) self.assertEqual(vlm_mock.call_count, total) + def test_killed_run_keeps_saved_batches_and_resume_skips_them(self): + path = Path(tempfile.mkdtemp()) / "long.pdf" + pdf = fitz.open() + for page_no in range(1, 26): + pdf.new_page().insert_text((72, 90), f"Chapter {page_no}\n\nBody text of page {page_no}.") + pdf.save(str(path)) + pdf.close() + with patch("wikify.engine.llm.has_openrouter", return_value=False): + sd = parse_pdf(str(path), title="Resume Test") + + class WorkerKilled(BaseException): + pass + + vlm_calls = [] + + def parse_page_image(data_url, **kwargs): + vlm_calls.append(data_url) + if len(vlm_calls) == 18: + raise WorkerKilled + return "vlm page" + + def canonical_sources(): + return frappe.get_all( + "Source Page", + filters={"source_document": sd}, + fields=["page_no", "canonical_source"], + order_by="page_no asc", + ) + + with ( + patch("frappe.db.commit") as commit_mock, + patch("wikify.engine.llm.has_openrouter", return_value=True), + patch("wikify.engine.llm.chat_completion", side_effect=_fake_chat), + patch( + "wikify.engine.remediate.clean_markdown", + side_effect=lambda md, model=None, project_context="", instruction="": md, + ), + patch("wikify.engine.remediate.vlm.parse_page_image", side_effect=parse_page_image), + ): + with self.assertRaises(WorkerKilled): + remediate_pdf(sd, str(path), scope="all") + + self.assertEqual(commit_mock.call_count, 1) + saved = [row.page_no for row in canonical_sources() if row.canonical_source] + self.assertEqual(saved, list(range(1, 11))) + + with patch("wikify.engine.remediate.vlm.parse_page_image", return_value="vlm page") as vlm_mock: + result = remediate_pdf(sd, str(path), scope="all", resume=True) + + self.assertEqual(vlm_mock.call_count, 15) + self.assertEqual(result["targets"], 25) + self.assertTrue(all(row.canonical_source for row in canonical_sources())) + def test_vlm_failure_falls_back_to_cleanup_with_note(self): sd, path = self._parse() with ( diff --git a/wikify/tests/test_stalled_imports.py b/wikify/tests/test_stalled_imports.py index 13d3da4..4a4c1b7 100644 --- a/wikify/tests/test_stalled_imports.py +++ b/wikify/tests/test_stalled_imports.py @@ -11,10 +11,19 @@ class TestStalledImports(FrappeTestCase): - def make_import(self, status: str, minutes_since_modified: int): + def make_import(self, status: str, minutes_since_modified: int, unremediated_page: bool = False): source_document = frappe.get_doc({"doctype": "Source Document", "title": "Stalled Test"}).insert( ignore_permissions=True ) + if unremediated_page: + frappe.get_doc( + { + "doctype": "Source Page", + "source_document": source_document.name, + "page_no": 1, + "kind": "text", + } + ).insert(ignore_permissions=True) imp = frappe.get_doc( { "doctype": "Wikify Import", @@ -52,6 +61,15 @@ def test_an_import_with_no_live_job_is_marked_failed(self): self.assertIn("Click Remediate to retry", error) self.assertTrue(frappe.db.exists("Import Log Entry", {"import": name, "level": "error"})) + def test_an_import_with_unremediated_pages_points_to_resume(self): + name = self.make_import("Remediating", 30, unremediated_page=True) + + self.fail_stalled_imports(job_is_enqueued=False) + + status, error = frappe.db.get_value("Wikify Import", name, ["status", "error"]) + self.assertEqual(status, "Failed") + self.assertIn("Click Resume import to continue", error) + def test_an_import_with_a_live_job_is_left_running(self): name = self.make_import("Remediating", 30)