diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index d1397bc3..3d24a6b5 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -445,6 +445,7 @@ def _build_page_dataframe( vlm_model=vlm_model, page_assets_by_page=page_assets_by_page, node_summary_max_pages=page_memory_config.node_summary_max_pages, + max_concurrent=page_memory_config.node_assembly_concurrency, ) logger.info( "[page_memory] C7 assembled {} node rows (verdict={})", @@ -649,6 +650,7 @@ def _run_hierarchy_scope( fat_leaf_pages=fat_leaf_pages, budget=None, vlm_model=vlm_model, + max_concurrent=page_memory_config.title_detection_concurrency, ) _record_trace_stage( trace_recorder, diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index ff4191a2..855ce3c6 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -242,6 +242,9 @@ def pages_by_leaf_count(views: list[NodePageView]) -> dict[int, list[LeafNode]]: # ── VLM-backed helpers ─────────────────────────────────────────────── +_NODE_VLM_RETRY_DELAY_SECONDS = 5.0 + + def resolve_page_text( *, page: int, @@ -254,18 +257,33 @@ def resolve_page_text( Electronic PDFs already have PyMuPDF text; scanned pages have (near) empty text and fall back to the shared ``transcribe()`` OCR primitive (§4.2). + Retries once on an empty transcription before giving up. """ text = (raw_text or "").strip() if text: return text if not vlm_model or not image_path or not os.path.exists(image_path): return "" - return transcribe( - image_paths=[image_path], - model=vlm_model, - max_tokens=1500, - usage_task="page_memory.node_ocr", - ) + + import gevent + + for attempt in range(2): + transcribed = transcribe( + image_paths=[image_path], + model=vlm_model, + max_tokens=1500, + usage_task="page_memory.node_ocr", + ) + if transcribed: + return transcribed + if attempt == 0: + logger.warning( + "[node_assembler] empty OCR transcription for page {}, retrying after {}s", + page, + _NODE_VLM_RETRY_DELAY_SECONDS, + ) + gevent.sleep(_NODE_VLM_RETRY_DELAY_SECONDS) + return "" def compute_node_summary( @@ -378,26 +396,36 @@ def _vlm_node_summary( if not image_paths: return None - result = summarize( - mode="page", - image_paths=image_paths, - model=vlm_model, - usage_task="page_memory.node_summary", - prompt_task="page-memory-node-summary", - prompt_paras={ - "max_tokens": 400, - "node_title": leaf.title, - "next_title": next_title or "", - "kw_num": 5, - }, - ) - if not result.summary and not result.entities: - return None - return ( - result.summary, - [e.text for e in result.entities], - [e.to_dict() for e in result.entities], - ) + import gevent + + for attempt in range(2): + result = summarize( + mode="page", + image_paths=image_paths, + model=vlm_model, + usage_task="page_memory.node_summary", + prompt_task="page-memory-node-summary", + prompt_paras={ + "max_tokens": 400, + "node_title": leaf.title, + "next_title": next_title or "", + "kw_num": 5, + }, + ) + if result.summary or result.entities: + return ( + result.summary, + [e.text for e in result.entities], + [e.to_dict() for e in result.entities], + ) + if attempt == 0: + logger.warning( + "[node_assembler] empty VLM node summary for {}, retrying after {}s", + leaf.section_path, + _NODE_VLM_RETRY_DELAY_SECONDS, + ) + gevent.sleep(_NODE_VLM_RETRY_DELAY_SECONDS) + return None # ── Orchestration ──────────────────────────────────────────────────── @@ -416,41 +444,85 @@ def build_node_rows( vlm_model: str | None = None, page_assets_by_page: dict[int, list[PageAsset]] | None = None, node_summary_max_pages: int = _NODE_SUMMARY_MAX_PAGES_DEFAULT, + max_concurrent: int | None = None, ) -> list[dict[str, Any]]: """Assemble one row per leaf section node (node-granularity chunks).""" + import gevent + from gevent.pool import Pool as GeventPool + + from shared.core.config import settings + available_pages = set(raw_text_by_page.keys()) leaves = identify_leaf_nodes(skeletons) views, page_owner = assign_pages_to_leaves(leaves, available_pages=available_pages) page_to_leaves = pages_by_leaf_count(views) + resolved_max_concurrent = max_concurrent or int( + getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 4) + ) + # Resolve body text once per owned page (PyMuPDF, OCR fallback for scanned). + # Loop 2 (node summaries) reads this map via build_node_content, so it + # must run only after this pool has fully joined. + owned_pages = [page for view in views for page in view.owned_pages] + + def _resolve_one(page: int) -> tuple[int, str]: + text = resolve_page_text( + page=page, + raw_text=raw_text_by_page.get(page, ""), + image_path=image_path_by_page.get(page), + vlm_model=vlm_model, + ) + return page, text + resolved_text: dict[int, str] = {} - for view in views: - for page in view.owned_pages: - resolved_text[page] = resolve_page_text( - page=page, - raw_text=raw_text_by_page.get(page, ""), - image_path=image_path_by_page.get(page), - vlm_model=vlm_model, - ) + if owned_pages: + pool = GeventPool(size=min(resolved_max_concurrent, len(owned_pages))) + greenlets = [pool.spawn(_resolve_one, page) for page in owned_pages] + gevent.joinall(greenlets) + for greenlet in greenlets: + if greenlet.value is None: + continue + page, text = greenlet.value + resolved_text[page] = text + + # Node summaries: computed concurrently, but merged into `summaries` by + # index so row order below matches `views` order regardless of the + # order greenlets complete in. + def _summarize_one(index: int) -> tuple[int, tuple[str, list[str], list[dict[str, str]]]]: + view = views[index] + return index, compute_node_summary( + view=view, + page_to_leaves=page_to_leaves, + tag_by_page=tag_by_page, + image_path_by_page=image_path_by_page, + vlm_model=vlm_model, + node_summary_max_pages=node_summary_max_pages, + ) + + summaries: list[tuple[str, list[str], list[dict[str, str]]]] = [ + ("", [], []) for _ in views + ] + if views: + pool = GeventPool(size=min(resolved_max_concurrent, len(views))) + greenlets = [pool.spawn(_summarize_one, index) for index in range(len(views))] + gevent.joinall(greenlets) + for greenlet in greenlets: + if greenlet.value is None: + continue + index, result = greenlet.value + summaries[index] = result rows: list[dict[str, Any]] = [] rows_by_path: dict[str, dict[str, Any]] = {} - for view in views: + for index, view in enumerate(views): leaf = view.leaf content = build_node_content( view, page_owner=page_owner, page_text=resolved_text, ) - summary, keywords, entities = compute_node_summary( - view=view, - page_to_leaves=page_to_leaves, - tag_by_page=tag_by_page, - image_path_by_page=image_path_by_page, - vlm_model=vlm_model, - node_summary_max_pages=node_summary_max_pages, - ) + summary, keywords, entities = summaries[index] know_id = f"node_{gen_str_codes(f'{filename}::{leaf.section_path}')}" row = { "content": content, diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index fcee9da5..56ac55f3 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -262,6 +262,7 @@ def tag_page_titles( fat_leaf_pages: set[int], budget: Any | None = None, vlm_model: str | None = None, + max_concurrent: int | None = None, ) -> list[PageTagResult]: """Run independent VLM title detection on fat-leaf pages. @@ -278,6 +279,8 @@ def tag_page_titles( Deprecated, ignored. Kept for call-site compatibility. vlm_model: VLM model name; falls back to ``$IMAGE_MODEL``. + max_concurrent: + Maximum concurrent title-detection VLM calls. Returns ------- @@ -292,31 +295,52 @@ def tag_page_titles( logger.warning("[page_tagger] no VLM model for title detection; skipping") return tag_results + import gevent + from gevent.pool import Pool as GeventPool + + from shared.core.config import settings + tag_map = {t.page_index: t for t in tag_results} page_map = {p.page_index: p for p in pages} - vlm_calls = 0 - titles_found = 0 - for page_idx in sorted(fat_leaf_pages): - page = page_map.get(page_idx) - tag = tag_map.get(page_idx) - if page is None or tag is None: - continue + page_indices = [ + page_idx + for page_idx in sorted(fat_leaf_pages) + if (page := page_map.get(page_idx)) is not None + and page_idx in tag_map + and page.image_path + and os.path.exists(page.image_path) + ] - # Skip pages without images (text_only / skip) - if not page.image_path or not os.path.exists(page.image_path): - continue + resolved_max_concurrent = max_concurrent or int( + getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 4) + ) + + def _detect_one(page_idx: int) -> tuple[int, list[dict[str, Any]]]: + page = page_map[page_idx] + return page_idx, _tag_vlm_titles(page, model=model) - observed = _tag_vlm_titles(page, model=model) - tag.observed_titles = observed + pool = GeventPool(size=min(resolved_max_concurrent, len(page_indices))) + greenlets = [pool.spawn(_detect_one, page_idx) for page_idx in page_indices] + gevent.joinall(greenlets) + + vlm_calls = 0 + titles_found = 0 + for greenlet in greenlets: + if greenlet.value is None: + continue + page_idx, observed = greenlet.value + tag_map[page_idx].observed_titles = observed vlm_calls += 1 titles_found += len(observed) logger.info( - "[page_tagger] title detection: {} VLM calls on {} fat-leaf pages, {} titles found", + "[page_tagger] title detection: {} VLM calls on {} fat-leaf pages, " + "{} titles found, concurrency={}", vlm_calls, len(fat_leaf_pages), titles_found, + resolved_max_concurrent, ) return tag_results diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index c66fc217..cc75df76 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -9,8 +9,13 @@ os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") os.environ.setdefault("S3_TEMP_PATH", "/tmp") +import gevent + +import app.services.page_memory.node_assembler as node_assembler from app.services.page_memory.node_assembler import ( SAME_AS_PREFIX, + LeafNode, + NodePageView, assign_pages_to_leaves, build_node_content, build_node_rows, @@ -20,6 +25,7 @@ from app.services.page_memory.page_assets import PageAsset from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton +from shared.services.ai.summary.model import BodySummary from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks import pandas as pd @@ -389,3 +395,177 @@ def test_page_connectto_normalizes_to_asset_chunk_id() -> None: "ref": "[tables/table_page_1_1.html]", } ] + + +def _many_leaf_skeletons(count: int) -> list[SectionSkeleton]: + """One single-page leaf section per page, 1..count, all siblings.""" + return [ + SectionSkeleton( + section_path=f"demo.pdf/Section {page}", + level=1, + start_page=page, + end_page=page, + title=f"Section {page}", + parent_path="demo.pdf", + ) + for page in range(1, count + 1) + ] + + +def test_build_node_rows_preserves_order_under_concurrent_vlm_summaries( + tmp_path, +) -> None: + """Node summaries are computed concurrently but must land in view order.""" + page_count = 8 + skeletons = _many_leaf_skeletons(page_count) + raw_text_by_page = {page: "" for page in range(1, page_count + 1)} + image_path_by_page: dict[int, str] = {} + for page in range(1, page_count + 1): + img = tmp_path / f"page-{page}.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n fake") + image_path_by_page[page] = str(img) + + # Every leaf here is single-page and unshared, so compute_node_summary + # short-circuits to the per-page tag (no VLM call needed) -- this test + # only needs to prove ordering survives the pooled merge, which applies + # regardless of which code path filled `summaries[index]`. + tag_by_page = { + page: PageTagResult( + page_index=page, summary=f"summary-{page}", keywords=[f"kw-{page}"] + ) + for page in range(1, page_count + 1) + } + + rows = build_node_rows( + skeletons=skeletons, + raw_text_by_page=raw_text_by_page, + image_path_by_page=image_path_by_page, + kind_by_page={}, + tag_by_page=tag_by_page, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + max_concurrent=4, + ) + + assert [r["path"] for r in rows] == [ + f"demo.pdf/Section {page}" for page in range(1, page_count + 1) + ] + assert [r["summary"] for r in rows] == [ + f"summary-{page}" for page in range(1, page_count + 1) + ] + + +def test_vlm_node_summary_retries_once_on_empty_result(monkeypatch, tmp_path) -> None: + """A first empty VLM response is retried once before giving up.""" + calls: list[int] = [] + + def _fake_summarize(**kwargs): + calls.append(1) + if len(calls) == 1: + return BodySummary(summary="", entities=[], kind="page") + return BodySummary(summary="retried summary", entities=[], kind="page") + + monkeypatch.setattr(node_assembler, "summarize", _fake_summarize) + monkeypatch.setattr(gevent, "sleep", lambda _seconds: None) + + img1 = tmp_path / "page-1.png" + img1.write_bytes(b"\x89PNG\r\n\x1a\n fake") + img2 = tmp_path / "page-2.png" + img2.write_bytes(b"\x89PNG\r\n\x1a\n fake") + + leaf = LeafNode( + section_path="demo.pdf/Section 1", + title="Section 1", + level=1, + start_page=1, + end_page=2, + ) + view = NodePageView(leaf=leaf, pages=[1, 2], owned_pages=[1, 2]) + + result = node_assembler._vlm_node_summary( + view=view, + page_to_leaves={1: [leaf], 2: [leaf]}, + image_path_by_page={1: str(img1), 2: str(img2)}, + vlm_model="fake-vlm", + node_summary_max_pages=5, + ) + + assert len(calls) == 2 + assert result is not None + summary, _keywords, _entities = result + assert summary == "retried summary" + + +def test_vlm_node_summary_returns_none_after_retry_exhausted(monkeypatch, tmp_path) -> None: + """Two consecutive empty VLM responses fall through to None (caller degrades).""" + monkeypatch.setattr( + node_assembler, + "summarize", + lambda **kwargs: BodySummary(summary="", entities=[], kind="page"), + ) + monkeypatch.setattr(gevent, "sleep", lambda _seconds: None) + + img = tmp_path / "page-1.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n fake") + + leaf = LeafNode( + section_path="demo.pdf/Section 1", + title="Section 1", + level=1, + start_page=1, + end_page=1, + ) + view = NodePageView(leaf=leaf, pages=[1], owned_pages=[1]) + + result = node_assembler._vlm_node_summary( + view=view, + page_to_leaves={1: [leaf]}, + image_path_by_page={1: str(img)}, + vlm_model="fake-vlm", + node_summary_max_pages=5, + ) + + assert result is None + + +def test_resolve_page_text_retries_once_on_empty_transcription(monkeypatch, tmp_path) -> None: + calls: list[int] = [] + + def _fake_transcribe(**kwargs): + calls.append(1) + return "" if len(calls) == 1 else "transcribed text" + + monkeypatch.setattr(node_assembler, "transcribe", _fake_transcribe) + monkeypatch.setattr(gevent, "sleep", lambda _seconds: None) + + img = tmp_path / "page-1.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n fake") + + text = node_assembler.resolve_page_text( + page=1, + raw_text="", + image_path=str(img), + vlm_model="fake-vlm", + ) + + assert len(calls) == 2 + assert text == "transcribed text" + + +def test_resolve_page_text_empty_after_retry_exhausted(monkeypatch, tmp_path) -> None: + monkeypatch.setattr(node_assembler, "transcribe", lambda **kwargs: "") + monkeypatch.setattr(gevent, "sleep", lambda _seconds: None) + + img = tmp_path / "page-1.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n fake") + + text = node_assembler.resolve_page_text( + page=1, + raw_text="", + image_path=str(img), + vlm_model="fake-vlm", + ) + + assert text == "" diff --git a/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py b/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py new file mode 100644 index 00000000..b278a620 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +import app.services.page_memory.page_tagger as page_tagger +from app.services.page_memory.page_renderer import PageRenderResult +from app.services.page_memory.page_tagger import PageTagResult, tag_page_titles + + +def _render_result(page_index: int, image_path: str) -> PageRenderResult: + return PageRenderResult( + page_index=page_index, + image_path=image_path, + raw_text="", + width=612.0, + height=792.0, + is_landscape=False, + ) + + +def test_tag_page_titles_preserves_order_under_concurrency(monkeypatch, tmp_path) -> None: + """Title detection runs concurrently but results land on the right page's tag.""" + page_count = 8 + pages: list[PageRenderResult] = [] + tag_results: list[PageTagResult] = [] + for page in range(1, page_count + 1): + img = tmp_path / f"page-{page}.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n fake") + pages.append(_render_result(page, str(img))) + tag_results.append(PageTagResult(page_index=page)) + + def _fake_tag_vlm_titles(page: PageRenderResult, *, model: str) -> list[dict]: + return [{"text": f"Title {page.page_index}", "prominence": 0.9}] + + monkeypatch.setattr(page_tagger, "_tag_vlm_titles", _fake_tag_vlm_titles) + + result = tag_page_titles( + pages=pages, + tag_results=tag_results, + fat_leaf_pages=set(range(1, page_count + 1)), + budget=None, + vlm_model="fake-vlm", + max_concurrent=4, + ) + + for page in range(1, page_count + 1): + tag = next(t for t in result if t.page_index == page) + assert tag.observed_titles == [{"text": f"Title {page}", "prominence": 0.9}] + + +def test_tag_page_titles_skips_pages_without_image(monkeypatch, tmp_path) -> None: + img = tmp_path / "page-1.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n fake") + + pages = [ + _render_result(1, str(img)), + _render_result(2, str(tmp_path / "missing.png")), + ] + tag_results = [ + PageTagResult(page_index=1), + PageTagResult(page_index=2), + ] + + calls: list[int] = [] + + def _fake_tag_vlm_titles(page: PageRenderResult, *, model: str) -> list[dict]: + calls.append(page.page_index) + return [{"text": "Title", "prominence": 0.9}] + + monkeypatch.setattr(page_tagger, "_tag_vlm_titles", _fake_tag_vlm_titles) + + result = tag_page_titles( + pages=pages, + tag_results=tag_results, + fat_leaf_pages={1, 2}, + budget=None, + vlm_model="fake-vlm", + max_concurrent=4, + ) + + assert calls == [1] + tag_1 = next(t for t in result if t.page_index == 1) + tag_2 = next(t for t in result if t.page_index == 2) + assert tag_1.observed_titles == [{"text": "Title", "prominence": 0.9}] + assert tag_2.observed_titles == [] + + +def test_tag_page_titles_returns_unchanged_when_no_fat_leaf_pages() -> None: + tag_results = [PageTagResult(page_index=1)] + result = tag_page_titles( + pages=[], + tag_results=tag_results, + fat_leaf_pages=set(), + budget=None, + vlm_model="fake-vlm", + ) + assert result is tag_results diff --git a/packages/shared-python/shared/models/schemas/page_memory_config.py b/packages/shared-python/shared/models/schemas/page_memory_config.py index 17944439..c8a4747d 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -13,6 +13,8 @@ class PageMemoryConfig: max_pages: int = 1500 scope_concurrency: int = 4 tag_concurrency: int = 4 + title_detection_concurrency: int = 4 + node_assembly_concurrency: int = 4 tag_mode: Literal["vlm", "text"] = "vlm" fine_min_pages: int = 4 hierarchy_model: str | None = None @@ -57,6 +59,14 @@ def from_mapping(cls, value: object) -> Self: value.get("tag_concurrency"), default.tag_concurrency, ), + title_detection_concurrency=_as_int( + value.get("title_detection_concurrency"), + default.title_detection_concurrency, + ), + node_assembly_concurrency=_as_int( + value.get("node_assembly_concurrency"), + default.node_assembly_concurrency, + ), tag_mode=resolved_tag_mode, fine_min_pages=_as_int( value.get("fine_min_pages"),