From 6773fe3ffc432670480df93c7e2e6f504bde0008 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Mon, 22 Jun 2026 17:26:24 +0800 Subject: [PATCH 01/27] feat: implement node assembler for page memory hierarchy to support shared page content and VLM-based node summarization --- .../connect_builder/summary_builder.py | 11 +- .../success_finalization.py | 9 + .../services/page_memory/fine_hierarchy.py | 392 +++++++++++++ .../services/page_memory/memory_service.py | 235 ++++++-- .../services/page_memory/node_assembler.py | 526 ++++++++++++++++++ .../app/services/page_memory/page_tagger.py | 243 +++++++- .../page_memory/skeleton_extractor.py | 43 +- .../test_document_agent_budget_contract.py | 44 +- ...est_page_memory_fine_hierarchy_contract.py | 183 ++++++ .../test_page_memory_navigation_contract.py | 144 +++++ ...est_page_memory_node_assembler_contract.py | 167 ++++++ .../shared/services/ai/prompt_service.py | 365 +++++++----- .../services/storage/zip_chunk_schema.py | 31 +- .../services/storage/zip_doc_navigation.py | 184 ++++++ .../services/storage/zip_result_schema.py | 17 +- .../services/storage/zip_result_service.py | 25 +- 16 files changed, 2376 insertions(+), 243 deletions(-) create mode 100644 apps/worker/app/services/page_memory/fine_hierarchy.py create mode 100644 apps/worker/app/services/page_memory/node_assembler.py create mode 100644 apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py create mode 100644 apps/worker/tests/contract/test_page_memory_navigation_contract.py create mode 100644 apps/worker/tests/contract/test_page_memory_node_assembler_contract.py diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 156cde97..a55c5520 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -114,6 +114,7 @@ def ensure_doc_nav_json( source_file_name: str = "", *, overwrite: bool = False, + skeletons: List[Dict[str, Any]] | None = None, ) -> str: """Materialize ``doc_nav.json`` from chunks when the parser did not emit one.""" nav_path = os.path.join(file_dir, DOC_NAV_FILENAME) @@ -122,7 +123,15 @@ def ensure_doc_nav_json( from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder - doc_nav = ZipResultSchemaBuilder().build_doc_nav(chunks, source_file_name) + schema = ZipResultSchemaBuilder() + if skeletons: + doc_nav = schema.build_doc_nav_from_skeletons( + skeletons, + chunks, + source_file_name, + ) + else: + doc_nav = schema.build_doc_nav(chunks, source_file_name) _save_doc_nav(file_dir, doc_nav) return nav_path diff --git a/apps/worker/app/services/document_ingestion/success_finalization.py b/apps/worker/app/services/document_ingestion/success_finalization.py index 93c5c301..91f9691d 100644 --- a/apps/worker/app/services/document_ingestion/success_finalization.py +++ b/apps/worker/app/services/document_ingestion/success_finalization.py @@ -138,12 +138,14 @@ def _enrich_document_navigation( section_summaries: dict[str, str] = {} add_dir = artifact.add_dir parsed_contents_df = artifact.dataframe + skeletons = _get_page_memory_skeletons(parsed_contents_df) if add_dir and source_file_name: if "path" in parsed_contents_df.columns: ensure_doc_nav_json( str(add_dir), chunks, source_file_name=source_file_name, + skeletons=skeletons, ) try: document_root_for_enrich = os.path.dirname(str(add_dir)) @@ -194,6 +196,13 @@ def _attach_document_top_summary( metadata["document_top_summary"] = document_top_summary +def _get_page_memory_skeletons(parsed_df: Any) -> list[dict[str, Any]]: + raw_skeletons = getattr(parsed_df, "attrs", {}).get("page_memory_skeletons") + if not isinstance(raw_skeletons, list): + return [] + return [skel for skel in raw_skeletons if isinstance(skel, dict)] + + def _record_processing_completion( *, job_id: str, diff --git a/apps/worker/app/services/page_memory/fine_hierarchy.py b/apps/worker/app/services/page_memory/fine_hierarchy.py new file mode 100644 index 00000000..af2a05cf --- /dev/null +++ b/apps/worker/app/services/page_memory/fine_hierarchy.py @@ -0,0 +1,392 @@ +"""Fine-grained hierarchy builder for page-memory native hierarchy (Step 3). + +Page-memory candidates are already confirmed visual headings from PDF/PPT pages. +They must not run through the chunk-track heading executor, whose compact/body +row semantics are designed for raw text. This module calls a dedicated +``page-memory-hierarchy`` prompt and builds deeper ``SectionSkeleton`` entries +under each coarse TOC leaf. +""" + +from __future__ import annotations + +import json +import os +import re +from pathlib import Path +from typing import Any + +from loguru import logger + +from app.services.page_memory.page_tagger import PageTagResult +from app.services.page_memory.skeleton_extractor import SectionSkeleton +from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.services.ai.prompt_service import build_prompt +from shared.services.ai.response_process_service import eval_response + + +def refine_fat_leaf_skeletons( + *, + coarse_skeletons: list[SectionSkeleton], + tag_results: list[PageTagResult], + fat_leaf_pages: set[int], + model_name: str | None = None, + output_dir: str | None = None, +) -> list[SectionSkeleton]: + """Refine coarse TOC leaf skeletons using VLM-observed title candidates. + + For each coarse skeleton that overlaps fat-leaf pages: + 1. Collect real ``observed_titles`` from that leaf using exclusive end + boundaries. + 2. Ask the page-memory hierarchy prompt to assign relative levels. + 3. Rebuild the nested tree and graft it under the coarse leaf. + + Thin leaves (not in ``fat_leaf_pages``) are passed through unchanged. + """ + if not fat_leaf_pages: + return coarse_skeletons + + refined: list[SectionSkeleton] = [] + + for index, skeleton in enumerate(coarse_skeletons): + exclusive_end = _exclusive_end(coarse_skeletons, index) + skel_pages = set(range(skeleton.start_page, exclusive_end + 1)) + overlap = skel_pages & fat_leaf_pages + + if not overlap: + refined.append(skeleton) + continue + + # Collect observed titles from tagged pages within this skeleton + candidates = _collect_candidates( + skeleton=skeleton, + exclusive_end=exclusive_end, + tag_results=tag_results, + fat_leaf_pages=fat_leaf_pages, + ) + + if not candidates: + refined.append(skeleton) + continue + + # Run hierarchy LLM and graft the resulting tree under the coarse leaf + deeper = _run_hierarchy_on_candidates( + candidates=candidates, + skeleton=skeleton, + model_name=model_name, + output_dir=output_dir, + ) + + if deeper: + refined.extend(deeper) + else: + refined.append(skeleton) + + logger.info( + "[page_memory.fine_hierarchy] refined {} → {} skeletons ({} fat-leaf pages)", + len(coarse_skeletons), + len(refined), + len(fat_leaf_pages), + ) + return refined + + +def compute_fat_leaf_pages( + skeletons: list[SectionSkeleton], + min_pages: int, +) -> set[int]: + """Compute the set of page indices belonging to fat-leaf sections. + + A "fat leaf" is a coarse skeleton (TOC leaf node) whose page range + exceeds ``min_pages``. Only these pages get VLM title detection. + + Uses exclusive end boundaries: each skeleton's scan range ends at + ``next_skeleton.start_page - 1`` to avoid scanning pages that belong + to the next sibling skeleton (the raw ``end_page`` from the hierarchy + locator uses closed-closed intervals that can overlap at boundaries). + """ + fat_pages: set[int] = set() + for idx, skel in enumerate(skeletons): + exclusive_end = _exclusive_end(skeletons, idx) + page_span = exclusive_end - skel.start_page + 1 + if page_span > min_pages: + fat_pages.update(range(skel.start_page, exclusive_end + 1)) + return fat_pages + + +# ── Internal helpers ───────────────────────────────────────────────── + + +def _collect_candidates( + *, + skeleton: SectionSkeleton, + exclusive_end: int, + tag_results: list[PageTagResult], + fat_leaf_pages: set[int], +) -> list[dict[str, Any]]: + """Gather VLM-observed titles within a skeleton's page range. + + Returns list of ``{id, heading, page, prominence}`` sorted by page then + prominence (strongest first within page). + """ + candidates: list[dict[str, Any]] = [] + tag_by_page = {t.page_index: t for t in tag_results} + candidate_id = 0 + seen: set[str] = set() + parent_key = _title_key(skeleton.title) + + for page in range(skeleton.start_page, exclusive_end + 1): + if page not in fat_leaf_pages: + continue + tag = tag_by_page.get(page) + if tag is None or not tag.observed_titles: + continue + + # Sort by prominence descending within page + sorted_titles = sorted( + tag.observed_titles, + key=lambda t: -(t.get("prominence") or 0.5), + ) + for title_entry in sorted_titles: + text = title_entry.get("text", "").strip() + if not text or len(text) < 2: + continue + key = _title_key(text) + if not key or key == parent_key or key in seen: + continue + seen.add(key) + candidate_id += 1 + candidates.append({ + "id": candidate_id, + "heading": text, + "page": page, + "prominence": title_entry.get("prominence"), + }) + + return candidates + + +def _run_hierarchy_on_candidates( + *, + candidates: list[dict[str, Any]], + skeleton: SectionSkeleton, + model_name: str | None, + output_dir: str | None, +) -> list[SectionSkeleton] | None: + """Run the page-memory hierarchy prompt and rebuild a nested skeleton tree. + + The LLM assigns *relative* levels (starting at 1) within this leaf. We + reconstruct parent/child nesting from those levels with a stack, so the + fine structure (e.g. L1→L2→L3 inside the leaf) is preserved instead of + being flattened, then offset every level by ``skeleton.level`` to make it + absolute under the coarse leaf. + """ + if not candidates: + return None + + max_depth = int(os.environ.get("KB_LLM_HEADING_MAX_DEPTH", "6")) + max_tokens = int(os.environ.get("PAGE_MEMORY_HIERARCHY_MAX_TOKENS", "2000")) + input_json = json.dumps( + [ + { + "id": cand["id"], + "page": cand["page"], + "prominence": cand.get("prominence"), + "heading": cand["heading"], + } + for cand in candidates + ], + ensure_ascii=False, + indent=2, + ) + coarse_context = ( + f"title={skeleton.title}\n" + f"path={skeleton.section_path}\n" + f"pages={skeleton.start_page}-{skeleton.end_page}" + ) + + try: + prompt, temperature, _top_p, prompt_max_tokens = build_prompt( + "page-memory-hierarchy", + input_json, + "", + paras={ + "max_depth": max_depth, + "max_tokens": max_tokens, + "coarse_context": coarse_context, + }, + ) + resolved_model = ( + model_name + or os.environ.get("PAGE_MEMORY_HIERARCHY_MODEL") + or os.environ.get( + "HIERARCHY_LLM_MODEL", + os.environ.get("NORMOL_MODEL"), + ) + ) + answer = get_openai_client(model=resolved_model).chat_completion( + messages=[ + {"role": "system", "content": "you are a document structure expert"}, + {"role": "user", "content": prompt}, + ], + model=resolved_model, + max_tokens=prompt_max_tokens, + temperature=temperature, + usage_task="page_memory.hierarchy", + ) + result = eval_response(answer) + except Exception as exc: + logger.warning( + "[page_memory.fine_hierarchy] LLM failed for skeleton {}: {}", + skeleton.section_path, + exc, + ) + return None + + levels_by_id = _parse_hierarchy_result(result, max_depth=max_depth) + if not levels_by_id: + logger.warning( + "[page_memory.fine_hierarchy] empty hierarchy result for skeleton {}", + skeleton.section_path, + ) + return None + + # Keep candidate order (== id order == the sequence fed to the LLM). + ordered: list[dict[str, Any]] = [] + candidate_by_id = {int(c["id"]): c for c in candidates} + for row_id in sorted(levels_by_id): + cand = candidate_by_id.get(row_id) + if cand is None: + continue + rel_level = levels_by_id[row_id] + if rel_level <= 0: + continue + ordered.append({ + "id": row_id, + "rel_level": rel_level, + "heading": str(cand["heading"]), + "page": int(cand.get("page") or skeleton.start_page), + }) + ordered.sort(key=lambda item: item["id"]) + + if not ordered: + return None + + # Stack-based nesting: a node hangs under the nearest preceding node whose + # relative level is shallower (smaller rel_level). + nodes: list[dict[str, Any]] = [] + stack: list[tuple[int, str]] = [] # (rel_level, section_path) + for item in ordered: + rel_level = item["rel_level"] + heading = item["heading"] + while stack and stack[-1][0] >= rel_level: + stack.pop() + parent_path = stack[-1][1] if stack else skeleton.section_path + section_path = f"{parent_path}/{heading}" + stack.append((rel_level, section_path)) + nodes.append({ + "section_path": section_path, + "parent_path": parent_path, + "rel_level": rel_level, + "abs_level": skeleton.level + rel_level, + "title": heading, + "start_page": item["page"], + }) + + # end_page: a section runs until the next section at the same-or-shallower + # relative level starts; descendants are nested inside that span. + deeper: list[SectionSkeleton] = [] + for idx, node in enumerate(nodes): + end_page = skeleton.end_page + for nxt in nodes[idx + 1:]: + if nxt["rel_level"] <= node["rel_level"]: + end_page = max(nxt["start_page"] - 1, node["start_page"]) + break + deeper.append(SectionSkeleton( + section_path=node["section_path"], + level=node["abs_level"], + start_page=node["start_page"], + end_page=end_page, + title=node["title"], + parent_path=node["parent_path"], + evidence={ + "source": "fine_hierarchy_vlm", + "observed_page": node["start_page"], + "parent_skeleton": skeleton.section_path, + "llm_relative_level": node["rel_level"], + }, + )) + + _save_debug_hierarchy( + output_dir=output_dir, + skeleton=skeleton, + candidates=candidates, + ordered=ordered, + ) + return deeper or None + + +def _exclusive_end(skeletons: list[SectionSkeleton], index: int) -> int: + skeleton = skeletons[index] + later_starts = [ + skel.start_page + for skel in skeletons + if skel.start_page > skeleton.start_page + ] + if not later_starts: + return skeleton.end_page + return min(skeleton.end_page, min(later_starts) - 1) + + +def _title_key(title: str) -> str: + normalized = re.sub(r"\s+", "", str(title or "")).casefold() + normalized = re.sub(r"[^\w\u4e00-\u9fff]+", "", normalized) + return normalized + + +def _parse_hierarchy_result(result: Any, *, max_depth: int) -> dict[int, int]: + if not isinstance(result, list): + logger.warning( + "[page_memory.fine_hierarchy] unexpected hierarchy response type: {}", + type(result).__name__, + ) + return {} + parsed: dict[int, int] = {} + for item in result: + if not isinstance(item, dict): + continue + try: + row_id = int(item["id"]) + level = int(item["level"]) + except (KeyError, TypeError, ValueError): + continue + if level < 1: + continue + parsed[row_id] = min(level, max_depth) + return parsed + + +def _save_debug_hierarchy( + *, + output_dir: str | None, + skeleton: SectionSkeleton, + candidates: list[dict[str, Any]], + ordered: list[dict[str, Any]], +) -> None: + if not output_dir: + return + payload = { + "parent": skeleton.to_dict(), + "candidates": candidates, + "hierarchy": ordered, + } + try: + debug_path = Path(output_dir) / "page_memory_fine_hierarchy.json" + debug_path.write_text( + json.dumps(payload, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + except Exception as exc: + logger.debug( + "[page_memory.fine_hierarchy] failed to save debug hierarchy: {}", + exc, + ) diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 06ad4922..77d59dc3 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -32,9 +32,9 @@ class PageMemoryInput: def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: """Run the page-memory track. - Supports three granularity verdicts: + Supports two granularity verdicts: - ``whole_doc`` (≤6 pages, no TOC) → single whole-document chunk - - ``page`` / ``shard_page`` → per-page chunks via the full C1-C7 pipeline + - ``page`` → per-page chunks via the full C1-C7 pipeline """ full_output_dir = _resolve_output_dir(request) os.makedirs(full_output_dir, exist_ok=True) @@ -61,7 +61,7 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: verdict=verdict, ) - # page / shard_page → unified per-page pipeline + # page → per-page pipeline return full_output_dir, _build_page_dataframe( pdf_path=pdf_path, filename=request.filename, @@ -77,11 +77,17 @@ def _resolve_output_dir(request: PageMemoryInput) -> str: def _decide_granularity(profile: Any) -> str: + """Decide granularity for a page-memory document. + + Page-based track processes pages individually via VLM — no physical + document splitting is needed regardless of page count. The old + ``shard_page`` verdict (>200 pages) was only meaningful for the + MinerU batch API pipeline (``_parse_pdf_via_shards``) which splits + long PDFs into physical sub-documents. + """ page_count = int(getattr(profile, "page_count", 0) or 0) toc = getattr(profile, "toc", None) has_toc = bool(getattr(toc, "has_toc", False)) - if page_count > 200: - return "shard_page" if page_count <= 6 and not has_toc: return "whole_doc" return "page" @@ -105,6 +111,8 @@ def _build_page_dataframe( C1 page_renderer → PageRenderResult[] C2 page_plan → PagePlan[] C3 page_tagger → PageTagResult[] + C3b title detection → observed_titles (native hierarchy only) + C4b fine_hierarchy → refined SectionSkeleton[] (native hierarchy only) C6 page_section_mapper → PageSectionMapping[] C7 assemble DataFrame """ @@ -115,31 +123,58 @@ def _build_page_dataframe( from app.services.page_memory.page_tagger import tag_pages from app.services.page_memory.skeleton_extractor import extract_section_skeletons + # ── Native hierarchy flag (Step 3) ──────────────────────────────── + # Page-memory handles PDF/PPT(→PDF) structure with its own page-native + # hierarchy. Chunk-track hierarchy detection remains isolated to non + # page-memory formats such as DOCX/MD. + native_hierarchy = os.environ.get( + "PAGE_MEMORY_NATIVE_HIERARCHY", "true" + ).strip().lower() in ("1", "true", "yes", "on") + anatomy = getattr(profile, "anatomy", None) page_count = max(int(profile.page_count or 0), 0) if page_count <= 0: return pd.DataFrame(columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) - # ── unified budget (page_locate + page_tagging share one tracker) ── + # ── unified budget (page_locate + page_tagging + title_detection) ── page_tagging_budget = int( os.environ.get("PAGE_MEMORY_TAG_BUDGET", str(page_count * 1200)) ) page_locate_budget = int( os.environ.get("PAGE_MEMORY_LOCATE_BUDGET", str(min(page_count * 1600, 2_000_000))) ) + title_detection_budget = int( + os.environ.get("PAGE_MEMORY_TITLE_BUDGET", str(page_count * 800)) + ) if native_hierarchy else 0 + total_visual = page_tagging_budget + page_locate_budget + title_detection_budget + + # plan_budget powers the LLM decision loop inside PageLocateSubAgent. + # Without it the sub-agent falls back to a deterministic path that + # cannot rewrite queries (e.g. drop trailing doc-reference codes), + # causing grep to miss titles whose TOC text differs from body text. + plan_budget = int( + os.environ.get("PAGE_MEMORY_PLAN_BUDGET", str(min(page_count * 800, 2_000_000))) + ) + + stage_envelopes = { + "page_locate": StageEnvelope( + min_guarantee=page_locate_budget, + cap=None, + ), + "page_tagging": StageEnvelope( + min_guarantee=page_tagging_budget, + cap=None, + ), + } + if native_hierarchy: + stage_envelopes["page_title_detection"] = StageEnvelope( + min_guarantee=title_detection_budget, + cap=None, + ) budget = BudgetTracker( - plan_budget=0, - visual_budget=page_tagging_budget + page_locate_budget, - visual_stage_envelopes={ - "page_locate": StageEnvelope( - min_guarantee=page_locate_budget, - cap=None, - ), - "page_tagging": StageEnvelope( - min_guarantee=page_tagging_budget, - cap=None, - ), - }, + plan_budget=plan_budget, + visual_budget=total_visual, + visual_stage_envelopes=stage_envelopes, ) # ── build ToolContext for sub-agent VLM calls ───────────────────── @@ -196,6 +231,53 @@ def _build_page_dataframe( vlm_model=vlm_model, ) + # ── C3b + C4b: Native hierarchy (Step 2 + Step 3) ──────────────── + if native_hierarchy and skeletons: + from app.services.page_memory.fine_hierarchy import ( + compute_fat_leaf_pages, + refine_fat_leaf_skeletons, + ) + from app.services.page_memory.page_tagger import ( + get_fine_min_pages, + tag_page_titles, + ) + + fine_min = get_fine_min_pages() + fat_leaf_pages = compute_fat_leaf_pages(skeletons, min_pages=fine_min) + logger.info( + "[page_memory] native hierarchy: {} fat-leaf pages (min={})", + len(fat_leaf_pages), fine_min, + ) + + # Step 2: VLM title detection on fat-leaf pages + if fat_leaf_pages: + tags = tag_page_titles( + pages=rendered, + tag_results=tags, + fat_leaf_pages=fat_leaf_pages, + budget=budget, + vlm_model=vlm_model, + ) + + # Step 3: Refine coarse skeletons with LLM hierarchy + skeletons = refine_fat_leaf_skeletons( + coarse_skeletons=skeletons, + tag_results=tags, + fat_leaf_pages=fat_leaf_pages, + model_name=os.environ.get( + "PAGE_MEMORY_HIERARCHY_MODEL", + os.environ.get( + "HIERARCHY_LLM_MODEL", + os.environ.get("NORMOL_MODEL"), + ), + ), + output_dir=output_dir, + ) + logger.info( + "[page_memory] C4b refined: {} sections after fine hierarchy", + len(skeletons), + ) + # ── C6: page → section mapping ─────────────────────────────────── mappings = map_pages_to_sections( page_count=page_count, @@ -213,33 +295,100 @@ def _build_page_dataframe( for lbl in page_labels: label_map[lbl.page] = lbl.kind - # Build page → observed_titles from C4 skeleton (primary sections) - skeleton_titles: dict[int, list[str]] = {} - for skel in skeletons: - titles = skeleton_titles.setdefault(skel.start_page, []) - if skel.title and skel.title not in titles: - titles.append(skel.title) + # Build page → observed_titles. Native hierarchy uses only true VLM + # observations; the skeleton-title fallback is kept only for legacy mode. + observed_titles_map: dict[int, list[str]] = {} + if native_hierarchy: + for tag in tags: + if tag.observed_titles: + observed_titles_map[tag.page_index] = [ + t["text"] for t in tag.observed_titles if t.get("text") + ] + else: + # Legacy: map skeleton titles to start pages + for skel in skeletons: + titles = observed_titles_map.setdefault(skel.start_page, []) + if skel.title and skel.title not in titles: + titles.append(skel.title) + + # Shared per-page lookups for node-granularity assembly. + raw_text_by_page: dict[int, str] = {} + image_uri_by_page: dict[int, str] = {} + image_path_by_page: dict[int, str] = {} + for page in range(1, page_count + 1): + rend = render_map.get(page) + raw_text_by_page[page] = (rend.raw_text if rend else page_texts.get(page, "")) or "" + if rend and rend.image_path and os.path.exists(rend.image_path): + image_path_by_page[page] = rend.image_path + image_uri_by_page[page] = str( + Path(rend.image_path).relative_to(output_dir) + ) + + if native_hierarchy and skeletons: + from app.services.page_memory.node_assembler import build_node_rows + + rows = build_node_rows( + skeletons=skeletons, + raw_text_by_page=raw_text_by_page, + image_uri_by_page=image_uri_by_page, + image_path_by_page=image_path_by_page, + kind_by_page=label_map, + tag_by_page=tag_map, + filename=filename, + verdict=verdict, + native_hierarchy=native_hierarchy, + budget=budget, + vlm_model=vlm_model, + ) + logger.info( + "[page_memory] C7 assembled {} node rows (verdict={})", + len(rows), verdict, + ) + else: + rows = _build_legacy_page_rows( + mappings=mappings, + tag_map=tag_map, + raw_text_by_page=raw_text_by_page, + image_uri_by_page=image_uri_by_page, + label_map=label_map, + observed_titles_map=observed_titles_map, + filename=filename, + verdict=verdict, + native_hierarchy=native_hierarchy, + ) + logger.info( + "[page_memory] C7 assembled {} page rows (verdict={})", + len(rows), verdict, + ) + df = pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) + df.attrs["page_memory_skeletons"] = [skel.to_dict() for skel in skeletons] + df.attrs["page_memory_native_hierarchy"] = native_hierarchy + return df + +def _build_legacy_page_rows( + *, + mappings: Any, + tag_map: dict[int, Any], + raw_text_by_page: dict[int, str], + image_uri_by_page: dict[int, str], + label_map: dict[int, str], + observed_titles_map: dict[int, list[str]], + filename: str, + verdict: str, + native_hierarchy: bool, +) -> list[dict[str, Any]]: + """Legacy per-page rows (one ``type=page`` chunk per physical page).""" rows: list[dict[str, Any]] = [] for mapping in mappings: page = mapping.page_index tag = tag_map.get(page) - rend = render_map.get(page) - raw_text = rend.raw_text if rend else page_texts.get(page, "") - content = raw_text.strip() + content = (raw_text_by_page.get(page, "") or "").strip() summary = tag.summary if tag else "" - keywords_list = tag.keywords if tag else [] - keywords_str = ";".join(keywords_list) + keywords_str = ";".join(tag.keywords) if tag else "" - doc_hash = gen_str_codes(f"{filename}::{page}") - know_id = f"page_{doc_hash}" - - image_uri = "" - if rend and rend.image_path and os.path.exists(rend.image_path): - image_uri = str( - Path(rend.image_path).relative_to(output_dir) - ) + know_id = f"page_{gen_str_codes(f'{filename}::{page}')}" rows.append({ "content": content, @@ -256,20 +405,16 @@ def _build_page_dataframe( "extra_metadata": { "granularity": "page", "page_index": page, - "page_image_uri": image_uri, + "page_image_uri": image_uri_by_page.get(page, ""), "strategy_used": tag.strategy_used if tag else "", "kind": label_map.get(page, "normal"), - "observed_titles": skeleton_titles.get(page, []), + "observed_titles": observed_titles_map.get(page, []), "section_roles": mapping.section_roles, "source_verdict": verdict, + "native_hierarchy": native_hierarchy, }, }) - - logger.info( - "[page_memory] C7 assembled {} page rows (verdict={})", - len(rows), verdict, - ) - return pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) + return rows def _build_page_ctx( diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py new file mode 100644 index 00000000..1e2a52e7 --- /dev/null +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -0,0 +1,526 @@ +"""Node-granularity assembly for the page-memory track. + +Page-track historically emitted one ``type=page`` chunk per physical page, +so a section spanning N pages produced N chunks sharing the same ``path``. +This module switches the unit of assembly to the **leaf section node**: + +- One chunk per leaf node (path). Internal/structural nodes live in the + navigation tree only (their summaries aggregate bottom-up). +- A page that belongs to multiple nodes is *referenced* by each of them; the + page image is never duplicated (``page_image_uris`` is a list of shared + references). +- A page's body text is stored **once**, under the first leaf (in reading + order) that covers it. Other nodes that share the page emit a + ``SAME-AS `` marker instead of repeating the text. Body text is + not a core asset in page-track — the page image is. + +Summary/keywords are settled per node (see ``node_summary``): a node covering +multiple pages is summarized as a whole, and a page hosting multiple nodes is +summarized per node using a title boundary so the slices do not overlap. +""" + +from __future__ import annotations + +import base64 +import json +import os +from dataclasses import dataclass, field +from typing import Any, cast + +from loguru import logger + +from app.services.document_agent.budget import BudgetTracker +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.page_memory.page_tagger import PageTagResult +from app.services.page_memory.skeleton_extractor import SectionSkeleton +from shared.services.ai.prompt_service import build_prompt +from shared.utils.token_estimate import estimate_tokens + +SAME_AS_PREFIX = "SAME-AS" + +_BUDGET_STAGE = "page_tagging" +_NODE_SUMMARY_MAX_PAGES_DEFAULT = 5 + + +@dataclass(frozen=True) +class LeafNode: + """A leaf section node = one node-granularity chunk.""" + + section_path: str + title: str + level: int + start_page: int + end_page: int + parent_path: str | None = None + + +@dataclass +class NodePageView: + """Resolved page assignment for a single leaf node.""" + + leaf: LeafNode + pages: list[int] = field(default_factory=list) + """All body pages covered by this leaf (owned + shared), in order.""" + + owned_pages: list[int] = field(default_factory=list) + """Pages whose body text is stored under this node.""" + + +def identify_leaf_nodes(skeletons: list[SectionSkeleton]) -> list[LeafNode]: + """Return leaf skeletons in reading order. + + A skeleton is a leaf when no other skeleton declares it as ``parent_path``. + Reading order = ``(start_page, original index)`` so that siblings sharing a + start page keep the top-to-bottom order produced by fine hierarchy. + """ + parent_paths = { + skel.parent_path + for skel in skeletons + if skel.parent_path + } + leaves: list[tuple[int, int, LeafNode]] = [] + for index, skel in enumerate(skeletons): + if skel.section_path in parent_paths: + continue + leaves.append( + ( + skel.start_page, + index, + LeafNode( + section_path=skel.section_path, + title=skel.title, + level=skel.level, + start_page=skel.start_page, + end_page=skel.end_page, + parent_path=skel.parent_path, + ), + ) + ) + leaves.sort(key=lambda item: (item[0], item[1])) + return [leaf for _, _, leaf in leaves] + + +def assign_pages_to_leaves( + leaves: list[LeafNode], + *, + available_pages: set[int], +) -> tuple[list[NodePageView], dict[int, LeafNode]]: + """Map every leaf to its pages and decide per-page text ownership. + + Parameters + ---------- + leaves: + Leaf nodes in reading order. + available_pages: + Pages that actually have a rendered/body chunk (excludes pages outside + the document or with no content). + + Returns + ------- + (views, page_owner) + ``views`` is one ``NodePageView`` per leaf (same order as ``leaves``). + ``page_owner`` maps a page index to the leaf that owns its body text + (the first leaf, in reading order, that covers it). + """ + page_owner: dict[int, LeafNode] = {} + views: list[NodePageView] = [] + for leaf in leaves: + pages = [ + page + for page in range(leaf.start_page, leaf.end_page + 1) + if page in available_pages + ] + owned: list[int] = [] + for page in pages: + if page not in page_owner: + page_owner[page] = leaf + owned.append(page) + views.append(NodePageView(leaf=leaf, pages=pages, owned_pages=owned)) + return views, page_owner + + +def build_node_content( + view: NodePageView, + *, + page_owner: dict[int, LeafNode], + page_text: dict[int, str], +) -> str: + """Assemble a node's body content with per-page text deduplication. + + Owned pages contribute their resolved body text; shared pages contribute a + ``SAME-AS p`` reference so the text is stored once. + """ + segments: list[str] = [] + for page in view.pages: + owner = page_owner.get(page) + if owner is not None and owner.section_path == view.leaf.section_path: + segments.append((page_text.get(page) or "").strip()) + elif owner is not None: + segments.append(f"[{SAME_AS_PREFIX} {owner.section_path} p{page}]") + return "\n\n".join(segment for segment in segments if segment).strip() + + +def next_title_on_page( + leaf: LeafNode, + *, + page: int, + leaves_on_page: list[LeafNode], +) -> str | None: + """Return the title of the next leaf starting on the same page after *leaf*. + + Used to bound the summary slice when a page hosts multiple nodes. Returns + ``None`` when *leaf* is the last node beginning on this page. + """ + ordered = [item for item in leaves_on_page if item.start_page == page] + for index, item in enumerate(ordered): + if item.section_path == leaf.section_path: + if index + 1 < len(ordered): + return ordered[index + 1].title + return None + return None + + +def pages_by_leaf_count(views: list[NodePageView]) -> dict[int, list[LeafNode]]: + """Map each page to the leaves that cover it (reading order).""" + page_to_leaves: dict[int, list[LeafNode]] = {} + for view in views: + for page in view.pages: + page_to_leaves.setdefault(page, []).append(view.leaf) + return page_to_leaves + + +# ── VLM-backed helpers ─────────────────────────────────────────────── + + +def _node_summary_max_pages() -> int: + return int( + os.environ.get( + "PAGE_MEMORY_NODE_SUMMARY_MAX_PAGES", + str(_NODE_SUMMARY_MAX_PAGES_DEFAULT), + ) + ) + + +def _read_image_b64(image_path: str) -> str | None: + try: + with open(image_path, "rb") as handle: + return base64.b64encode(handle.read()).decode() + except Exception as exc: # pragma: no cover - filesystem edge + logger.warning("[node_assembler] failed to read image {}: {}", image_path, exc) + return None + + +def resolve_page_text( + *, + page: int, + raw_text: str, + image_path: str | None, + vlm_model: str | None, + budget: BudgetTracker | None, +) -> str: + """Body text for an owned page: PyMuPDF text, or VLM OCR for scanned pages. + + Electronic PDFs already have PyMuPDF text; scanned pages have (near) empty + text and fall back to a one-shot VLM OCR transcription. + """ + 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 "" + + img_b64 = _read_image_b64(image_path) + if img_b64 is None: + return "" + + prompt, temperature, _top_p, max_tokens = build_prompt( + "page-memory-vlm-ocr", "", "", paras={"max_tokens": 1500} + ) + est = estimate_tokens(prompt) + 1000 + if budget is not None and not budget.try_reserve("visual", est, stage=_BUDGET_STAGE): + logger.debug("[node_assembler] OCR budget exhausted for page {}", page) + return "" + + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=vlm_model) + content_parts = [ + {"type": "text", "text": prompt}, + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + }, + ] + raw_response, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=vlm_model, + temperature=temperature, + max_tokens=max_tokens, + response_format={"type": "json_object"}, + usage_task="page_memory.node_ocr", + ) + if budget is not None: + budget.commit( + "visual", + actual=usage.get("total_tokens", est), + est=est, + stage=_BUDGET_STAGE, + ) + data = json.loads(raw_response) + return str(data.get("text", "")).strip() + except Exception as exc: + logger.warning("[node_assembler] OCR failed for page {}: {}", page, exc) + if budget is not None: + budget.refund("visual", est=est, stage=_BUDGET_STAGE) + return "" + + +def compute_node_summary( + *, + view: NodePageView, + page_to_leaves: dict[int, list[LeafNode]], + tag_by_page: dict[int, PageTagResult], + image_path_by_page: dict[int, str], + vlm_model: str | None, + budget: BudgetTracker | None, +) -> tuple[str, list[str]]: + """Settle a node's summary/keywords. + + Reuses the per-page tag when the node is a single page that no sibling leaf + shares. Otherwise asks the VLM to summarize the node as a whole, bounding + the slice with the next sibling title when the page hosts multiple nodes. + Falls back to combining per-page tags when the VLM is unavailable. + """ + pages = view.pages + if not pages: + return "", [] + + single_page = len(pages) == 1 + shared = any(len(page_to_leaves.get(page, [])) > 1 for page in pages) + + if single_page and not shared: + tag = tag_by_page.get(pages[0]) + if tag is not None: + return tag.summary, list(tag.keywords) + return "", [] + + if vlm_model: + result = _vlm_node_summary( + view=view, + page_to_leaves=page_to_leaves, + image_path_by_page=image_path_by_page, + vlm_model=vlm_model, + budget=budget, + ) + if result is not None: + return result + + return _combine_page_tags(pages=pages, tag_by_page=tag_by_page) + + +def _combine_page_tags( + *, + pages: list[int], + tag_by_page: dict[int, PageTagResult], +) -> tuple[str, list[str]]: + summaries: list[str] = [] + keywords: list[str] = [] + seen: set[str] = set() + for page in pages: + tag = tag_by_page.get(page) + if tag is None: + continue + summary = (tag.summary or "").strip() + if summary and summary.upper() != "EMPTY": + summaries.append(summary) + for keyword in tag.keywords: + key = keyword.strip().casefold() + if key and key not in seen: + seen.add(key) + keywords.append(keyword.strip()) + return " ".join(summaries).strip(), keywords + + +def _vlm_node_summary( + *, + view: NodePageView, + page_to_leaves: dict[int, list[LeafNode]], + image_path_by_page: dict[int, str], + vlm_model: str, + budget: BudgetTracker | None, +) -> tuple[str, list[str]] | None: + leaf = view.leaf + pages = view.pages[: _node_summary_max_pages()] + + # Boundary title: only meaningful when this node's start page hosts a later + # sibling node, so the VLM can stop at that boundary. + next_title = next_title_on_page( + leaf, + page=leaf.start_page, + leaves_on_page=page_to_leaves.get(leaf.start_page, []), + ) + + image_parts: list[dict[str, Any]] = [] + for page in pages: + path = image_path_by_page.get(page) + if not path or not os.path.exists(path): + continue + img_b64 = _read_image_b64(path) + if img_b64 is None: + continue + image_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + if not image_parts: + return None + + prompt, temperature, _top_p, max_tokens = build_prompt( + "page-memory-node-summary", + "", + "", + paras={ + "max_tokens": 400, + "node_title": leaf.title, + "next_title": next_title or "", + "kw_num": 5, + }, + ) + est = estimate_tokens(prompt) + 800 * len(image_parts) + if budget is not None and not budget.try_reserve("visual", est, stage=_BUDGET_STAGE): + logger.debug( + "[node_assembler] node summary budget exhausted for {}", + leaf.section_path, + ) + return None + + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=vlm_model) + content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + content_parts.extend(image_parts) + raw_response, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=vlm_model, + temperature=temperature, + max_tokens=max_tokens, + response_format={"type": "json_object"}, + usage_task="page_memory.node_summary", + ) + if budget is not None: + budget.commit( + "visual", + actual=usage.get("total_tokens", est), + est=est, + stage=_BUDGET_STAGE, + ) + data = json.loads(raw_response) + kw_str = str(data.get("keywords", "")) + keywords = [k.strip() for k in kw_str.split(";") if k.strip()] + return str(data.get("summary", "")).strip(), keywords + except Exception as exc: + logger.warning( + "[node_assembler] node summary VLM failed for {}: {}", + leaf.section_path, + exc, + ) + if budget is not None: + budget.refund("visual", est=est, stage=_BUDGET_STAGE) + return None + + +# ── Orchestration ──────────────────────────────────────────────────── + + +def build_node_rows( + *, + skeletons: list[SectionSkeleton], + raw_text_by_page: dict[int, str], + image_uri_by_page: dict[int, str], + image_path_by_page: dict[int, str], + kind_by_page: dict[int, str], + tag_by_page: dict[int, PageTagResult], + filename: str, + verdict: str, + native_hierarchy: bool, + budget: BudgetTracker | None = None, + vlm_model: str | None = None, +) -> list[dict[str, Any]]: + """Assemble one row per leaf section node (node-granularity chunks).""" + 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) + + # Resolve body text once per owned page (PyMuPDF, OCR fallback for scanned). + 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, + budget=budget, + ) + + rows: list[dict[str, Any]] = [] + for view in views: + leaf = view.leaf + content = build_node_content( + view, + page_owner=page_owner, + page_text=resolved_text, + ) + summary, keywords = 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, + budget=budget, + ) + page_image_uris = [ + image_uri_by_page[page] + for page in view.pages + if image_uri_by_page.get(page) + ] + know_id = f"node_{gen_str_codes(f'{filename}::{leaf.section_path}')}" + rows.append( + { + "content": content, + "path": leaf.section_path, + "type": "page", + "length": len(content), + "keywords": ";".join(keywords), + "summary": summary, + "know_id": know_id, + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": ",".join(str(page) for page in view.pages), + "extra_metadata": { + "granularity": "node", + "section_path": leaf.section_path, + "section_level": leaf.level, + "page_indices": list(view.pages), + "owned_pages": list(view.owned_pages), + "page_image_uris": page_image_uris, + "kind": kind_by_page.get(leaf.start_page, "normal"), + "source_verdict": verdict, + "native_hierarchy": native_hierarchy, + }, + } + ) + + logger.info( + "[node_assembler] assembled {} node rows from {} leaves ({} pages)", + len(rows), + len(leaves), + len(available_pages), + ) + return rows diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index eda9178a..9428de10 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -1,4 +1,4 @@ -"""Page tagger: VLM per-page annotation for summary and keywords. +"""Page tagger: VLM per-page annotation for summary, keywords, and title candidates. For ``vlm_lite`` pages, sends the page PNG to the VLM and expects a JSON response with ``summary`` and ``keywords``. @@ -6,6 +6,11 @@ to extract summary + keywords from raw text. For ``skip_tagging`` pages, content is preserved but summary is omitted. +Step 2 of page-memory native hierarchy adds: +- Independent VLM title candidate extraction (``observed_titles``) +- Fat-leaf gating: only pages in TOC leaves with > N pages trigger title detection +- Title extraction uses a dedicated verbatim-only prompt (temp=0, small max_tokens) + Budget is drawn from the ``page_tagging`` stage envelope. """ @@ -22,6 +27,7 @@ from app.services.document_agent.budget import BudgetTracker from app.services.page_memory.page_plan import PagePlan, PageProcessingStrategy from app.services.page_memory.page_renderer import PageRenderResult +from shared.services.ai.prompt_service import build_prompt from shared.utils.token_estimate import estimate_tokens @@ -33,32 +39,19 @@ class PageTagResult: summary: str = "" keywords: list[str] = field(default_factory=list) strategy_used: str = "" + observed_titles: list[dict[str, Any]] = field(default_factory=list) + """Step 2: verbatim title candidates observed on this page. + Each entry is ``{"text": str, "prominence": float | None}``. + Empty list means no titles were detected (or title detection was skipped). + """ -# ── VLM prompt: outputs summary + keywords only ───────────────────── - -_VLM_TAG_PROMPT = """\ -You are annotating a single PDF page screenshot for a document memory system. -Return strict JSON with exactly these keys: - -{ - "summary": "<1-3 sentence summary of the page content>", - "keywords": ";;" -} - -Rules: -- "summary": describe the main content visible on the page in 1-3 sentences. - If the page contains tables, mention the table topic and key columns. - If the page contains figures or charts, describe what they depict. -- "keywords": extract the most important thematic keywords (up to 5), - separated by semicolons ";". Keywords must be in the same language as - the visible page content. -- Return ONLY the JSON object, no markdown fences or extra text. -""" _BUDGET_STAGE = "page_tagging" +_BUDGET_STAGE_TITLES = "page_title_detection" _MAX_JSON_RETRIES = 1 _RAW_TEXT_SUMMARY_LIMIT = 500 +_DEFAULT_FINE_MIN_PAGES = 4 def tag_pages( @@ -157,7 +150,6 @@ def _tag_text_only(page: PageRenderResult) -> PageTagResult: # Try the existing summary-full LLM call (same as text chunk pipeline) try: - from shared.services.ai.prompt_service import build_prompt from shared.services.ai.openai_compatible_client_sync import get_openai_client text_model = os.environ.get("NORMOL_MODEL", "deepseek-v4-flash") @@ -211,7 +203,13 @@ def _tag_vlm_lite( budget: BudgetTracker | None, ) -> PageTagResult: """Send page PNG to VLM and parse JSON response.""" - est = estimate_tokens(_VLM_TAG_PROMPT) + 800 # ~800 tokens for image + prompt, temperature, _top_p, max_tokens = build_prompt( + "page-memory-vlm-tag", + "", + "", + paras={"max_tokens": 600}, + ) + est = estimate_tokens(prompt) + 800 # ~800 tokens for image if budget is not None: if not budget.try_reserve("visual", est, stage=_BUDGET_STAGE): @@ -250,7 +248,7 @@ def _tag_vlm_lite( return _tag_text_only(page) content_parts: list[dict[str, Any]] = [ - {"type": "text", "text": _VLM_TAG_PROMPT}, + {"type": "text", "text": prompt}, { "type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}, @@ -266,8 +264,8 @@ def _tag_vlm_lite( raw_response, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, - temperature=0.0, - max_tokens=600, + temperature=temperature, + max_tokens=max_tokens, response_format={"type": "json_object"}, usage_task="page_memory.tag", ) @@ -319,3 +317,196 @@ def _tag_vlm_lite( # Should not reach here, but safety net return _tag_text_only(page) + + +# ── Step 2: Independent title candidate extraction ─────────────────── + + +def get_fine_min_pages() -> int: + """Fat-leaf gating threshold from env ``PAGE_MEMORY_FINE_MIN_PAGES``.""" + return int(os.environ.get("PAGE_MEMORY_FINE_MIN_PAGES", str(_DEFAULT_FINE_MIN_PAGES))) + + +def tag_page_titles( + *, + pages: list[PageRenderResult], + tag_results: list[PageTagResult], + fat_leaf_pages: set[int], + budget: BudgetTracker | None = None, + vlm_model: str | None = None, +) -> list[PageTagResult]: + """Run independent VLM title detection on fat-leaf pages. + + Parameters + ---------- + pages: + Rendered page results. + tag_results: + Existing tag results from ``tag_pages()`` (will be updated in-place). + fat_leaf_pages: + Set of page indices belonging to fat-leaf TOC sections + (those with > ``PAGE_MEMORY_FINE_MIN_PAGES`` pages). + budget: + Optional budget tracker. + vlm_model: + VLM model name; falls back to ``$IMAGE_MODEL``. + + Returns + ------- + list[PageTagResult] + Updated tag results with ``observed_titles`` populated for fat-leaf pages. + """ + if not fat_leaf_pages: + return tag_results + + model = vlm_model or os.environ.get("IMAGE_MODEL") + if not model: + logger.warning("[page_tagger] no VLM model for title detection; skipping") + return tag_results + + 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 + + # Skip pages without images (text_only / skip) + if not page.image_path or not os.path.exists(page.image_path): + continue + + observed = _tag_vlm_titles(page, model=model, budget=budget) + tag.observed_titles = observed + vlm_calls += 1 + titles_found += len(observed) + + logger.info( + "[page_tagger] title detection: {} VLM calls on {} fat-leaf pages, {} titles found", + vlm_calls, + len(fat_leaf_pages), + titles_found, + ) + return tag_results + + +def _tag_vlm_titles( + page: PageRenderResult, + *, + model: str, + budget: BudgetTracker | None, +) -> list[dict[str, Any]]: + """Send page PNG to VLM with the title-only prompt and parse results.""" + prompt, temperature, _top_p, max_tokens = build_prompt( + "page-memory-vlm-title", + "", + "", + paras={"max_tokens": 300}, + ) + est = estimate_tokens(prompt) + 800 # ~800 tokens for image + + if budget is not None: + if not budget.try_reserve("visual", est, stage=_BUDGET_STAGE_TITLES): + logger.debug( + "[page_tagger] title budget exhausted for page {}", + page.page_index, + ) + return [] + + try: + with open(page.image_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + except Exception as exc: + logger.warning( + "[page_tagger] failed to read PNG for title detection page {}: {}", + page.page_index, exc, + ) + if budget is not None: + budget.refund("visual", est=est, stage=_BUDGET_STAGE_TITLES) + return [] + + content_parts: list[dict[str, Any]] = [ + {"type": "text", "text": prompt}, + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + }, + ] + + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + + for attempt in range(_MAX_JSON_RETRIES + 1): + try: + raw_response, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=temperature, + max_tokens=max_tokens, + response_format={"type": "json_object"}, + usage_task="page_memory.title_detection", + ) + if budget is not None: + budget.commit( + "visual", + actual=usage.get("total_tokens", est), + est=est, + stage=_BUDGET_STAGE_TITLES, + ) + + data = json.loads(raw_response) + titles_raw = data.get("titles", []) + if not isinstance(titles_raw, list): + return [] + + observed: list[dict[str, Any]] = [] + for item in titles_raw: + if isinstance(item, dict) and item.get("text"): + text = str(item["text"]).strip() + + is_table = item.get("is_in_table") is True + is_header = item.get("is_in_header_footer") is True + + if is_table or is_header: + logger.debug( + "[page_tagger] filtered CoT title on page {}: '{}' (table={}, header={})", + page.page_index, text, is_table, is_header + ) + continue + + if text: + prominence = None + try: + prominence = float(item.get("prominence", 0.5)) + except (TypeError, ValueError): + pass + observed.append({ + "text": text, + "prominence": prominence, + "is_in_table": is_table, + "is_in_header_footer": is_header + }) + return observed + + except json.JSONDecodeError: + if attempt < _MAX_JSON_RETRIES: + continue + logger.warning( + "[page_tagger] title JSON retry exhausted for page {}", + page.page_index, + ) + return [] + except Exception as exc: + logger.warning( + "[page_tagger] title VLM failed for page {}: {}", + page.page_index, exc, + ) + if budget is not None: + budget.refund("visual", est=est, stage=_BUDGET_STAGE_TITLES) + return [] + + return [] diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 561d3985..c4e70c59 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -1,4 +1,16 @@ -"""Build page-memory section skeletons from profile-time anatomy.""" +"""Build page-memory section skeletons from profile-time anatomy. + +Step 1 of the page-memory native hierarchy plan: +- Full TOC-depth grep anchoring + on-demand VLM confirmation +- Section boundaries come purely from TOC anchoring + +Page-based track processes pages individually via VLM, so no physical +document splitting (shard windowing) is needed. The ``shard_plan`` +concept only exists to support the MinerU batch API pipeline +(``_parse_pdf_via_shards``), which requires splitting long PDFs into +physical sub-documents before sending them to MinerU. Chunk-track's +native text formats (DOCX/MD) don't use shards either. +""" from __future__ import annotations @@ -8,7 +20,6 @@ from app.services.document_agent.manifest import ( H1Candidate, PageAnatomyMap, - Shard, ToolContext, ) from app.services.document_agent.structure.page_locate_agent import ( @@ -51,7 +62,11 @@ def extract_section_skeletons( ctx: ToolContext | None = None, hierarchy_nodes: list[TitleNode] | None = None, ) -> list[SectionSkeleton]: - """Convert PageAnatomyMap hierarchy evidence into leaf section skeletons.""" + """Convert PageAnatomyMap hierarchy evidence into section skeletons. + + Section page ranges are anchored purely from the TOC hierarchy (every + level, every document). + """ page_count = _page_count(anatomy) root_path = f"{filename}/Root" if page_count <= 0: @@ -109,7 +124,6 @@ def extract_section_skeletons( item, filename=filename, page_count=page_count, - shards=_shards(anatomy), locate_summary=locate_result.summary, toc_selection=toc_selection, ) @@ -124,16 +138,11 @@ def _range_to_skeleton( *, filename: str, page_count: int, - shards: list[Shard], locate_summary: dict[str, Any], toc_selection: dict[str, Any], ) -> SectionSkeleton: start_page = _clamp_page(item.start_page, page_count) - end_page = _clamp_to_shard( - start_page=start_page, - end_page=_clamp_page(item.end_page, page_count), - shards=shards, - ) + end_page = _clamp_page(item.end_page, page_count) path_titles = [clean_toc_title(title) or title for title in item.path_titles] section_path = "/".join([filename, *path_titles]) parent_path = "/".join([filename, *path_titles[:-1]]) if len(path_titles) > 1 else filename @@ -145,8 +154,6 @@ def _range_to_skeleton( } if toc_selection: evidence["toc_selection"] = toc_selection - if end_page != item.end_page: - evidence["shard_clamped"] = True return SectionSkeleton( section_path=section_path, level=item.level, @@ -328,18 +335,6 @@ def _title_key(title: str) -> str: return normalize_heading_text(clean_toc_title(title) or title).casefold() -def _shards(anatomy: Any | None) -> list[Shard]: - shard_plan = getattr(anatomy, "shard_plan", None) - return list(getattr(shard_plan, "shards", []) or []) - - -def _clamp_to_shard(*, start_page: int, end_page: int, shards: list[Shard]) -> int: - for shard in shards: - if shard.page_start <= start_page <= shard.page_end: - return min(end_page, shard.page_end) - return end_page - - def _clamp_page(page: int, page_count: int) -> int: return min(max(page, 1), max(page_count, 1)) diff --git a/apps/worker/tests/contract/test_document_agent_budget_contract.py b/apps/worker/tests/contract/test_document_agent_budget_contract.py index 73670d76..a3076b56 100644 --- a/apps/worker/tests/contract/test_document_agent_budget_contract.py +++ b/apps/worker/tests/contract/test_document_agent_budget_contract.py @@ -13,6 +13,7 @@ from app.services.document_agent.budget import BudgetTracker, StageEnvelope from app.services.page_memory import memory_service from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks +from shared.services.storage.zip_chunk_schema import ZipChunkSchemaBuilder import pandas as pd @@ -89,6 +90,47 @@ def test_dataframe_converter_accepts_page_chunks_with_extra_metadata() -> None: assert chunks[0]["metadata"]["page_nums"] == [1, 2] +def test_zip_chunk_schema_preserves_page_memory_section_roles() -> None: + chunks = [ + { + "chunk_id": "page-231", + "type": "page", + "content": "page body", + "path": "demo.pdf/3 基本规定/3.2 管理规定", + "metadata": { + "summary": "page summary", + "page_nums": [231], + "page_index": 231, + "section_roles": [ + { + "section_path": "demo.pdf/3 基本规定", + "role": "primary", + }, + { + "section_path": "demo.pdf/3 基本规定/3.1 职责", + "role": "primary", + }, + { + "section_path": "demo.pdf/3 基本规定/3.2 管理规定", + "role": "primary", + }, + ], + }, + } + ] + + formatted = ZipChunkSchemaBuilder().format_chunks( + chunks, + image_files_map={}, + table_files_map={}, + ) + + metadata = formatted[0]["metadata"] + assert metadata["page_nums"] == [231] + assert metadata["page_index"] == 231 + assert metadata["section_roles"] == chunks[0]["metadata"]["section_roles"] + + def test_page_memory_granularity_routes_supported_page_modes() -> None: assert ( memory_service._decide_granularity( # noqa: SLF001 @@ -106,5 +148,5 @@ def test_page_memory_granularity_routes_supported_page_modes() -> None: memory_service._decide_granularity( # noqa: SLF001 SimpleNamespace(page_count=201, toc=SimpleNamespace(has_toc=False)) ) - == "shard_page" + == "page" ) diff --git a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py new file mode 100644 index 00000000..c46eac98 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +import json +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") + +from app.services.page_memory import fine_hierarchy +from app.services.page_memory.page_tagger import PageTagResult +from app.services.page_memory.skeleton_extractor import SectionSkeleton + + +class _FakeClient: + def __init__(self, response: list[dict[str, int]]) -> None: + self.response = response + + def chat_completion(self, **_kwargs) -> str: + return json.dumps(self.response) + + +def test_compute_fat_leaf_pages_uses_exclusive_boundaries() -> None: + skeletons = [ + SectionSkeleton( + section_path="demo.pdf/A", + level=1, + start_page=10, + end_page=12, + title="A", + parent_path="demo.pdf", + ), + SectionSkeleton( + section_path="demo.pdf/B", + level=1, + start_page=12, + end_page=15, + title="B", + parent_path="demo.pdf", + ), + ] + + assert fine_hierarchy.compute_fat_leaf_pages(skeletons, min_pages=1) == { + 10, + 11, + 12, + 13, + 14, + 15, + } + + +def test_refine_fat_leaf_skeletons_excludes_next_section_start_when_unordered( + monkeypatch, +) -> None: + previous = SectionSkeleton( + section_path="demo.pdf/Section A", + level=3, + start_page=225, + end_page=302, + title="Section A", + parent_path="demo.pdf", + ) + later = SectionSkeleton( + section_path="demo.pdf/Section C", + level=3, + start_page=320, + end_page=330, + title="Section C", + parent_path="demo.pdf", + ) + next_section = SectionSkeleton( + section_path="demo.pdf/Section B", + level=3, + start_page=302, + end_page=319, + title="Section B", + parent_path="demo.pdf", + ) + tags = [ + PageTagResult( + page_index=301, + observed_titles=[{"text": "A.1 Last Heading", "prominence": 1.0}], + ), + PageTagResult( + page_index=302, + observed_titles=[ + { + "text": "B.1 Boundary Heading", + "prominence": 1.0, + } + ], + ), + ] + + monkeypatch.setattr( + fine_hierarchy, + "get_openai_client", + lambda model=None: _FakeClient([{"id": 1, "level": 1}]), + ) + + refined = fine_hierarchy.refine_fat_leaf_skeletons( + coarse_skeletons=[previous, later, next_section], + tag_results=tags, + fat_leaf_pages={301, 302}, + model_name="test-model", + ) + + assert [item.title for item in refined] == [ + "A.1 Last Heading", + "Section C", + "B.1 Boundary Heading", + ] + assert refined[0].parent_path == "demo.pdf/Section A" + assert refined[2].parent_path == "demo.pdf/Section B" + + +def test_refine_fat_leaf_skeletons_uses_page_memory_prompt_without_demoting_siblings( + monkeypatch, +) -> None: + skeleton = SectionSkeleton( + section_path="demo.pdf/安全风险分级管控", + level=1, + start_page=225, + end_page=245, + title="安全风险分级管控", + parent_path="demo.pdf", + ) + tags = [ + PageTagResult( + page_index=225, + observed_titles=[ + {"text": "安全风险分级管控", "prominence": 1.0}, + {"text": "1 总则", "prominence": 1.0}, + ], + ), + PageTagResult( + page_index=226, + observed_titles=[ + {"text": "2 术语", "prominence": 1.0}, + {"text": "2.1 定义", "prominence": 0.8}, + ], + ), + PageTagResult( + page_index=228, + observed_titles=[{"text": "3 基本规定", "prominence": 1.0}], + ), + ] + + monkeypatch.setattr( + fine_hierarchy, + "get_openai_client", + lambda model=None: _FakeClient( + [ + {"id": 1, "level": 1}, + {"id": 2, "level": 1}, + {"id": 3, "level": 2}, + {"id": 4, "level": 1}, + ] + ), + ) + + refined = fine_hierarchy.refine_fat_leaf_skeletons( + coarse_skeletons=[skeleton], + tag_results=tags, + fat_leaf_pages={225, 226, 227, 228}, + model_name="test-model", + ) + + assert [item.title for item in refined] == [ + "1 总则", + "2 术语", + "2.1 定义", + "3 基本规定", + ] + assert [item.level for item in refined] == [2, 2, 3, 2] + assert refined[0].end_page == 225 + assert refined[1].end_page == 227 + assert refined[2].parent_path.endswith("/2 术语") + assert all(item.title != skeleton.title for item in refined) diff --git a/apps/worker/tests/contract/test_page_memory_navigation_contract.py b/apps/worker/tests/contract/test_page_memory_navigation_contract.py new file mode 100644 index 00000000..75160aa4 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_navigation_contract.py @@ -0,0 +1,144 @@ +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 pandas as pd + +from shared.services.storage.zip_doc_navigation import build_doc_nav_from_skeletons +from shared.services.storage.zip_result_service import ZipResultService + + +def test_doc_nav_from_skeletons_preserves_same_page_sibling_sections() -> None: + skeletons = _same_page_sibling_skeletons() + chunks = _page_chunks() + + doc_nav = build_doc_nav_from_skeletons( + skeletons, + chunks, + "demo.pdf", + ) + + sections = doc_nav["sections"] + assert len(sections) == 1 + parent = sections[0] + assert parent["title"] == "3 基本规定" + assert parent["chunk_count"] == 2 + + child_counts = { + child["title"]: child["chunk_count"] for child in parent["children"] + } + assert child_counts == { + "3.1 职责": 1, + "3.2 管理规定": 2, + } + + +def test_doc_nav_from_skeletons_synthesizes_missing_ancestor_sections() -> None: + skeletons = [ + { + "section_path": "demo.pdf/安全类/SJSYJ-SC103/3 基本规定/3.1 职责", + "level": 4, + "start_page": 231, + "end_page": 231, + "title": "3.1 职责", + "parent_path": "demo.pdf/安全类/SJSYJ-SC103/3 基本规定", + }, + { + "section_path": "demo.pdf/安全类/SJSYJ-SC103/3 基本规定/3.2 管理规定", + "level": 4, + "start_page": 231, + "end_page": 232, + "title": "3.2 管理规定", + "parent_path": "demo.pdf/安全类/SJSYJ-SC103/3 基本规定", + }, + ] + + doc_nav = build_doc_nav_from_skeletons( + skeletons, + _page_chunks(), + "demo.pdf", + ) + + assert doc_nav["sections"][0]["title"] == "安全类" + assert doc_nav["sections"][0]["chunk_count"] == 2 + standard = doc_nav["sections"][0]["children"][0] + basic = standard["children"][0] + assert basic["title"] == "3 基本规定" + assert [child["title"] for child in basic["children"]] == [ + "3.1 职责", + "3.2 管理规定", + ] + + +def test_zip_result_service_uses_page_memory_skeletons_for_manifest_hierarchy() -> None: + parsed_df = pd.DataFrame() + parsed_df.attrs["page_memory_skeletons"] = _same_page_sibling_skeletons() + + doc_nav, hierarchy = ZipResultService()._build_navigation_outputs( # noqa: SLF001 + formatted_chunks=_page_chunks(), + source_file_name="demo.pdf", + parsed_df=parsed_df, + ) + + assert doc_nav is not None + assert hierarchy == { + "3 基本规定": { + "3.1 职责": {}, + "3.2 管理规定": {}, + } + } + + +def _same_page_sibling_skeletons() -> list[dict[str, object]]: + return [ + { + "section_path": "demo.pdf/3 基本规定", + "level": 1, + "start_page": 231, + "end_page": 232, + "title": "3 基本规定", + "parent_path": None, + }, + { + "section_path": "demo.pdf/3 基本规定/3.1 职责", + "level": 2, + "start_page": 231, + "end_page": 231, + "title": "3.1 职责", + "parent_path": "demo.pdf/3 基本规定", + }, + { + "section_path": "demo.pdf/3 基本规定/3.2 管理规定", + "level": 2, + "start_page": 231, + "end_page": 232, + "title": "3.2 管理规定", + "parent_path": "demo.pdf/3 基本规定", + }, + ] + + +def _page_chunks() -> list[dict[str, object]]: + return [ + { + "chunk_id": "page-231", + "type": "page", + "content": "page 231", + "path": "demo.pdf/3 基本规定/3.2 管理规定", + "metadata": {"page_nums": [231], "summary": "page 231 summary"}, + }, + { + "chunk_id": "page-232", + "type": "page", + "content": "page 232", + "path": "demo.pdf/3 基本规定/3.2 管理规定", + "metadata": {"page_nums": [232], "summary": "page 232 summary"}, + }, + ] 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 new file mode 100644 index 00000000..cf8dbb8a --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -0,0 +1,167 @@ +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") + +from app.services.page_memory.node_assembler import ( + SAME_AS_PREFIX, + assign_pages_to_leaves, + build_node_content, + build_node_rows, + identify_leaf_nodes, +) +from app.services.page_memory.page_tagger import PageTagResult +from app.services.page_memory.skeleton_extractor import SectionSkeleton + + +def _same_page_sibling_skeletons() -> list[SectionSkeleton]: + parent = SectionSkeleton( + section_path="demo.pdf/3 基本规定", + level=1, + start_page=231, + end_page=232, + title="3 基本规定", + parent_path="demo.pdf", + ) + child_a = SectionSkeleton( + section_path="demo.pdf/3 基本规定/3.1 职责", + level=2, + start_page=231, + end_page=231, + title="3.1 职责", + parent_path="demo.pdf/3 基本规定", + ) + child_b = SectionSkeleton( + section_path="demo.pdf/3 基本规定/3.2 管理规定", + level=2, + start_page=231, + end_page=232, + title="3.2 管理规定", + parent_path="demo.pdf/3 基本规定", + ) + return [parent, child_a, child_b] + + +def test_identify_leaf_nodes_drops_internal_parents() -> None: + leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) + assert [leaf.title for leaf in leaves] == ["3.1 职责", "3.2 管理规定"] + + +def test_page_ownership_first_leaf_owns_shared_page() -> None: + leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) + views, page_owner = assign_pages_to_leaves(leaves, available_pages={231, 232}) + + assert page_owner[231].title == "3.1 职责" + assert page_owner[232].title == "3.2 管理规定" + + by_title = {view.leaf.title: view for view in views} + assert by_title["3.1 职责"].owned_pages == [231] + assert by_title["3.2 管理规定"].owned_pages == [232] + assert by_title["3.2 管理规定"].pages == [231, 232] + + +def test_build_node_content_uses_same_as_for_shared_page() -> None: + leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) + views, page_owner = assign_pages_to_leaves(leaves, available_pages={231, 232}) + by_title = {view.leaf.title: view for view in views} + page_text = {231: "text-231", 232: "text-232"} + + content_a = build_node_content( + by_title["3.1 职责"], page_owner=page_owner, page_text=page_text + ) + content_b = build_node_content( + by_title["3.2 管理规定"], page_owner=page_owner, page_text=page_text + ) + + assert content_a == "text-231" + assert content_b.startswith(f"[{SAME_AS_PREFIX} demo.pdf/3 基本规定/3.1 职责 p231]") + assert "text-232" in content_b + assert "text-231" not in content_b + + +def test_build_node_rows_reuses_tags_without_vlm() -> None: + rows = build_node_rows( + skeletons=_same_page_sibling_skeletons(), + raw_text_by_page={231: "text-231", 232: "text-232"}, + image_uri_by_page={231: "pages/page-231.png", 232: "pages/page-232.png"}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={ + 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), + 232: PageTagResult(page_index=232, summary="s232", keywords=["k2"]), + }, + filename="demo.pdf", + verdict="page", + native_hierarchy=True, + budget=None, + vlm_model=None, + ) + + assert [r["path"] for r in rows] == [ + "demo.pdf/3 基本规定/3.1 职责", + "demo.pdf/3 基本规定/3.2 管理规定", + ] + by_path = {r["path"]: r for r in rows} + + leaf_a = by_path["demo.pdf/3 基本规定/3.1 职责"] + assert leaf_a["page_nums"] == "231" + assert leaf_a["content"] == "text-231" + assert leaf_a["extra_metadata"]["page_image_uris"] == ["pages/page-231.png"] + assert leaf_a["extra_metadata"]["granularity"] == "node" + + leaf_b = by_path["demo.pdf/3 基本规定/3.2 管理规定"] + assert leaf_b["page_nums"] == "231,232" + assert SAME_AS_PREFIX in leaf_b["content"] + assert "text-232" in leaf_b["content"] + assert leaf_b["extra_metadata"]["page_image_uris"] == [ + "pages/page-231.png", + "pages/page-232.png", + ] + assert leaf_b["extra_metadata"]["owned_pages"] == [232] + + +def test_build_node_rows_uses_vlm_node_summary_with_boundary(monkeypatch, tmp_path) -> None: + captured: dict[str, object] = {} + + class _FakeClient: + def chat_completion_with_usage(self, **kwargs): + captured["messages"] = kwargs.get("messages") + captured["usage_task"] = kwargs.get("usage_task") + return ('{"summary": "node summary", "keywords": "ka;kb"}', {"total_tokens": 10}) + + # node_assembler imports get_openai_client lazily from this module. + import shared.services.ai.openai_compatible_client_sync as client_mod + + monkeypatch.setattr(client_mod, "get_openai_client", lambda model=None: _FakeClient()) + + img = tmp_path / "page-231.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n fake") + + rows = build_node_rows( + skeletons=_same_page_sibling_skeletons(), + raw_text_by_page={231: "text-231", 232: "text-232"}, + image_uri_by_page={231: "pages/page-231.png", 232: "pages/page-232.png"}, + image_path_by_page={231: str(img), 232: str(img)}, + kind_by_page={}, + tag_by_page={ + 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), + 232: PageTagResult(page_index=232, summary="s232", keywords=["k2"]), + }, + filename="demo.pdf", + verdict="page", + native_hierarchy=True, + budget=None, + vlm_model="fake-vlm", + ) + + by_path = {r["path"]: r for r in rows} + leaf_a = by_path["demo.pdf/3 基本规定/3.1 职责"] + assert leaf_a["summary"] == "node summary" + assert leaf_a["keywords"] == "ka;kb" + assert captured["usage_task"] == "page_memory.node_summary" diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 0e65926e..960dadb5 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -195,149 +195,6 @@ def build_prompt(task, texts, query, **kwargs): """ # ==================== Heading/Structure Prompts ==================== - - # --------------------------------------------------------------------- - # LEGACY `eval-headings` prompt — designed for FULL-TEXT input, before - # `_compact_for_llm` collapses consecutive body rows into placeholders. - # Kept as reference; DO NOT delete. The live prompt below targets the - # COMPACT input shape used when `KB_LAYOUT_LLM_COMPACT_INPUT` is on - # (default). The live prompt below is the publication baseline. - # --------------------------------------------------------------------- - # elif task == 'eval-headings': - # temperature = 0 - # top_p = 0.01 - # max_depth = kwargs['paras']['max_depth'] - # max_tokens = kwargs['paras']['max_tokens'] - # toc_context = kwargs['paras'].get('toc_context', '') - # - # # developing toc context (if any) - # if toc_context: - # toc_section = f""" - # ***Important Reference: Table of Contents (TOC)*** - # The following is the document's table of contents with predefined levels. Use this as a reference when assigning levels: - # - # ''' - # {toc_context} - # ''' - # - # - If a row's heading matches a TOC entry, use the TOC's predefined level - # - If a row appears to be a sub-section of a TOC entry, assign a deeper level - # - IMPORTANT: If a row does NOT appear in the TOC, it CAN ONLY be set as either a body text (level = -1) or sub-section with a deeper level than the nearest TOC heading above it - # """ - # else: - # toc_section = "" - # - # prompt = f""" - # You are a document structure auditing expert. You will receive a Markdown table with text rows, where each row may be a heading or body text, including: - # 1. id column: line number - # 2. heading column: text content - # 3. level column: preliminary estimated level (may be inaccurate or missing), where: - # 1 represents `

` (highest), 2 represents `

`, and so on - # -1 indicates the text is estimated as body text (not a heading) - # "Not Sure" indicates the level is undetermined - # - # Data to be adjusted: - # ''' - # {texts} - # ''' - # - # {toc_section} - # - # ***Placeholder Rows*** - # Some rows may appear as "[N BODY LINES]" with an id like "55-63" (a range) or - # "56" (a single line), and level column rendered as "-". These are NOT real - # candidates — they are compact markers representing N consecutive body-text - # lines that have been collapsed to save space. Treat them only as positional - # context (they tell you how many body lines sit between two adjacent heading - # candidates). - # - You MUST NOT emit placeholder rows in your output. - # - The output id field MUST be a single integer; never return an id containing - # a hyphen ("-") or the level placeholder "-". - # - Only evaluate rows whose id is a single integer. - # - # ***Process in THREE steps:*** - # - # **STEP 1 — Global Pattern Scan (before assigning any levels)** - # Scan ALL candidate heading rows across the entire input. - # Identify every distinct structural/numbering pattern that signals hierarchy depth, for example: - # - Decimal numbering: "1", "1.1", "1.1.1" → depth increases with dot count - # - Enumeration styles such as Chinese numerals, numbered bullets, - # or circled digits map from shallower to deeper levels - # - Chapter/section keywords: "Chapter X", "Part X", and Chinese - # chapter/section markers - # - Indentation or formatting cues visible in the text prefix - # Rank these patterns from shallowest to deepest to form a pattern → level mapping. - # - # **STEP 2 — Assign levels using the following rules (in priority order)** - # Rows marked as "Not Sure" should be treated like any other candidate row: - # use the same rules below to decide whether they are true headings (level >= 1) - # or body text (level = -1). - # - # Rule 0 — Figure/Image rows are always body text (highest priority, no exceptions): - # Any row whose heading text is exactly "Figure/Image" MUST be assigned level = -1. - # These represent embedded images, figures, or inline resource references in the document. - # Do NOT include these rows in the output (they are automatically treated as level = -1). - # Rule 1 — Normalize to start at level 1: - # The shallowest heading pattern found in this document MUST be assigned level 1. - # Do NOT preserve preliminary estimates that start at level 2, 3, or deeper - # if those headings are actually the top-level headings of the document. - # Rule 2 — Global consistency (highest priority among content rules): - # Headings that share the same structural pattern SHOULD receive the SAME level - # throughout the ENTIRE document, regardless of their position or textual content. - # (e.g., all "X.Y" two-part numbers must have the same level; all "X.Y.Z" - # three-part numbers must share a different, deeper level.) - # Rule 3 — Pattern over semantics: - # When determining a heading's level, its numbering/structural pattern takes - # precedence over its text length or semantic meaning. - # Parenthetical annotations or long descriptions inside a heading text do NOT - # indicate a different hierarchy level. - # Rule 4 — Parent-child continuity and no level skipping: - # Each heading must be consistent with adjacent headings. - # A heading may stay at the same level, return to an ancestor level, - # or go only ONE level deeper than its nearest valid ancestor heading. - # Level jumps such as level 1 directly to level 3 are invalid. - # Rule 5 — Body text demotion: - # If a row does not truly serve as a section title in the document outline, - # set its level to -1. - # Strong body-text cues include: - # - a full sentence or clause ending with sentence punctuation - # - an isolated broken word, broken phrase, label fragment, data value, or body continuation - # - a single Chinese character, digit, or very short fragment that clearly combines - # with the next row to form one continuous phrase rather than a standalone heading - # Rule 6 — Semantic heading promotion: - # A row with NO obvious numbering or structural-format markers can still - # be a heading, but ONLY when ALL of the following conditions are met: - # (a) The text is short and title-like (not a full sentence with punctuation). - # (b) It is NOT a broken fragment that simply continues into the next row - # (those belong to Rule 5 body-text demotion). - # (c) Multiple longer body-text rows follow it, and the row clearly - # organizes, summarizes, or introduces the topic of those rows — - # i.e., removing it would leave the following rows without a - # meaningful section label. - # Being short alone is NOT sufficient; the row must demonstrably serve - # as a section boundary that groups the content below it. - # When promoting, assign a level consistent with the surrounding - # hierarchy — typically one level deeper than the nearest heading above. - # - # **STEP 3 — Consistency check (one pass) before writing output** - # Scan the level assignments you are about to output and confirm: - # - All headings sharing the same structural or semantic pattern have been assigned the same level. - # If any inconsistency is found, normalise to the most representative level for that pattern. - # - # ***Output requirements*** - # - Output must be a [JSON array] only - # - **Only include rows that you judge to be headings** (level >= 1). Do NOT include body text rows (level = -1) in the output - # - Any row not present in your output will be automatically treated as body text (level = -1) - # - Each element must contain the following fields in order: - # - "id": original line number (integer) - # - "level": the corrected heading level (integer from 1 to {max_depth}) - # - # ***Format requirements*** - # - Output only valid JSON — do not add markdown fences (no ```json) - # - Do not add escaped newlines or other control characters - # - Do not add any explanations, comments, or descriptive text - # """ - elif task == "eval-headings": # COMPACT-input variant. Input is pre-compressed by `compact_for_llm` # so that consecutive body-text rows are folded into a single @@ -588,6 +445,228 @@ def build_prompt(task, texts, query, **kwargs): - No explanations, comments, or descriptive text """ + # ==================== Page-Memory Native Hierarchy Prompts ==================== + + elif task == "page-memory-vlm-tag": + temperature = 0 + top_p = 0.01 + max_tokens = kwargs.get("paras", {}).get("max_tokens", 600) + prompt = """\ +You are annotating a single PDF page screenshot for a document memory system. +Return strict JSON with exactly these keys: + +{ + "summary": "<1-3 sentence summary of the page content>", + "keywords": ";;" +} + +Rules: +- "summary": describe the main content visible on the page in 1-3 sentences. + If the page contains tables, mention the table topic and key columns. + If the page contains figures or charts, describe what they depict. +- "keywords": extract the most important thematic keywords (up to 5), + separated by semicolons ";". Keywords must be in the same language as + the visible page content. +- Return ONLY the JSON object, no markdown fences or extra text. +""" + + elif task == "page-memory-vlm-title": + temperature = 0 + top_p = 0.01 + max_tokens = kwargs.get("paras", {}).get("max_tokens", 300) + prompt = """\ +You are extracting document-outline-level headings from a PDF page screenshot. +Your goal is to find ONLY the headings that would appear in a Table of Contents. +Most pages will have ZERO such headings — returning an empty list is expected +and correct for the majority of pages. + +Return strict JSON: +{ + "titles": [ + { + "text": "", + "prominence": <0.0-1.0>, + "is_in_table": , + "is_in_header_footer": + } + ] +} + +═══ MANDATORY BOOLEAN FLAGS (CRITICAL) ═══ +For EVERY extracted heading, you MUST accurately evaluate these two flags: +1. is_in_table (boolean): Set to `true` if the text is ANYWHERE inside a table. +2. is_in_header_footer (boolean): Set to `true` if the text is located in the top margin (header) or bottom margin (footer) of the page. + +═══ WHAT TO EXTRACT ═══ + +Only extract text that satisfies ALL three criteria: + +1. HEADING FUNCTION (primary — must be true): + The text serves as a TITLE for the body content that follows it. + It introduces or labels a block of subsequent paragraphs, clauses, + or sub-sections. If you removed this text, the following body content + would lose its topic label. + +2. STANDALONE LINE (must be true): + The text occupies its own line, clearly separated from surrounding + body paragraphs. It is NOT inside a table, NOT part of a list, + and NOT embedded within a sentence. + +3. VISUAL DISTINCTION (supporting): + The text is visually set apart from body text — larger font, bold, + centered, or has extra vertical spacing. + +"prominence": 1.0 = most prominent; 0.5 = medium; 0.1 = minor. +Return titles in TOP-TO-BOTTOM order. Text must be EXACT verbatim. + +═══ WHAT TO EXCLUDE (critical — read carefully) ═══ + +1. TABLE CONTENT — Any text that is part of a table. If the + text is surrounded by grid lines, borders, or cell boundaries, or if + its neighboring content is arranged in rows and columns, it is table + content and MUST BE EXCLUDED. This applies even when the text is bold, + large, or spans a merged cell. Specifically exclude: + - Column headers, row category labels, merged-cell group labels + - Any label inside a tabular layout, regardless of visual prominence + +2. PAGE PERIPHERY — Text in margins or corners of the page: + organization/document names repeated as running headers, page numbers, + book/volume titles used as running headers or footers. + +3. BODY TEXT — Numbered clauses, list items, paragraphs, or running + prose, even if bold or indented. + +4. CAPTIONS — Figure/table captions, footnotes. + +5. TOC ENTRIES — If the page is itself a Table of Contents or index, + do NOT extract its listed entries. A TOC page lists other sections + with page numbers — those entries are references, not headings. + +═══ IMPORTANT ═══ +Many pages consist entirely of tables, body text, or appendix forms. +These pages have NO qualifying headings. Return {"titles": []} for them. +Do NOT force-extract table labels or body text as headings. + +Return ONLY the JSON object, no markdown fences. +""" + + elif task == "page-memory-hierarchy": + temperature = 0 + top_p = 0.01 + max_depth = kwargs["paras"].get("max_depth", 6) + max_tokens = kwargs["paras"].get("max_tokens", 2000) + coarse_context = kwargs["paras"].get("coarse_context", "") + coarse_section = f""" +Confirmed coarse parent section: +''' +{coarse_context} +''' +""" if coarse_context else "" + prompt = f""" +You are constructing a fine-grained document hierarchy for ONE already-bounded +PDF segment. The input rows are NOT raw body text. They are clean title +candidates observed directly from page screenshots by a VLM. + +Your task: +- Assign a relative hierarchy level to each real section/table/form heading. +- Level 1 means top-level inside this segment, level 2 is its child, etc. +- Preserve all legitimate sibling headings. Consecutive same-level headings are + normal and MUST NOT be demoted just because no body text appears between rows. +- Use page order as reading order. The "prominence" value is visual strength, + but numbering and structural pattern are more important. + +{coarse_section} + +Input rows: +''' +{texts} +''' + +Rules: +1. Trust structural numbering patterns first: + - "1", "2", "3" style headings at the same granularity are siblings. + - "3.1" is a child of "3"; "3.2.1" is a child of "3.2". + - "附录 A/B/C" are top-level siblings inside the segment unless the parent + context says otherwise. + - "表/附表" entries under an appendix are usually children of that appendix. +2. Do not skip levels. A child can be at most one level deeper than its nearest + valid parent. +3. Filter only obvious noise: + - duplicate/repeated variants of the same heading on adjacent rows; + - TOC/index headings such as "Contents", "目录", "目次"; + - front matter such as "前言" when it is outside the segment's body outline. +4. When two rows are near-duplicates, keep the clearer/more complete one and + omit the duplicate from output. +5. Do not invent headings. Only return ids that exist in the input. + +Output requirements: +- Output ONLY a valid JSON array. No markdown fences, no explanations. +- Include each retained heading as: + {{"id": , "level": }} +- Omitted ids are treated as filtered noise. +""" + + elif task == "page-memory-node-summary": + temperature = 0 + top_p = 0.01 + max_tokens = kwargs.get("paras", {}).get("max_tokens", 400) + node_title = kwargs.get("paras", {}).get("node_title", "") + next_title = kwargs.get("paras", {}).get("next_title", "") + kw_num = kwargs.get("paras", {}).get("kw_num", 5) + if next_title: + scope = ( + f"Summarize ONLY the content that belongs to the section titled " + f"\"{node_title}\". The section ends where the next section " + f"\"{next_title}\" begins on the page(s). Ignore everything that " + f"belongs to \"{next_title}\" or to other sections." + ) + else: + scope = ( + f"Summarize the content of the section titled \"{node_title}\" " + f"across the provided page image(s) as a single coherent section." + ) + prompt = f"""\ +You are summarizing one section of a document for a navigation/memory system. +You are given the page screenshot(s) that this section spans. + +{scope} + +Return strict JSON with exactly these keys: +{{ + "summary": "<1-4 sentence summary of THIS section's content>", + "keywords": ";;..." +}} + +Rules: +- "summary": describe what this section is about, in the same language as the + visible page content. If the section is mostly a table, describe the table's + topic and key columns. Do not summarize content that belongs to other + sections on the same page. +- "keywords": up to {kw_num} thematic keywords separated by ";". +- Return ONLY the JSON object, no markdown fences or extra text. +""" + + elif task == "page-memory-vlm-ocr": + temperature = 0 + top_p = 0.01 + max_tokens = kwargs.get("paras", {}).get("max_tokens", 1500) + prompt = """\ +You are transcribing a single scanned PDF page screenshot for a document +memory system. Extract the page's body text as faithfully as possible. + +Return strict JSON with exactly this key: +{ + "text": "" +} + +Rules: +- Preserve the reading order (top-to-bottom, left-to-right). +- Transcribe tables row by row using a simple readable layout. +- Do NOT add commentary, translation, or summary — transcription only. +- Omit pure decorative running headers/footers and page numbers. +- Return ONLY the JSON object, no markdown fences or extra text. +""" + # ==================== Image Processing Prompts ==================== elif task == "summary-images": diff --git a/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py index ece59d38..5f7b6a9a 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -107,7 +107,8 @@ def format_chunks( metadata["tokens"] = [] elif chunk_type == "page": # Page chunks carry page_image_uri and section_roles in - # extra_metadata; minimal metadata here. + # extra_metadata; preserve them for page-memory navigation. + _merge_page_chunk_metadata(metadata, existing_metadata) metadata["keywords"] = existing_metadata.get("keywords") or [] metadata["tokens"] = [] @@ -129,6 +130,34 @@ def _normalize_chunk_type(value: Any) -> str: return raw_type.split("\n", 1)[0].lower() +_PAGE_CHUNK_METADATA_KEYS = ( + "granularity", + "page_index", + "page_indices", + "owned_pages", + "section_path", + "section_level", + "page_image_uri", + "page_image_uris", + "strategy_used", + "kind", + "status", + "observed_titles", + "section_roles", + "source_verdict", + "native_hierarchy", +) + + +def _merge_page_chunk_metadata( + metadata: dict[str, Any], + existing_metadata: dict[str, Any], +) -> None: + for key in _PAGE_CHUNK_METADATA_KEYS: + if key in existing_metadata: + metadata[key] = existing_metadata[key] + + def _base_chunk_metadata( existing_metadata: dict[str, Any], chunk: dict[str, Any], diff --git a/packages/shared-python/shared/services/storage/zip_doc_navigation.py b/packages/shared-python/shared/services/storage/zip_doc_navigation.py index ae22f47d..9ad5870a 100644 --- a/packages/shared-python/shared/services/storage/zip_doc_navigation.py +++ b/packages/shared-python/shared/services/storage/zip_doc_navigation.py @@ -199,3 +199,187 @@ def _section_tree_to_output( } ) return result + + +def _coerce_int(value: Any) -> int | None: + try: + return int(value) + except (TypeError, ValueError): + return None + + +def build_doc_nav_from_skeletons( + skeletons: list[dict[str, Any]], + chunks: list[dict[str, Any]], + source_file_name: str, +) -> dict[str, Any]: + """Build doc_nav.json from SectionSkeleton dicts + page chunks. + + Unlike ``ZipDocNavigationBuilder.build_doc_nav`` which infers the tree + from chunk paths (losing sections whose start page is shared with a + sibling), this builder uses skeletons as the authoritative tree structure + and computes chunk_count by page-range overlap. + + Parameters + ---------- + skeletons: + List of ``SectionSkeleton.to_dict()`` dicts. Must have + ``section_path``, ``title``, ``level``, ``start_page``, ``end_page``, + and ``parent_path``. + chunks: + Formatted chunk dicts (from ``dataframe_to_chunks``). Only chunks + with ``type == "page"`` are counted; their ``metadata.page_nums[0]`` + identifies the page index. + source_file_name: + Original filename for the doc_nav envelope. + """ + page_indices: set[int] = set() + chunk_summaries: dict[int, str] = {} + for chunk in chunks: + meta = chunk.get("metadata") or {} + page_nums = meta.get("page_nums") or [] + if page_nums: + page = _coerce_int(page_nums[0]) + if page is None: + continue + page_indices.add(page) + if not chunk_summaries.get(page): + chunk_summaries[page] = ( + (meta.get("summary") or "").strip() + or (chunk.get("content") or "").strip()[:200] + ) + + sorted_skels = sorted( + skeletons, + key=lambda s: ( + _coerce_int(s.get("start_page")) or 0, + _coerce_int(s.get("level")) or 0, + str(s.get("section_path") or ""), + ), + ) + nodes: dict[str, dict[str, Any]] = {} + root_children_order: list[str] = [] + + def _link_child(parent_path: str, child_path: str) -> None: + if not parent_path: + if child_path not in root_children_order: + root_children_order.append(child_path) + return + parent = nodes.get(parent_path) + if parent is None: + return + children = parent["_children_paths"] + if child_path not in children: + children.append(child_path) + + def _ensure_path_node(path: str) -> str: + root_parts, section_parts = split_document_path( + path, + source_file_name=source_file_name, + ) + if not section_parts: + return "" + + parent_path = "" + path_parts = list(root_parts) + for index, title in enumerate(section_parts): + path_parts.append(title) + current_path = "/".join(path_parts) + if current_path not in nodes: + nodes[current_path] = { + "title": title, + "path": current_path, + "level": index + 1, + "start_page": None, + "end_page": None, + "summary": "", + "parent_path": parent_path, + "owned_pages": set(), + "_children_paths": [], + } + _link_child(parent_path, current_path) + parent_path = current_path + return parent_path + + for skel in sorted_skels: + sp = str(skel.get("section_path") or "").strip() + if not sp: + continue + current_path = _ensure_path_node(sp) + if not current_path: + continue + + node = nodes[current_path] + start_page = _coerce_int(skel.get("start_page")) + end_page = _coerce_int(skel.get("end_page")) + if start_page is not None and end_page is not None: + node["start_page"] = start_page + node["end_page"] = end_page + node["owned_pages"].update(page_indices & set(range(start_page, end_page + 1))) + node["summary"] = chunk_summaries.get(start_page, node["summary"]) + node["title"] = str(skel.get("title") or node["title"]) + node["level"] = _coerce_int(skel.get("level")) or node["level"] + + def _collect_owned_pages(path: str) -> set[int]: + node = nodes[path] + pages = set(node["owned_pages"]) + for child_path in node["_children_paths"]: + pages.update(_collect_owned_pages(child_path)) + node["owned_pages"] = pages + return pages + + for path in root_children_order: + if path in nodes: + _collect_owned_pages(path) + + def _to_output(node: dict[str, Any], level: int = 1) -> dict[str, Any]: + children_out = [ + _to_output(nodes[child_path], level + 1) + for child_path in node["_children_paths"] + if child_path in nodes + ] + return { + "title": node["title"], + "path": node["path"], + "level": level, + "summary": node["summary"], + "chunk_count": len(node["owned_pages"]), + "children": children_out, + } + + sections = [_to_output(nodes[sp]) for sp in root_children_order if sp in nodes] + + # Stats + stats = { + "total_chunks": len(chunks), + "text_chunks": sum(1 for c in chunks if c.get("type") == "text"), + "image_chunks": sum(1 for c in chunks if c.get("type") == "image"), + "table_chunks": sum(1 for c in chunks if c.get("type") == "table"), + "page_chunks": sum(1 for c in chunks if c.get("type") == "page"), + "max_depth": _max_depth(sections), + } + + # Resources + image_resources = [] + table_resources = [] + for chunk in chunks: + ct = chunk.get("type", "") + path = chunk.get("path", "") + meta = chunk.get("metadata") or {} + summary = (meta.get("summary") or "").strip() + if ct == "image": + image_resources.append({"path": path, "summary": summary}) + elif ct == "table": + table_resources.append({"path": path, "summary": summary}) + + return { + "version": "1.0", + "file_name": source_file_name or "", + "stats": stats, + "sections": sections, + "resources": { + "images": image_resources, + "tables": table_resources, + }, + } + diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py index e444475a..667cab70 100644 --- a/packages/shared-python/shared/services/storage/zip_result_schema.py +++ b/packages/shared-python/shared/services/storage/zip_result_schema.py @@ -5,7 +5,10 @@ from typing import Any from shared.services.storage.zip_chunk_schema import ZipChunkSchemaBuilder -from shared.services.storage.zip_doc_navigation import ZipDocNavigationBuilder +from shared.services.storage.zip_doc_navigation import ( + ZipDocNavigationBuilder, + build_doc_nav_from_skeletons, +) from shared.services.storage.zip_manifest_schema import ZipManifestBuilder @@ -70,3 +73,15 @@ def build_doc_nav( formatted_chunks, source_file_name, ) + + def build_doc_nav_from_skeletons( + self, + skeletons: list[dict[str, Any]], + formatted_chunks: list[dict[str, Any]], + source_file_name: str, + ) -> dict[str, Any]: + return build_doc_nav_from_skeletons( + skeletons, + formatted_chunks, + source_file_name, + ) diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index b2fee6e1..20e83e25 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -71,6 +71,7 @@ def generate_zip_package( doc_nav, hierarchy = self._build_navigation_outputs( formatted_chunks=formatted_chunks, source_file_name=source_file_name, + parsed_df=parsed_df, ) manifest = self._schema.generate_manifest( job_id=job_id, @@ -121,11 +122,33 @@ def _build_navigation_outputs( *, formatted_chunks: list[dict[str, Any]], source_file_name: str, + parsed_df: pd.DataFrame | None = None, ) -> tuple[dict[str, Any] | None, dict[str, Any]]: try: - doc_nav = self._schema.build_doc_nav(formatted_chunks, source_file_name) + skeletons = _get_page_memory_skeletons(parsed_df) + if skeletons: + doc_nav = self._schema.build_doc_nav_from_skeletons( + skeletons, + formatted_chunks, + source_file_name, + ) + else: + doc_nav = self._schema.build_doc_nav(formatted_chunks, source_file_name) hierarchy = self._schema.build_hierarchy_dict(doc_nav.get("sections", [])) return doc_nav, hierarchy except Exception as exc: logger.warning(f"generate doc_nav.json fail {exc}") return None, {} + + +def _get_page_memory_skeletons( + parsed_df: pd.DataFrame | None, +) -> list[dict[str, Any]]: + if parsed_df is None: + return [] + + raw_skeletons = parsed_df.attrs.get("page_memory_skeletons") + if not isinstance(raw_skeletons, list): + return [] + + return [skel for skel in raw_skeletons if isinstance(skel, dict)] From 0035b4afd7cabe77206751382f5108a3d521cccf Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 23 Jun 2026 15:13:16 +0800 Subject: [PATCH 02/27] feat: implement page asset integration into node assembly and chunk connections with support for custom probes --- apps/api/.env.example | 2 + apps/worker/.env.example | 12 + .../services/page_memory/memory_service.py | 48 +- .../services/page_memory/node_assembler.py | 153 +++-- .../app/services/page_memory/page_assets.py | 537 ++++++++++++++++++ apps/worker/experiments/chart_asset_probe.py | 452 +++++++++++++++ .../experiments/vlm_bbox_tabula_probe.py | 294 ++++++++++ apps/worker/pyproject.toml | 1 + apps/worker/requirements.txt | 18 +- ...est_page_memory_node_assembler_contract.py | 113 +++- .../shared-python/shared/core/config/ai.py | 11 +- .../shared/services/ai/prompt_service.py | 48 ++ .../services/chunks/chunk_connections.py | 4 + .../chunks/dataframe_chunk_converter.py | 2 +- uv.lock | 16 + 15 files changed, 1662 insertions(+), 49 deletions(-) create mode 100644 apps/worker/app/services/page_memory/page_assets.py create mode 100644 apps/worker/experiments/chart_asset_probe.py create mode 100644 apps/worker/experiments/vlm_bbox_tabula_probe.py diff --git a/apps/api/.env.example b/apps/api/.env.example index 2db79471..02795f9f 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -81,6 +81,8 @@ ARK_API_KEY= # ARK_URL=https://ark.cn-beijing.volces.com/api/v3/chat/completions # NORMOL_MODEL=deepseek-v4-flash # HIERARCHY_LLM_MODEL= +# VLM for image/OCR flows. Default qwen3.6-flash; alternate qwen3-vl-32b-instruct +# (open-weights, self-hostable). DashScope pricing: help.aliyun.com/zh/model-studio/model-pricing # IMAGE_MODEL=qwen3.6-flash # IMAGE_MODEL_MAX=qwen3.6-flash diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 2b7a622f..36658855 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -86,6 +86,10 @@ ARK_API_KEY= # ARK_URL=https://ark.cn-beijing.volces.com/api/v3/chat/completions # NORMOL_MODEL=deepseek-v4-flash # HIERARCHY_LLM_MODEL= +# VLM for page tagging, OCR, chart/table bbox probes, and image flows. +# Default: qwen3.6-flash (DashScope China ~¥1.2 in / ¥7.2 out per 1M tokens). +# Alternate: qwen3-vl-32b-instruct — open-weights, self-hostable (vLLM/SGLang); +# DashScope China ~¥2 in / ¥8 out per 1M tokens (see model-pricing doc). # IMAGE_MODEL=qwen3.6-flash # IMAGE_MODEL_MAX=qwen3.6-flash @@ -122,6 +126,14 @@ OVERSIZED_PDF_SHARD_ENABLED=true OVERSIZED_PDF_SOFT_LIMIT=1500 PDF_PAGE_TOC_ENABLED=true RETRIEVAL_PAGE_MEMORY_ENABLED=false +# Page-memory asset extraction is opt-in. Requires VLM credentials; table HTML +# extraction additionally requires tabula-py and a Java runtime. +PAGE_MEMORY_ASSET_EXTRACTION_ENABLED=false +# PAGE_MEMORY_ASSET_MODEL= +# PAGE_MEMORY_ASSET_MAX_PAGES=50 +# PAGE_MEMORY_ASSET_BUDGET= +# PAGE_MEMORY_ASSET_CONFIDENCE_THRESHOLD=0.3 +# PAGE_MEMORY_TABLE_ENGINE=tabula MINERU_SHARD_CONCURRENCY=3 PARSE_AGENT_PLAN_BUDGET=50000 PARSE_AGENT_VISUAL_BUDGET=80000 diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 77d59dc3..02c0c37f 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -122,6 +122,13 @@ def _build_page_dataframe( from app.services.page_memory.page_section_mapper import map_pages_to_sections from app.services.page_memory.page_tagger import tag_pages from app.services.page_memory.skeleton_extractor import extract_section_skeletons + from app.services.page_memory.page_assets import ( + extract_page_assets_from_renders, + get_asset_budget, + get_asset_confidence_threshold, + get_asset_max_pages, + page_asset_extraction_enabled, + ) # ── Native hierarchy flag (Step 3) ──────────────────────────────── # Page-memory handles PDF/PPT(→PDF) structure with its own page-native @@ -143,10 +150,21 @@ def _build_page_dataframe( page_locate_budget = int( os.environ.get("PAGE_MEMORY_LOCATE_BUDGET", str(min(page_count * 1600, 2_000_000))) ) - title_detection_budget = int( - os.environ.get("PAGE_MEMORY_TITLE_BUDGET", str(page_count * 800)) - ) if native_hierarchy else 0 - total_visual = page_tagging_budget + page_locate_budget + title_detection_budget + title_detection_budget = ( + int(os.environ.get("PAGE_MEMORY_TITLE_BUDGET", str(page_count * 800))) + if native_hierarchy + else 0 + ) + asset_extraction_enabled = page_asset_extraction_enabled() + asset_extraction_budget = ( + get_asset_budget(page_count) if asset_extraction_enabled else 0 + ) + total_visual = ( + page_tagging_budget + + page_locate_budget + + title_detection_budget + + asset_extraction_budget + ) # plan_budget powers the LLM decision loop inside PageLocateSubAgent. # Without it the sub-agent falls back to a deterministic path that @@ -171,6 +189,11 @@ def _build_page_dataframe( min_guarantee=title_detection_budget, cap=None, ) + if asset_extraction_enabled: + stage_envelopes["page_asset_extraction"] = StageEnvelope( + min_guarantee=asset_extraction_budget, + cap=None, + ) budget = BudgetTracker( plan_budget=plan_budget, visual_budget=total_visual, @@ -214,6 +237,21 @@ def _build_page_dataframe( page_texts=page_texts, ) + page_assets_by_page: dict[int, list[Any]] = {} + if asset_extraction_enabled: + asset_model = os.environ.get("PAGE_MEMORY_ASSET_MODEL") or os.environ.get( + "IMAGE_MODEL" + ) + page_assets_by_page = extract_page_assets_from_renders( + pdf_path=pdf_path, + rendered_pages=rendered, + output_dir=output_dir, + model_name=asset_model, + budget=budget, + max_pages=get_asset_max_pages(page_count), + confidence_threshold=get_asset_confidence_threshold(), + ) + # ── C2: page plan ───────────────────────────────────────────────── page_labels = anatomy.page_labels if anatomy else [] plans = derive_page_processing_plan( @@ -339,6 +377,7 @@ def _build_page_dataframe( native_hierarchy=native_hierarchy, budget=budget, vlm_model=vlm_model, + page_assets_by_page=page_assets_by_page, ) logger.info( "[page_memory] C7 assembled {} node rows (verdict={})", @@ -525,4 +564,3 @@ def _render_page_images( for item in rendered if item.get("png_path") ] - diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index 1e2a52e7..35f61646 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -31,6 +31,11 @@ from app.services.document_agent.budget import BudgetTracker from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.page_memory.page_assets import ( + PageAsset, + asset_reference, + build_asset_rows, +) from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton from shared.services.ai.prompt_service import build_prompt @@ -73,11 +78,7 @@ def identify_leaf_nodes(skeletons: list[SectionSkeleton]) -> list[LeafNode]: Reading order = ``(start_page, original index)`` so that siblings sharing a start page keep the top-to-bottom order produced by fine hierarchy. """ - parent_paths = { - skel.parent_path - for skel in skeletons - if skel.parent_path - } + parent_paths = {skel.parent_path for skel in skeletons if skel.parent_path} leaves: list[tuple[int, int, LeafNode]] = [] for index, skel in enumerate(skeletons): if skel.section_path in parent_paths: @@ -237,7 +238,9 @@ def resolve_page_text( "page-memory-vlm-ocr", "", "", paras={"max_tokens": 1500} ) est = estimate_tokens(prompt) + 1000 - if budget is not None and not budget.try_reserve("visual", est, stage=_BUDGET_STAGE): + if budget is not None and not budget.try_reserve( + "visual", est, stage=_BUDGET_STAGE + ): logger.debug("[node_assembler] OCR budget exhausted for page {}", page) return "" @@ -267,8 +270,12 @@ def resolve_page_text( est=est, stage=_BUDGET_STAGE, ) - data = json.loads(raw_response) - return str(data.get("text", "")).strip() + try: + data = json.loads(raw_response) + return str(data.get("text", "")).strip() + except (json.JSONDecodeError, TypeError): + # VLM returned non-JSON; use raw text directly + return raw_response.strip() if raw_response else "" except Exception as exc: logger.warning("[node_assembler] OCR failed for page {}: {}", page, exc) if budget is not None: @@ -390,7 +397,9 @@ def _vlm_node_summary( }, ) est = estimate_tokens(prompt) + 800 * len(image_parts) - if budget is not None and not budget.try_reserve("visual", est, stage=_BUDGET_STAGE): + if budget is not None and not budget.try_reserve( + "visual", est, stage=_BUDGET_STAGE + ): logger.debug( "[node_assembler] node summary budget exhausted for {}", leaf.section_path, @@ -449,6 +458,7 @@ def build_node_rows( native_hierarchy: bool, budget: BudgetTracker | None = None, vlm_model: str | None = None, + page_assets_by_page: dict[int, list[PageAsset]] | None = None, ) -> list[dict[str, Any]]: """Assemble one row per leaf section node (node-granularity chunks).""" available_pages = set(raw_text_by_page.keys()) @@ -469,6 +479,7 @@ def build_node_rows( ) rows: list[dict[str, Any]] = [] + rows_by_path: dict[str, dict[str, Any]] = {} for view in views: leaf = view.leaf content = build_node_content( @@ -490,37 +501,107 @@ def build_node_rows( if image_uri_by_page.get(page) ] know_id = f"node_{gen_str_codes(f'{filename}::{leaf.section_path}')}" - rows.append( - { - "content": content, - "path": leaf.section_path, - "type": "page", - "length": len(content), - "keywords": ";".join(keywords), - "summary": summary, - "know_id": know_id, - "tokens": "", - "connectto": "", - "addtime": get_str_time(), - "page_nums": ",".join(str(page) for page in view.pages), - "extra_metadata": { - "granularity": "node", - "section_path": leaf.section_path, - "section_level": leaf.level, - "page_indices": list(view.pages), - "owned_pages": list(view.owned_pages), - "page_image_uris": page_image_uris, - "kind": kind_by_page.get(leaf.start_page, "normal"), - "source_verdict": verdict, - "native_hierarchy": native_hierarchy, - }, - } + row = { + "content": content, + "path": leaf.section_path, + "type": "page", + "length": len(content), + "keywords": ";".join(keywords), + "summary": summary, + "know_id": know_id, + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": ",".join(str(page) for page in view.pages), + "extra_metadata": { + "granularity": "node", + "section_path": leaf.section_path, + "section_level": leaf.level, + "page_indices": list(view.pages), + "owned_pages": list(view.owned_pages), + "page_image_uris": page_image_uris, + "kind": kind_by_page.get(leaf.start_page, "normal"), + "source_verdict": verdict, + "native_hierarchy": native_hierarchy, + }, + } + rows.append(row) + rows_by_path[leaf.section_path] = row + + asset_rows: list[dict[str, Any]] = [] + if page_assets_by_page: + asset_rows = build_asset_rows(page_assets_by_page) + _attach_asset_connections( + page_assets_by_page=page_assets_by_page, + page_owner=page_owner, + page_to_leaves=page_to_leaves, + rows_by_path=rows_by_path, ) logger.info( - "[node_assembler] assembled {} node rows from {} leaves ({} pages)", + "[node_assembler] assembled {} asset rows + {} node rows from {} leaves ({} pages)", + len(asset_rows), len(rows), len(leaves), len(available_pages), ) - return rows + return asset_rows + rows + + +def _attach_asset_connections( + *, + page_assets_by_page: dict[int, list[PageAsset]], + page_owner: dict[int, LeafNode], + page_to_leaves: dict[int, list[LeafNode]], + rows_by_path: dict[str, dict[str, Any]], +) -> None: + for page_index, assets in page_assets_by_page.items(): + owner_leaf = page_owner.get(page_index) + leaves_on_page = page_to_leaves.get(page_index, []) + for asset in assets: + ref = asset_reference(asset) + if not ref: + continue + if owner_leaf is not None: + owner_row = rows_by_path.get(owner_leaf.section_path) + if owner_row is not None: + _append_connect_to( + owner_row, + { + "target": ref, + "relation": "embeds", + "ref": ref, + }, + ) + for leaf in leaves_on_page: + if ( + owner_leaf is not None + and leaf.section_path == owner_leaf.section_path + ): + continue + row = rows_by_path.get(leaf.section_path) + if row is None: + continue + connection: dict[str, Any] = { + "target": ref, + "relation": "related", + "ref": ref, + } + if owner_leaf is not None: + connection["same_as_owner"] = owner_leaf.section_path + _append_connect_to(row, connection) + + +def _append_connect_to(row: dict[str, Any], connection: dict[str, Any]) -> None: + existing = row.get("connectto") + if isinstance(existing, str) and existing.strip(): + try: + connections = json.loads(existing) + except json.JSONDecodeError: + connections = [] + elif isinstance(existing, list): + connections = list(existing) + else: + connections = [] + connections.append(connection) + row["connectto"] = json.dumps(connections, ensure_ascii=False) diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py new file mode 100644 index 00000000..b1f5ffd6 --- /dev/null +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -0,0 +1,537 @@ +"""Page-memory asset extraction for charts, figures, and tables.""" + +from __future__ import annotations + +import base64 +import importlib +import json +import os +import shutil +import subprocess +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, cast + +import pandas as pd +from loguru import logger +from PIL import Image + +from app.services.document_agent.budget import BudgetTracker +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.page_memory.page_renderer import PageRenderResult +from shared.services.ai.prompt_service import build_prompt +from shared.utils.token_estimate import estimate_tokens + +_BUDGET_STAGE = "page_asset_extraction" +_GRID_SIZE = 1000 +_VALID_KINDS = {"table", "chart", "figure"} + + +@dataclass +class PageAsset: + asset_id: str + page_index: int + asset_index: int + kind: str + bbox_px: list[int] + width_px: int + height_px: int + width_pt: float + height_pt: float + caption: str = "" + confidence: float = 0.0 + title: str = "" + summary: str = "" + keywords: list[str] = field(default_factory=list) + image_uri: str = "" + html_uri: str = "" + image_path: str = "" + html_path: str = "" + extraction_status: str = "pending" + + +def page_asset_extraction_enabled() -> bool: + return os.environ.get( + "PAGE_MEMORY_ASSET_EXTRACTION_ENABLED", "false" + ).strip().lower() in {"1", "true", "yes", "on"} + + +def get_asset_confidence_threshold() -> float: + try: + return float(os.environ.get("PAGE_MEMORY_ASSET_CONFIDENCE_THRESHOLD", "0.3")) + except ValueError: + return 0.3 + + +def get_asset_max_pages(page_count: int) -> int: + try: + value = int(os.environ.get("PAGE_MEMORY_ASSET_MAX_PAGES", "50")) + except ValueError: + value = 50 + return max(0, min(page_count, value)) + + +def get_asset_budget(page_count: int) -> int: + default_budget = str(page_count * 600) + try: + return max(0, int(os.environ.get("PAGE_MEMORY_ASSET_BUDGET", default_budget))) + except ValueError: + return max(0, page_count * 600) + + +def detect_page_assets( + *, + page: PageRenderResult, + source_name: str, + model_name: str | None, + budget: BudgetTracker | None, + confidence_threshold: float, +) -> list[PageAsset]: + """Detect asset regions on one rendered page via VLM.""" + if not model_name or not page.image_path or not os.path.exists(page.image_path): + return [] + + try: + with Image.open(page.image_path) as image: + width_px, height_px = image.size + with open(page.image_path, "rb") as handle: + image_b64 = base64.b64encode(handle.read()).decode() + except Exception as exc: + logger.warning( + "[page_assets] failed to read rendered page {}: {}", + page.page_index, + exc, + ) + return [] + + prompt, temperature, _top_p, max_tokens = build_prompt( + "page-memory-asset-detect", + "", + "", + paras={"max_tokens": 1200, "grid_size": _GRID_SIZE}, + ) + est = estimate_tokens(prompt) + 1000 + if budget is not None and not budget.try_reserve( + "visual", est, stage=_BUDGET_STAGE + ): + logger.debug( + "[page_assets] asset budget exhausted before page {}", page.page_index + ) + return [] + + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model_name) + raw_response, usage = client.chat_completion_with_usage( + messages=cast( + Any, + [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{image_b64}" + }, + }, + ], + } + ], + ), + model=model_name, + temperature=temperature, + max_tokens=max_tokens, + response_format={"type": "json_object"}, + usage_task="page_memory.asset_detect", + ) + if budget is not None: + budget.commit( + "visual", + actual=usage.get("total_tokens", est), + est=est, + stage=_BUDGET_STAGE, + ) + data = json.loads(raw_response) + except Exception as exc: + logger.warning( + "[page_assets] VLM asset detection failed on page {}: {}", + page.page_index, + exc, + ) + if budget is not None: + budget.refund("visual", est=est, stage=_BUDGET_STAGE) + return [] + + regions = data.get("regions") if isinstance(data, dict) else None + if not isinstance(regions, list): + return [] + + assets: list[PageAsset] = [] + for region in regions: + if not isinstance(region, dict): + continue + kind = str(region.get("kind") or region.get("type") or "").strip().lower() + if kind not in _VALID_KINDS: + continue + confidence = _safe_float(region.get("confidence"), default=0.0) + if confidence < confidence_threshold: + continue + bbox_px = _norm_bbox_to_px(region.get("bbox"), width_px, height_px) + if bbox_px is None: + continue + asset_index = len(assets) + 1 + asset_id = "asset_" + gen_str_codes( + f"{source_name}::{page.page_index}::{asset_index}::{kind}::{bbox_px}::" + f"{region.get('title') or ''}::{region.get('summary') or ''}" + ) + keywords = _normalize_keywords(region.get("keywords")) + assets.append( + PageAsset( + asset_id=asset_id, + page_index=page.page_index, + asset_index=asset_index, + kind=kind, + bbox_px=bbox_px, + width_px=width_px, + height_px=height_px, + width_pt=float(page.width or 0.0), + height_pt=float(page.height or 0.0), + caption=str(region.get("caption") or "").strip(), + confidence=confidence, + title=str(region.get("title") or "").strip(), + summary=str(region.get("summary") or "").strip(), + keywords=keywords, + extraction_status="detected", + ) + ) + return assets + + +def crop_page_asset( + *, + asset: PageAsset, + page_image_path: str, + output_dir: str, + margin_px: int = 4, +) -> PageAsset: + image_dir = Path(output_dir) / "images" + image_dir.mkdir(parents=True, exist_ok=True) + filename = f"image_page_{asset.page_index}_{asset.kind}_{asset.asset_index}.png" + output_path = image_dir / filename + try: + with Image.open(page_image_path) as image: + x1, y1, x2, y2 = asset.bbox_px + box = ( + max(0, x1 - margin_px), + max(0, y1 - margin_px), + min(image.width, x2 + margin_px), + min(image.height, y2 + margin_px), + ) + image.crop(box).save(output_path) + asset.image_path = str(output_path) + asset.image_uri = f"images/{filename}" + if asset.extraction_status == "detected": + asset.extraction_status = "cropped" + except Exception as exc: + logger.warning( + "[page_assets] crop failed for page {} asset {}: {}", + asset.page_index, + asset.asset_index, + exc, + ) + asset.extraction_status = "crop_failed" + return asset + + +def extract_table_html( + *, + asset: PageAsset, + pdf_path: str, + output_dir: str, +) -> PageAsset: + """Extract table HTML with tabula stream mode; degrade on missing Java/package.""" + if asset.kind != "table": + return asset + if os.environ.get("PAGE_MEMORY_TABLE_ENGINE", "tabula").strip().lower() != "tabula": + asset.extraction_status = f"{asset.extraction_status}:table_engine_disabled" + return asset + if not _has_working_java(): + asset.extraction_status = f"{asset.extraction_status}:java_unavailable" + return asset + + try: + tabula = cast(Any, importlib.import_module("tabula")) + except Exception as exc: + logger.warning("[page_assets] tabula-py unavailable: {}", exc) + asset.extraction_status = f"{asset.extraction_status}:tabula_unavailable" + return asset + + area_pt = _bbox_px_to_area_pt( + asset.bbox_px, + width_px=asset.width_px, + height_px=asset.height_px, + width_pt=asset.width_pt, + height_pt=asset.height_pt, + ) + if area_pt is None: + asset.extraction_status = f"{asset.extraction_status}:missing_page_dimensions" + return asset + + try: + frames = cast( + list[Any], + tabula.read_pdf( + pdf_path, + pages=asset.page_index, + area=area_pt, + guess=False, + lattice=False, + stream=True, + multiple_tables=True, + pandas_options={"header": None}, + ), + ) + except Exception as exc: + logger.warning( + "[page_assets] tabula extraction failed page {} asset {}: {}", + asset.page_index, + asset.asset_index, + exc, + ) + asset.extraction_status = f"{asset.extraction_status}:tabula_failed" + return asset + + frame = next( + ( + item + for item in frames + if isinstance(item, pd.DataFrame) and _is_meaningful_frame(item) + ), + None, + ) + if frame is None: + asset.extraction_status = f"{asset.extraction_status}:tabula_empty" + return asset + + table_dir = Path(output_dir) / "tables" + table_dir.mkdir(parents=True, exist_ok=True) + filename = f"table_page_{asset.page_index}_{asset.asset_index}.html" + output_path = table_dir / filename + html = frame.fillna("").astype(str).to_html(index=False, header=False) + output_path.write_text(html, encoding="utf-8") + asset.html_path = str(output_path) + asset.html_uri = f"tables/{filename}" + asset.extraction_status = "table_html_extracted" + return asset + + +def extract_page_assets_from_renders( + *, + pdf_path: str, + rendered_pages: list[PageRenderResult], + output_dir: str, + model_name: str | None, + budget: BudgetTracker | None, + max_pages: int, + confidence_threshold: float, +) -> dict[int, list[PageAsset]]: + pages = sorted(rendered_pages, key=lambda item: item.page_index)[:max_pages] + source_name = Path(pdf_path).name + assets_by_page: dict[int, list[PageAsset]] = {} + for page in pages: + detected = detect_page_assets( + page=page, + source_name=source_name, + model_name=model_name, + budget=budget, + confidence_threshold=confidence_threshold, + ) + page_assets: list[PageAsset] = [] + for asset in detected: + crop_page_asset( + asset=asset, + page_image_path=page.image_path, + output_dir=output_dir, + ) + if asset.kind == "table": + extract_table_html( + asset=asset, + pdf_path=pdf_path, + output_dir=output_dir, + ) + if asset.image_uri or asset.html_uri: + page_assets.append(asset) + if page_assets: + assets_by_page[page.page_index] = page_assets + logger.info( + "[page_assets] extracted {} assets across {} pages", + sum(len(items) for items in assets_by_page.values()), + len(assets_by_page), + ) + return assets_by_page + + +def build_asset_rows( + page_assets_by_page: dict[int, list[PageAsset]], +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for page_index in sorted(page_assets_by_page): + for asset in page_assets_by_page[page_index]: + ref_uri = ( + asset.html_uri + if asset.kind == "table" and asset.html_uri + else asset.image_uri + ) + if not ref_uri: + continue + row_type = "table" if asset.kind == "table" and asset.html_uri else "image" + content = _asset_content(asset, row_type=row_type) + summary = "\n".join( + part for part in [asset.title.strip(), asset.summary.strip()] if part + ) + rows.append( + { + "content": content, + "path": ref_uri, + "type": row_type, + "length": len(content), + "keywords": ";".join(asset.keywords), + "summary": summary, + "know_id": asset.asset_id, + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": str(asset.page_index), + "extra_metadata": { + "page_index": asset.page_index, + "asset_index": asset.asset_index, + "asset_kind": asset.kind, + "bbox_px": asset.bbox_px, + "confidence": asset.confidence, + "caption": asset.caption, + "image_uri": asset.image_uri, + "html_uri": asset.html_uri, + "extraction_status": asset.extraction_status, + "table_engine": "tabula" if asset.kind == "table" else "", + }, + } + ) + return rows + + +def asset_reference(asset: PageAsset) -> str: + uri = ( + asset.html_uri if asset.kind == "table" and asset.html_uri else asset.image_uri + ) + return f"[{uri}]" if uri else "" + + +def _asset_content(asset: PageAsset, *, row_type: str) -> str: + if row_type == "table" and asset.html_path: + try: + return Path(asset.html_path).read_text(encoding="utf-8") + except Exception as exc: + logger.warning( + "[page_assets] failed to read table HTML {}: {}", asset.html_path, exc + ) + parts = [asset.title, asset.summary, asset.caption] + if asset.image_uri: + parts.append(f"[{asset.image_uri}]") + return "\n".join(part.strip() for part in parts if part and part.strip()) + + +def _norm_bbox_to_px(value: object, width_px: int, height_px: int) -> list[int] | None: + if not isinstance(value, list) or len(value) != 4: + return None + try: + x1, y1, x2, y2 = [float(item) for item in value] + except (TypeError, ValueError): + return None + if x2 <= x1 or y2 <= y1: + return None + x1_px = round(max(0.0, min(_GRID_SIZE, x1)) / _GRID_SIZE * width_px) + x2_px = round(max(0.0, min(_GRID_SIZE, x2)) / _GRID_SIZE * width_px) + y1_px = round(max(0.0, min(_GRID_SIZE, y1)) / _GRID_SIZE * height_px) + y2_px = round(max(0.0, min(_GRID_SIZE, y2)) / _GRID_SIZE * height_px) + if x2_px <= x1_px or y2_px <= y1_px: + return None + return [x1_px, y1_px, x2_px, y2_px] + + +def _bbox_px_to_area_pt( + bbox_px: list[int], + *, + width_px: int, + height_px: int, + width_pt: float, + height_pt: float, + margin_pt: float = 2.0, +) -> list[float] | None: + if width_px <= 0 or height_px <= 0 or width_pt <= 0 or height_pt <= 0: + return None + x1, y1, x2, y2 = bbox_px + left = x1 / width_px * width_pt + right = x2 / width_px * width_pt + top = y1 / height_px * height_pt + bottom = y2 / height_px * height_pt + return [ + round(max(0.0, top - margin_pt), 2), + round(max(0.0, left - margin_pt), 2), + round(min(height_pt, bottom + margin_pt), 2), + round(min(width_pt, right + margin_pt), 2), + ] + + +def _has_working_java() -> bool: + if shutil.which("java") is None: + return False + try: + result = subprocess.run( + ["java", "-version"], + check=False, + capture_output=True, + text=True, + timeout=10, + ) + except Exception: + return False + combined = f"{result.stdout}\n{result.stderr}" + return result.returncode == 0 and "Unable to locate a Java Runtime" not in combined + + +def _is_meaningful_frame(frame: pd.DataFrame) -> bool: + if frame.empty or frame.shape[1] == 0: + return False + non_empty = frame.fillna("").astype(str).map(lambda text: bool(text.strip())) + return bool(non_empty.values.sum() >= 2) + + +def _normalize_keywords(value: object) -> list[str]: + if isinstance(value, list): + return [str(item).strip() for item in value if str(item).strip()] + if isinstance(value, str): + delimiter = ";" if ";" in value else "," + return [item.strip() for item in value.split(delimiter) if item.strip()] + return [] + + +def _safe_float(value: object, *, default: float) -> float: + try: + return float(str(value)) + except (TypeError, ValueError): + return default + + +__all__ = [ + "PageAsset", + "asset_reference", + "build_asset_rows", + "extract_page_assets_from_renders", + "get_asset_budget", + "get_asset_confidence_threshold", + "get_asset_max_pages", + "page_asset_extraction_enabled", +] diff --git a/apps/worker/experiments/chart_asset_probe.py b/apps/worker/experiments/chart_asset_probe.py new file mode 100644 index 00000000..74175034 --- /dev/null +++ b/apps/worker/experiments/chart_asset_probe.py @@ -0,0 +1,452 @@ +"""Experimental: VLM-driven chart/table region detection + cropping. + +Goal of this experiment +----------------------- +Test whether a VLM can directly locate table/chart bounding boxes on a +*rendered page image* (the same way PAGE-TRACK renders pages), so we can +**crop** those regions out and later hand the crop to a dedicated table +model (e.g. tabular / table-transformer) instead of asking the VLM to +transcribe the whole table verbatim (error-prone + expensive output). + +This script does NOT touch production code. It only: + + 1. Renders selected PDF pages to PNG at a fixed DPI (mirrors PAGE-TRACK, + default 144 DPI) so the pixel<->point mapping is fully controlled. + 2. (VLM) Asks the model for table/chart regions as normalized [0,1000] + boxes, maps them to pixels, crops, and draws an annotated overlay. + 3. (Baseline) Runs PyMuPDF's free, pixel-accurate detectors + (``find_tables`` + vector ``drawings`` rects + embedded ``images``) + so VLM output can be judged against a non-LLM reference. + +Outputs (under --out): + pages/page-N.png full page render + crops/page-N_vlm-K_.png VLM crops + crops/page-N_base-K_.png baseline crops + annotated/page-N.png page with VLM(red) + baseline(green) boxes + results.json all regions + metadata + report.md human-readable summary + +Run: + cd apps/worker + uv run python experiments/chart_asset_probe.py \ + --pdf "/path/to/doc.pdf" \ + --pages all --baseline + # Uses $IMAGE_MODEL (default qwen3.6-flash) unless --model is set. + # Alternate: --model qwen3-vl-32b-instruct (open-weights, local-deployable). + +VLM model guidance (chart/table bbox grounding) +------------------------------------------------ +* **Default (cloud):** ``qwen3.6-flash`` via ``$IMAGE_MODEL`` — cheapest, + strong bbox quality on our probe PDFs. +* **Alternate (cloud or self-hosted):** ``qwen3-vl-32b-instruct`` — open + weights, can be deployed locally (vLLM / SGLang / Ollama); DashScope + China pricing (2026-04): input ¥2 / output ¥8 per 1M tokens (see + https://help.aliyun.com/zh/model-studio/model-pricing ). +* Coordinates are requested in a normalized 0-1000 space to be robust + to whatever internal resize the API performs. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import re +import sys +from dataclasses import dataclass, field, asdict +from pathlib import Path +from typing import Any + +import fitz # PyMuPDF +from PIL import Image, ImageDraw, ImageFont + + +# ── coordinate convention ───────────────────────────────────────────── +# We ask the VLM for boxes in a normalized integer space [0, 1000] for +# BOTH axes, with origin at the top-left of the page image. This is robust +# to whatever internal resize the API performs. +NORM = 1000 + +_PROMPT = ( + "You are a precise document layout detector. The attached image is a " + "single rendered page.\n\n" + "Task: locate every TABLE and every FIGURE (bar/line/pie charts, " + "plots, diagrams, schematic images). Do NOT transcribe their content. " + "Only return their bounding boxes.\n\n" + "Coordinate system: treat the page image as a {n}x{n} grid. The " + "top-left corner is (0,0) and the bottom-right is ({n},{n}). For each " + "region return [x1,y1,x2,y2] integers in that 0-{n} space, where " + "(x1,y1) is the top-left and (x2,y2) the bottom-right of a box that " + "TIGHTLY encloses the region INCLUDING its title/caption and axis " + "labels but EXCLUDING surrounding body paragraphs.\n\n" + "Return STRICT JSON only, no prose:\n" + '{{"regions":[{{"type":"table|chart|figure","bbox":[x1,y1,x2,y2],' + '"caption":"short caption text if visible else empty",' + '"confidence":0.0}}]}}\n' + "If there are no tables/charts, return {{\"regions\":[]}}." +).format(n=NORM) + + +@dataclass +class Region: + source: str # "vlm" | "baseline" + page: int + kind: str # table | chart | figure | drawing | image + bbox_px: list[int] # [x1,y1,x2,y2] in rendered-image pixels + caption: str = "" + confidence: float = 0.0 + crop_path: str = "" + + +@dataclass +class PageResult: + page: int + width_px: int + height_px: int + width_pt: float + height_pt: float + image_path: str + vlm_regions: list[Region] = field(default_factory=list) + baseline_regions: list[Region] = field(default_factory=list) + vlm_error: str = "" + + +# ── rendering ───────────────────────────────────────────────────────── + + +def render_page(page: fitz.Page, dpi: int, out_path: Path) -> tuple[int, int]: + zoom = dpi / 72.0 + mat = fitz.Matrix(zoom, zoom) + pix = page.get_pixmap(matrix=mat, alpha=False) + pix.save(str(out_path)) + return pix.width, pix.height + + +# ── VLM call ────────────────────────────────────────────────────────── + + +def call_vlm(image_path: Path, model: str) -> dict[str, Any]: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + with open(image_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + + content_parts = [ + {"type": "text", "text": _PROMPT}, + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + }, + ] + # For this standalone experiment we pass a direct Qwen key when available. + # The production Ali token pool depends on local Redis; direct mode keeps + # the probe runnable on a laptop without changing production behavior. + client = get_openai_client(model=model, api_key=_direct_api_key_for_model(model)) + raw, usage = client.chat_completion_with_usage( + messages=[{"role": "user", "content": content_parts}], + model=model, + temperature=0.0, + max_tokens=1200, + response_format={"type": "json_object"}, + usage_task="experiment.chart_asset_probe", + ) + data = json.loads(raw) + data["_usage"] = usage + return data + + +def _direct_api_key_for_model(model: str) -> str | None: + model_lower = model.lower() + if "qwen" not in model_lower: + return None + single = os.environ.get("ALI_API_KEY", "").strip() + if single: + return single + keys = os.environ.get("ALI_API_KEYS", "").strip() + if not keys: + return None + for item in re.split(r"[,;\s]+", keys): + if item.strip(): + return item.strip() + return None + + +def norm_to_px(box: list[float], w: int, h: int) -> list[int]: + x1, y1, x2, y2 = box + px = [ + int(round(x1 / NORM * w)), + int(round(y1 / NORM * h)), + int(round(x2 / NORM * w)), + int(round(y2 / NORM * h)), + ] + # normalize ordering + clamp + x1, x2 = sorted((px[0], px[2])) + y1, y2 = sorted((px[1], px[3])) + x1 = max(0, min(x1, w)) + x2 = max(0, min(x2, w)) + y1 = max(0, min(y1, h)) + y2 = max(0, min(y2, h)) + return [x1, y1, x2, y2] + + +# ── PyMuPDF baseline (free, pixel-accurate) ─────────────────────────── + + +def baseline_regions(page: fitz.Page, dpi: int, w: int, h: int) -> list[Region]: + zoom = dpi / 72.0 + regions: list[Region] = [] + + # 1) Native table finder + try: + tf = page.find_tables() + for t in tf.tables: + r = t.bbox # (x0,y0,x1,y1) in points + regions.append( + Region( + source="baseline", + page=page.number + 1, + kind="table", + bbox_px=_pt_box_to_px(r, zoom, w, h), + confidence=1.0, + ) + ) + except Exception as exc: # noqa: BLE001 + print(f" [baseline] find_tables failed: {exc}", file=sys.stderr) + + # 2) Embedded raster images (charts often exported as images) + try: + for info in page.get_image_info(): + bbox = info.get("bbox") + if bbox: + regions.append( + Region( + source="baseline", + page=page.number + 1, + kind="image", + bbox_px=_pt_box_to_px(bbox, zoom, w, h), + confidence=0.5, + ) + ) + except Exception as exc: # noqa: BLE001 + print(f" [baseline] get_image_info failed: {exc}", file=sys.stderr) + + return regions + + +def _pt_box_to_px(box: Any, zoom: float, w: int, h: int) -> list[int]: + x1, y1, x2, y2 = box + px = [int(round(x1 * zoom)), int(round(y1 * zoom)), + int(round(x2 * zoom)), int(round(y2 * zoom))] + px[0] = max(0, min(px[0], w)) + px[2] = max(0, min(px[2], w)) + px[1] = max(0, min(px[1], h)) + px[3] = max(0, min(px[3], h)) + return px + + +# ── cropping + overlay ──────────────────────────────────────────────── + + +def crop_region(img: Image.Image, region: Region, margin: int, out_path: Path) -> None: + x1, y1, x2, y2 = region.bbox_px + x1 = max(0, x1 - margin) + y1 = max(0, y1 - margin) + x2 = min(img.width, x2 + margin) + y2 = min(img.height, y2 + margin) + if x2 - x1 < 4 or y2 - y1 < 4: + return + img.crop((x1, y1, x2, y2)).save(out_path) + region.crop_path = str(out_path) + + +def draw_overlay(img: Image.Image, page_res: PageResult, out_path: Path) -> None: + canvas = img.convert("RGB").copy() + draw = ImageDraw.Draw(canvas) + try: + font = ImageFont.truetype("/System/Library/Fonts/Supplemental/Arial.ttf", 22) + except Exception: # noqa: BLE001 + font = ImageFont.load_default() + + for r in page_res.baseline_regions: + x1, y1, x2, y2 = r.bbox_px + draw.rectangle([x1, y1, x2, y2], outline=(0, 170, 0), width=3) + draw.text((x1 + 2, y1 + 2), f"base:{r.kind}", fill=(0, 120, 0), font=font) + + for i, r in enumerate(page_res.vlm_regions): + x1, y1, x2, y2 = r.bbox_px + draw.rectangle([x1, y1, x2, y2], outline=(220, 0, 0), width=3) + label = f"vlm:{r.kind} {r.confidence:.2f}" + draw.text((x1 + 2, max(0, y1 - 24)), label, fill=(200, 0, 0), font=font) + + canvas.save(out_path) + + +# ── page range parsing ──────────────────────────────────────────────── + + +def parse_pages(spec: str, total: int) -> list[int]: + if spec.strip().lower() == "all": + return list(range(1, total + 1)) + pages: set[int] = set() + for part in spec.split(","): + part = part.strip() + if not part: + continue + if "-" in part: + a, b = part.split("-", 1) + pages.update(range(int(a), int(b) + 1)) + else: + pages.add(int(part)) + return sorted(p for p in pages if 1 <= p <= total) + + +# ── main ────────────────────────────────────────────────────────────── + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--pdf", required=True) + ap.add_argument("--pages", default="all", help="e.g. 'all', '1-10', '1,3,5'") + ap.add_argument("--dpi", type=int, default=144, help="render DPI (PAGE-TRACK uses 144)") + ap.add_argument( + "--model", + default=os.environ.get("IMAGE_MODEL", ""), + help=( + "VLM model for bbox detection. Defaults to $IMAGE_MODEL " + "(qwen3.6-flash). Alternate: qwen3-vl-32b-instruct " + "(open-weights, local-deployable; DashScope ¥2/¥8 per 1M in/out)." + ), + ) + ap.add_argument("--margin", type=int, default=8, help="crop padding px") + ap.add_argument("--baseline", action="store_true", help="also run PyMuPDF baseline") + ap.add_argument("--no-vlm", action="store_true", help="skip VLM (baseline only)") + ap.add_argument("--out", default="") + args = ap.parse_args() + + pdf_path = Path(args.pdf) + if not pdf_path.exists(): + print(f"PDF not found: {pdf_path}", file=sys.stderr) + return 2 + if not args.no_vlm and not args.model: + print("No --model and IMAGE_MODEL unset; pass --model or use --no-vlm", + file=sys.stderr) + return 2 + + out_dir = Path(args.out) if args.out else ( + Path.home() / ".knowhere" / "_debug_parse" / pdf_path.stem / "asset_probe" + ) + (out_dir / "pages").mkdir(parents=True, exist_ok=True) + (out_dir / "crops").mkdir(parents=True, exist_ok=True) + (out_dir / "annotated").mkdir(parents=True, exist_ok=True) + + doc = fitz.open(str(pdf_path)) + page_nums = parse_pages(args.pages, doc.page_count) + print(f"PDF: {pdf_path.name} | {doc.page_count} pages | probing {len(page_nums)} " + f"| dpi={args.dpi} | model={args.model or '(none)'} | baseline={args.baseline}") + + results: list[PageResult] = [] + total_tokens = 0 + + for pno in page_nums: + page = doc[pno - 1] + img_path = out_dir / "pages" / f"page-{pno}.png" + w, h = render_page(page, args.dpi, img_path) + pr = PageResult( + page=pno, width_px=w, height_px=h, + width_pt=page.rect.width, height_pt=page.rect.height, + image_path=str(img_path), + ) + print(f"\n[page {pno}] {w}x{h}px") + + if not args.no_vlm: + try: + data = call_vlm(img_path, args.model) + usage = data.pop("_usage", {}) + total_tokens += int(usage.get("total_tokens", 0) or 0) + for k, reg in enumerate(data.get("regions", [])): + box = reg.get("bbox") or reg.get("bbox_norm") + if not box or len(box) != 4: + continue + region = Region( + source="vlm", page=pno, + kind=str(reg.get("type", "region")), + bbox_px=norm_to_px([float(v) for v in box], w, h), + caption=str(reg.get("caption", "")), + confidence=float(reg.get("confidence", 0.0) or 0.0), + ) + pr.vlm_regions.append(region) + print(f" [vlm] {len(pr.vlm_regions)} region(s); " + f"tokens={usage.get('total_tokens', '?')}") + except Exception as exc: # noqa: BLE001 + pr.vlm_error = str(exc) + print(f" [vlm] ERROR: {exc}", file=sys.stderr) + + if args.baseline: + pr.baseline_regions = baseline_regions(page, args.dpi, w, h) + print(f" [baseline] {len(pr.baseline_regions)} region(s)") + + # crops + overlay + with Image.open(img_path) as im: + for k, r in enumerate(pr.vlm_regions): + crop_region(im, r, args.margin, + out_dir / "crops" / f"page-{pno}_vlm-{k}_{r.kind}.png") + for k, r in enumerate(pr.baseline_regions): + crop_region(im, r, args.margin, + out_dir / "crops" / f"page-{pno}_base-{k}_{r.kind}.png") + if pr.vlm_regions or pr.baseline_regions: + draw_overlay(im, pr, out_dir / "annotated" / f"page-{pno}.png") + results.append(pr) + + doc.close() + + # results.json + (out_dir / "results.json").write_text( + json.dumps([asdict(r) for r in results], ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + # report.md + _write_report(out_dir, pdf_path, args, results, total_tokens) + print(f"\nDone. Output: {out_dir}") + print(f" - annotated overlays: {out_dir/'annotated'}") + print(f" - crops: {out_dir/'crops'}") + print(f" - results.json / report.md") + return 0 + + +def _write_report(out_dir: Path, pdf_path: Path, args: Any, + results: list[PageResult], total_tokens: int) -> None: + n_vlm = sum(len(r.vlm_regions) for r in results) + n_base = sum(len(r.baseline_regions) for r in results) + lines = [ + f"# Chart/Table Asset Probe — {pdf_path.name}", + "", + f"- pages probed: **{len(results)}**", + f"- dpi: **{args.dpi}**, model: **{args.model or '(no-vlm)'}**", + f"- VLM regions: **{n_vlm}**, baseline regions: **{n_base}**", + f"- total VLM tokens: **{total_tokens}**", + "", + "| page | px | vlm | base | vlm kinds | vlm error |", + "|-----:|----|----:|-----:|-----------|-----------|", + ] + for r in results: + kinds = ",".join(sorted({x.kind for x in r.vlm_regions})) or "-" + err = (r.vlm_error[:40] + "…") if r.vlm_error else "" + lines.append( + f"| {r.page} | {r.width_px}x{r.height_px} | {len(r.vlm_regions)} " + f"| {len(r.baseline_regions)} | {kinds} | {err} |" + ) + lines += [ + "", + "## How to read", + "- `annotated/page-N.png`: red = VLM boxes, green = PyMuPDF baseline.", + "- Judge VLM by: does the red box tightly enclose the table/chart " + "(incl. caption, excl. body text)? Compare against green baseline.", + "- `crops/`: the actual extracted assets to feed a table model next.", + ] + (out_dir / "report.md").write_text("\n".join(lines), encoding="utf-8") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/experiments/vlm_bbox_tabula_probe.py b/apps/worker/experiments/vlm_bbox_tabula_probe.py new file mode 100644 index 00000000..53938bb5 --- /dev/null +++ b/apps/worker/experiments/vlm_bbox_tabula_probe.py @@ -0,0 +1,294 @@ +"""Experimental: VLM bbox + Tabula PDF table extraction. + +This probe tests the next step after ``chart_asset_probe.py``: + +1. Read only VLM-detected ``kind=table`` regions from ``results.json``. +2. Convert rendered-image pixel boxes back to PDF point coordinates. +3. Pass each area to Tabula (`tabula-py`) against the original PDF text layer. +4. Export candidate DataFrames as HTML/CSV for manual inspection. + +Important: Tabula does not read PNG crops. It needs a text-based PDF, so the +VLM box is used only as a precise `area=[top,left,bottom,right]` constraint. + +Run: + cd apps/worker + uv run --with tabula-py python experiments/vlm_bbox_tabula_probe.py \ + --pdf "/path/to/doc.pdf" \ + --results "/path/to/asset_probe/results.json" +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import pandas as pd + + +@dataclass +class TabulaCandidate: + page: int + region_index: int + mode: str + caption: str + bbox_px: list[int] + area_pt: list[float] + ok: bool + rows: int = 0 + cols: int = 0 + html_path: str = "" + csv_path: str = "" + error: str = "" + + +def _px_box_to_tabula_area( + bbox_px: list[int], + *, + width_px: int, + height_px: int, + width_pt: float, + height_pt: float, + margin_pt: float, +) -> list[float]: + """Convert [x1,y1,x2,y2] px box into Tabula [top,left,bottom,right] points.""" + x1, y1, x2, y2 = bbox_px + left = x1 / width_px * width_pt + right = x2 / width_px * width_pt + top = y1 / height_px * height_pt + bottom = y2 / height_px * height_pt + return [ + round(max(0.0, top - margin_pt), 2), + round(max(0.0, left - margin_pt), 2), + round(min(height_pt, bottom + margin_pt), 2), + round(min(width_pt, right + margin_pt), 2), + ] + + +def _load_vlm_table_regions(results_path: Path) -> list[dict[str, Any]]: + pages = json.loads(results_path.read_text(encoding="utf-8")) + table_regions: list[dict[str, Any]] = [] + for page in pages: + for region_index, region in enumerate(page.get("vlm_regions", [])): + if str(region.get("kind", "")).lower() != "table": + continue + table_regions.append({ + "page": int(page["page"]), + "region_index": region_index, + "caption": str(region.get("caption", "")), + "bbox_px": list(region["bbox_px"]), + "width_px": int(page["width_px"]), + "height_px": int(page["height_px"]), + "width_pt": float(page["width_pt"]), + "height_pt": float(page["height_pt"]), + }) + return table_regions + + +def _is_meaningful_frame(df: pd.DataFrame) -> bool: + if df.empty or df.shape[1] == 0: + return False + non_empty = df.fillna("").astype(str).map(lambda s: bool(s.strip())) + return bool(non_empty.values.sum() >= 2) + + +def _safe_name(page: int, region_index: int, mode: str, candidate_index: int) -> str: + return f"page-{page}_vlm-{region_index}_{mode}-{candidate_index}" + + +def _has_working_java() -> bool: + if shutil.which("java") is None: + return False + try: + result = subprocess.run( + ["java", "-version"], + check=False, + capture_output=True, + text=True, + timeout=10, + ) + except Exception: + return False + combined = f"{result.stdout}\n{result.stderr}" + return result.returncode == 0 and "Unable to locate a Java Runtime" not in combined + + +def _extract_one( + *, + tabula: Any, + pdf_path: Path, + out_dir: Path, + region: dict[str, Any], + mode: str, + area_pt: list[float], +) -> list[TabulaCandidate]: + lattice = mode == "lattice" + stream = mode == "stream" + try: + frames = tabula.read_pdf( + str(pdf_path), + pages=region["page"], + area=area_pt, + guess=False, + lattice=lattice, + stream=stream, + multiple_tables=True, + pandas_options={"header": None}, + ) + except Exception as exc: # noqa: BLE001 + return [ + TabulaCandidate( + page=region["page"], + region_index=region["region_index"], + mode=mode, + caption=region["caption"], + bbox_px=region["bbox_px"], + area_pt=area_pt, + ok=False, + error=str(exc), + ) + ] + + candidates: list[TabulaCandidate] = [] + if not frames: + return [ + TabulaCandidate( + page=region["page"], + region_index=region["region_index"], + mode=mode, + caption=region["caption"], + bbox_px=region["bbox_px"], + area_pt=area_pt, + ok=False, + error="tabula returned no tables", + ) + ] + + for candidate_index, df in enumerate(frames): + name = _safe_name(region["page"], region["region_index"], mode, candidate_index) + candidate = TabulaCandidate( + page=region["page"], + region_index=region["region_index"], + mode=mode, + caption=region["caption"], + bbox_px=region["bbox_px"], + area_pt=area_pt, + ok=_is_meaningful_frame(df), + rows=int(df.shape[0]), + cols=int(df.shape[1]), + ) + if candidate.ok: + html_path = out_dir / "html" / f"{name}.html" + csv_path = out_dir / "csv" / f"{name}.csv" + df.to_html(html_path, index=False, header=False, na_rep="") + df.to_csv(csv_path, index=False, header=False) + candidate.html_path = str(html_path) + candidate.csv_path = str(csv_path) + else: + candidate.error = "empty or near-empty dataframe" + candidates.append(candidate) + return candidates + + +def _write_report(out_dir: Path, candidates: list[TabulaCandidate]) -> None: + ok_count = sum(1 for item in candidates if item.ok) + lines = [ + "# VLM BBox + Tabula Probe", + "", + f"- candidates: **{len(candidates)}**", + f"- successful non-empty tables: **{ok_count}**", + "", + "| page | vlm | mode | ok | shape | caption | error |", + "|-----:|----:|------|----|-------|---------|-------|", + ] + for item in candidates: + error = item.error.replace("\n", " ")[:80] + lines.append( + f"| {item.page} | {item.region_index} | {item.mode} | " + f"{'yes' if item.ok else 'no'} | {item.rows}x{item.cols} | " + f"{item.caption} | {error} |" + ) + lines.append("") + lines.append("Only VLM `kind=table` regions were used; baseline regions were ignored.") + (out_dir / "report.md").write_text("\n".join(lines), encoding="utf-8") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pdf", required=True) + parser.add_argument("--results", required=True) + parser.add_argument("--out", default="") + parser.add_argument("--margin-pt", type=float, default=2.0) + args = parser.parse_args() + + pdf_path = Path(args.pdf) + results_path = Path(args.results) + out_dir = ( + Path(args.out) + if args.out + else results_path.parent / "tabula_from_vlm_bbox" + ) + (out_dir / "html").mkdir(parents=True, exist_ok=True) + (out_dir / "csv").mkdir(parents=True, exist_ok=True) + + if not pdf_path.exists(): + raise FileNotFoundError(pdf_path) + if not results_path.exists(): + raise FileNotFoundError(results_path) + if not _has_working_java(): + raise RuntimeError( + "Java runtime is required by tabula-py but was not found. " + "Install a JRE (e.g. OpenJDK/Temurin) and rerun this probe." + ) + + try: + import tabula + except ImportError as exc: + raise RuntimeError( + "tabula-py is required. Run with: uv run --with tabula-py python ..." + ) from exc + + regions = _load_vlm_table_regions(results_path) + print(f"Loaded {len(regions)} VLM table regions from {results_path}") + + candidates: list[TabulaCandidate] = [] + for region in regions: + area_pt = _px_box_to_tabula_area( + region["bbox_px"], + width_px=region["width_px"], + height_px=region["height_px"], + width_pt=region["width_pt"], + height_pt=region["height_pt"], + margin_pt=args.margin_pt, + ) + for mode in ("lattice", "stream"): + extracted = _extract_one( + tabula=tabula, + pdf_path=pdf_path, + out_dir=out_dir, + region=region, + mode=mode, + area_pt=area_pt, + ) + candidates.extend(extracted) + for item in extracted: + print( + f"page={item.page} vlm={item.region_index} " + f"mode={item.mode} ok={item.ok} shape={item.rows}x{item.cols}" + ) + + (out_dir / "results.json").write_text( + json.dumps([asdict(item) for item in candidates], ensure_ascii=False, indent=2), + encoding="utf-8", + ) + _write_report(out_dir, candidates) + print(f"Done. Output: {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/pyproject.toml b/apps/worker/pyproject.toml index 53dcc813..2af6afec 100644 --- a/apps/worker/pyproject.toml +++ b/apps/worker/pyproject.toml @@ -27,6 +27,7 @@ dependencies = [ "gevent>=24.11.1", "psycogreen>=1.0.2", "httpcore>=1.0.6", + "tabula-py>=2.10.0", ] [dependency-groups] diff --git a/apps/worker/requirements.txt b/apps/worker/requirements.txt index 1cc6338b..b777a158 100644 --- a/apps/worker/requirements.txt +++ b/apps/worker/requirements.txt @@ -41,6 +41,8 @@ attrs==26.1.0 # via aiohttp authlib==1.6.11 # via knowhere-shared +backoff==2.2.1 + # via posthog bcrypt==4.3.0 # via pwdlib beautifulsoup4==4.13.4 @@ -111,7 +113,10 @@ cryptography==46.0.7 defusedxml==0.7.1 # via markitdown distro==1.9.0 - # via openai + # via + # openai + # posthog + # tabula-py dnspython==2.8.0 # via email-validator email-validator==2.2.0 @@ -226,6 +231,7 @@ numpy==2.2.6 # pymupdf-layout # rank-bm25 # scipy + # tabula-py onnxruntime==1.25.0 # via # magika @@ -300,7 +306,9 @@ packaging==26.1 # opentelemetry-instrumentation-sqlalchemy # pytest pandas==2.3.1 - # via knowhere-worker-app + # via + # knowhere-worker-app + # tabula-py pgvector==0.2.4 # via knowhere-shared pillow==12.2.0 @@ -311,6 +319,8 @@ pillow==12.2.0 # python-pptx pluggy==1.6.0 # via pytest +posthog==7.18.1 + # via knowhere-shared pptx2md==2.0.6 # via knowhere-worker-app prompt-toolkit==3.0.52 @@ -424,6 +434,7 @@ requests==2.33.0 # markitdown # opentelemetry-exporter-otlp-proto-http # oss2 + # posthog rich==15.0.0 # via logfire s3transfer==0.13.1 @@ -451,6 +462,8 @@ starlette==0.52.1 # via fastapi syntok==1.4.4 # via knowhere-shared +tabula-py==2.10.0 + # via knowhere-worker-app tabulate==0.10.0 # via pymupdf4llm tenacity==9.1.4 @@ -474,6 +487,7 @@ typing-extensions==4.14.1 # opentelemetry-exporter-otlp-proto-http # opentelemetry-sdk # opentelemetry-semantic-conventions + # posthog # pydantic # pydantic-core # pytest-asyncio 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 cf8dbb8a..cb8fa7db 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 @@ -16,8 +16,13 @@ build_node_rows, identify_leaf_nodes, ) +from app.services.document_parser.support.parser_rows import PARSER_ROW_COLUMNS +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.chunks.dataframe_chunk_converter import dataframe_to_chunks + +import pandas as pd def _same_page_sibling_skeletons() -> list[SectionSkeleton]: @@ -126,19 +131,26 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: assert leaf_b["extra_metadata"]["owned_pages"] == [232] -def test_build_node_rows_uses_vlm_node_summary_with_boundary(monkeypatch, tmp_path) -> None: +def test_build_node_rows_uses_vlm_node_summary_with_boundary( + monkeypatch, tmp_path +) -> None: captured: dict[str, object] = {} class _FakeClient: def chat_completion_with_usage(self, **kwargs): captured["messages"] = kwargs.get("messages") captured["usage_task"] = kwargs.get("usage_task") - return ('{"summary": "node summary", "keywords": "ka;kb"}', {"total_tokens": 10}) + return ( + '{"summary": "node summary", "keywords": "ka;kb"}', + {"total_tokens": 10}, + ) # node_assembler imports get_openai_client lazily from this module. import shared.services.ai.openai_compatible_client_sync as client_mod - monkeypatch.setattr(client_mod, "get_openai_client", lambda model=None: _FakeClient()) + monkeypatch.setattr( + client_mod, "get_openai_client", lambda model=None: _FakeClient() + ) img = tmp_path / "page-231.png" img.write_bytes(b"\x89PNG\r\n\x1a\n fake") @@ -165,3 +177,98 @@ def chat_completion_with_usage(self, **kwargs): assert leaf_a["summary"] == "node summary" assert leaf_a["keywords"] == "ka;kb" assert captured["usage_task"] == "page_memory.node_summary" + + +def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: + asset = PageAsset( + asset_id="asset_table_1", + page_index=231, + asset_index=1, + kind="table", + bbox_px=[10, 20, 200, 120], + width_px=1000, + height_px=1400, + width_pt=500, + height_pt=700, + title="资产表", + summary="表格资产摘要", + keywords=["资产", "表格"], + image_uri="images/image_page_231_table_1.png", + html_uri="tables/table_page_231_1.html", + extraction_status="table_html_extracted", + ) + + rows = build_node_rows( + skeletons=_same_page_sibling_skeletons(), + raw_text_by_page={231: "text-231", 232: "text-232"}, + image_uri_by_page={231: "pages/page-231.png", 232: "pages/page-232.png"}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={ + 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), + 232: PageTagResult(page_index=232, summary="s232", keywords=["k2"]), + }, + filename="demo.pdf", + verdict="page", + native_hierarchy=True, + budget=None, + vlm_model=None, + page_assets_by_page={231: [asset]}, + ) + + assert [row["type"] for row in rows] == ["table", "page", "page"] + assert rows[0]["path"] == "tables/table_page_231_1.html" + assert rows[0]["know_id"] == "asset_table_1" + + by_path = {row["path"]: row for row in rows} + owner = by_path["demo.pdf/3 基本规定/3.1 职责"] + shared = by_path["demo.pdf/3 基本规定/3.2 管理规定"] + assert '"relation": "embeds"' in owner["connectto"] + assert '"target": "[tables/table_page_231_1.html]"' in owner["connectto"] + assert '"relation": "related"' in shared["connectto"] + assert '"same_as_owner": "demo.pdf/3 基本规定/3.1 职责"' in shared["connectto"] + + +def test_page_connectto_normalizes_to_asset_chunk_id() -> None: + rows = [ + { + "content": "
A
", + "path": "tables/table_page_1_1.html", + "type": "table", + "length": 34, + "keywords": "asset", + "summary": "asset summary", + "know_id": "asset_table_1", + "tokens": "", + "connectto": "", + "addtime": "", + "page_nums": "1", + "extra_metadata": {}, + }, + { + "content": "page text", + "path": "demo.pdf/Section", + "type": "page", + "length": 9, + "keywords": "", + "summary": "", + "know_id": "page_1", + "tokens": "", + "connectto": '[{"target":"[tables/table_page_1_1.html]","relation":"embeds","ref":"[tables/table_page_1_1.html]"}]', + "addtime": "", + "page_nums": "1", + "extra_metadata": {}, + }, + ] + df = pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) + + chunks = dataframe_to_chunks(df) + + page_chunk = next(chunk for chunk in chunks if chunk["type"] == "page") + assert page_chunk["metadata"]["connect_to"] == [ + { + "target": "asset_table_1", + "relation": "embeds", + "ref": "[tables/table_page_1_1.html]", + } + ] diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index a5a48869..d823b8a9 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -30,12 +30,19 @@ class AIConfig(BaseModel): ) IMAGE_MODEL: str = Field( default="qwen3.6-flash", - description="Image model for image summary, atlas, and OCR flows", + description=( + "Default VLM for page tagging, OCR, atlas, and chart/table bbox " + "probes. Alternate: qwen3-vl-32b-instruct (open-weights, " + "self-hostable via vLLM/SGLang)." + ), ) IMAGE_MODEL_MAX: str = Field( default="qwen3.6-flash", - description="Higher-capability image model for OCR and image type classification", + description=( + "Higher-capability VLM for OCR and image classification. " + "Same alternates as IMAGE_MODEL (e.g. qwen3-vl-32b-instruct)." + ), ) RETRIEVAL_DECOMPOSITION_ENABLED: bool = Field( default=False, diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 960dadb5..0e1f93e1 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -665,6 +665,54 @@ def build_prompt(task, texts, query, **kwargs): - Do NOT add commentary, translation, or summary — transcription only. - Omit pure decorative running headers/footers and page numbers. - Return ONLY the JSON object, no markdown fences or extra text. +""" + + elif task == "page-memory-asset-detect": + temperature = 0 + top_p = 0.01 + max_tokens = kwargs.get("paras", {}).get("max_tokens", 1200) + grid_size = kwargs.get("paras", {}).get("grid_size", 1000) + prompt = f"""\ +You are a precise document layout detector. The attached image is a single PDF +page screenshot. + +Find visually distinct tables, charts, and figures that should become reusable +document assets. Return strict JSON: +{{ + "regions": [ + {{ + "kind": "table|chart|figure", + "bbox": [x1, y1, x2, y2], + "caption": "", + "title": "", + "summary": "<1-3 sentence searchable summary>", + "keywords": ["", ""], + "confidence": 0.0 + }} + ] +}} + +Coordinate system: +- Treat the page image as a {grid_size}x{grid_size} grid. +- Origin is the top-left corner. +- bbox values must be integers in [0, {grid_size}]. +- bbox must tightly include the whole asset: title, caption, legend, axes, + labels, table headers, and footnotes that are part of the asset. +- Exclude surrounding body paragraphs, page headers, page footers, and page + numbers. + +Rules: +- "table": rows/columns of data, forms, financial tables, appendix tables. +- "chart": plotted data such as bar/line/pie/scatter charts. +- "figure": diagrams, flowcharts, architecture drawings, embedded images, or + visual schematics that are not data charts. +- Do not transcribe entire tables. Summarize their topic and visible columns or + row groups. +- "keywords" must be an array of up to 8 strings in the same language as the + visible asset text. +- Use confidence 0.0-1.0. Only include assets you can localize. +- If there are no assets, return {{"regions":[]}}. +- Return ONLY the JSON object, no markdown fences or explanations. """ # ==================== Image Processing Prompts ==================== diff --git a/packages/shared-python/shared/services/chunks/chunk_connections.py b/packages/shared-python/shared/services/chunks/chunk_connections.py index efb0804f..dcfa865e 100644 --- a/packages/shared-python/shared/services/chunks/chunk_connections.py +++ b/packages/shared-python/shared/services/chunks/chunk_connections.py @@ -20,6 +20,7 @@ class ConnectionPayload(TypedDict, total=False): position: PositionPayload score: float keywords: list[str] + same_as_owner: str RelationshipRef: TypeAlias = str | ChunkRefSpan @@ -139,6 +140,9 @@ def normalize_connect_to_targets( ref = item.get("ref") if ref: normalized_item["ref"] = str(ref) + same_as_owner = item.get("same_as_owner") + if same_as_owner: + normalized_item["same_as_owner"] = str(same_as_owner) position = item.get("position") if isinstance(position, dict): start = position.get("start") diff --git a/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py b/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py index e837b85d..c2e5c9e7 100644 --- a/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py +++ b/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py @@ -342,7 +342,7 @@ def dataframe_to_chunks(df: _ParserDataFrame | None) -> list[Dict[str, JsonValue for chunk in chunks: metadata = chunk["metadata"] relationship_refs = metadata.pop("_relationship_refs", []) - if chunk["type"] != "text": + if chunk["type"] not in {"text", "page"}: continue embed_connections = convert_refs_to_embed_connections( relationship_refs, resource_target_map diff --git a/uv.lock b/uv.lock index 93f90694..884ec07d 100644 --- a/uv.lock +++ b/uv.lock @@ -1592,6 +1592,7 @@ dependencies = [ { name = "pypdf" }, { name = "python-docx" }, { name = "python-pptx" }, + { name = "tabula-py" }, { name = "tqdm" }, ] @@ -1628,6 +1629,7 @@ requires-dist = [ { name = "pypdf", specifier = "==6.10.2" }, { name = "python-docx", specifier = "==1.2.0" }, { name = "python-pptx", specifier = "==1.0.2" }, + { name = "tabula-py", specifier = ">=2.10.0" }, { name = "tqdm", specifier = "==4.67.1" }, ] @@ -4055,6 +4057,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/47/0d/e9700408c99275692c57ac786edfbede03cf2a9818bd3ecbc4cbf21a4448/syntok-1.4.4-py3-none-any.whl", hash = "sha256:263546bea559cb693bc23dc436307a44e17a242bb9f2a8c5e39eb86906aaf831", size = 24548, upload-time = "2022-03-12T18:06:11.254Z" }, ] +[[package]] +name = "tabula-py" +version = "2.10.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "distro" }, + { name = "numpy" }, + { name = "pandas" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e2/31/2a14a5048f681c404ae0b32a00d141128dd3065965190fdcae3b33e2bcae/tabula_py-2.10.0.tar.gz", hash = "sha256:75968a83fe978e5d56ccf23f0f0255a459c256b7b52db7cabe5ac795bb3b12df", size = 12459408, upload-time = "2024-10-17T02:51:19.668Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2f/80/10bc6f303054d1a06eb8628f90e5997f4b1272956a477230f3fa95637c28/tabula_py-2.10.0-py3-none-any.whl", hash = "sha256:c7596c559fc813e313eb4fbc7aabe7e4290dbd04717c4cbe4aa4a2cafd00ab63", size = 12021009, upload-time = "2024-10-17T02:51:16.427Z" }, +] + [[package]] name = "tabulate" version = "0.10.0" From 6d7042ccabe869f5a1878a0299770d18c996f197 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 23 Jun 2026 23:07:15 +0800 Subject: [PATCH 03/27] feat: enhance document profiling and page memory integration with skip shard plan option - Added `skip_shard_plan` parameter to `run_lightweight_anatomy` and `profile_document` to allow bypassing LLM shard decisions during profiling. - Updated `memory_service` to utilize `skip_shard_plan` for page memory processing. - Introduced asset annotation functionality for visualizing extracted assets on pages. - Refactored page section mapping and node assembly to support internal section body pages. - Improved environment configuration for Java dependencies in Docker setup. --- apps/worker/.env.example | 8 +- .../services/document_agent/coordinator.py | 45 +++-- .../document_parser/profiling/doc_profiler.py | 20 +- .../services/page_memory/memory_service.py | 186 +++++------------- .../services/page_memory/node_assembler.py | 72 +++++-- .../app/services/page_memory/page_assets.py | 111 ++++++++++- .../page_memory/page_section_mapper.py | 128 ------------ apps/worker/experiments/chart_asset_probe.py | 12 +- .../test_doc_profile_anatomy_contract.py | 6 +- .../test_document_agent_budget_contract.py | 34 ++-- .../test_page_memory_asset_java_contract.py | 172 ++++++++++++++++ ...est_page_memory_node_assembler_contract.py | 54 ++++- deploy/docker/Dockerfile.worker | 4 + .../services/storage/zip_chunk_schema.py | 10 +- 14 files changed, 512 insertions(+), 350 deletions(-) delete mode 100644 apps/worker/app/services/page_memory/page_section_mapper.py create mode 100644 apps/worker/tests/contract/test_page_memory_asset_java_contract.py diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 36658855..d969f12f 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -129,11 +129,15 @@ RETRIEVAL_PAGE_MEMORY_ENABLED=false # Page-memory asset extraction is opt-in. Requires VLM credentials; table HTML # extraction additionally requires tabula-py and a Java runtime. PAGE_MEMORY_ASSET_EXTRACTION_ENABLED=false -# PAGE_MEMORY_ASSET_MODEL= -# PAGE_MEMORY_ASSET_MAX_PAGES=50 +# Default: qwen3-vl-32b-instruct. Override to IMAGE_MODEL for cheaper probes. +# PAGE_MEMORY_ASSET_MODEL=qwen3-vl-32b-instruct +# Defaults to the full page count when asset extraction is enabled. +# PAGE_MEMORY_ASSET_MAX_PAGES= # PAGE_MEMORY_ASSET_BUDGET= # PAGE_MEMORY_ASSET_CONFIDENCE_THRESHOLD=0.3 # PAGE_MEMORY_TABLE_ENGINE=tabula +# macOS/Homebrew local debugging only; worker Docker image installs OpenJDK. +# JAVA_HOME=/opt/homebrew/opt/java MINERU_SHARD_CONCURRENCY=3 PARSE_AGENT_PLAN_BUDGET=50000 PARSE_AGENT_VISUAL_BUDGET=80000 diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 38f5093b..b20d2552 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -27,6 +27,7 @@ from app.services.document_agent.state import AgentBlackboard, DocumentAgentState from app.services.document_agent import tools as _registered_tools # noqa: F401 from app.services.document_agent.trace import ParseRunRecorder +from app.services.document_agent.validators import single_shard_plan class ProfileCoordinator: @@ -132,9 +133,11 @@ def run_toc(self) -> TocResult: self.blackboard.toc_hierarchies = None return self.blackboard.toc_result - def run_lightweight_anatomy(self) -> PageAnatomyMap: + def run_lightweight_anatomy( + self, *, skip_shard_plan: bool = False + ) -> PageAnatomyMap: try: - return self._run_lightweight_anatomy() + return self._run_lightweight_anatomy(skip_shard_plan=skip_shard_plan) except Exception as exc: self._record_failure(exc) raise @@ -192,7 +195,9 @@ def _run_toc(self) -> TocResult: ) return self.blackboard.toc_result - def _run_lightweight_anatomy(self) -> PageAnatomyMap: + def _run_lightweight_anatomy( + self, *, skip_shard_plan: bool = False + ) -> PageAnatomyMap: self.state = DocumentAgentState.RUNNING if not self.blackboard.page_features: self._run_bootstrap() @@ -205,18 +210,28 @@ def _run_lightweight_anatomy(self) -> PageAnatomyMap: else: self._ensure_disabled_toc_placeholder() self._run_h1_boundary_pipeline() - result = REGISTRY.dispatch("propose.shard_plan", self.ctx, {}) - self.trace.record_step( - round_index=self.round_index, - actor="anatomy:propose.shard_plan", - action_type="anatomy", - result=result, - tool_name="propose.shard_plan", - tool_args={}, - ) - if result.status not in {"ok", "invalid"}: - raise RuntimeError(result.error or "propose.shard_plan failed") - self.round_index += 1 + if skip_shard_plan: + # Page-based track processes pages individually via VLM and never + # consumes the shard plan; only build_anatomy_map's invariant needs + # it. Populate a single-shard placeholder to skip the LLM shard + # decision + H2 refinement (kept global for chunk-track oversized + # MinerU sharding). + self.blackboard.shard_plan = single_shard_plan( + self.blackboard.page_count + ) + else: + result = REGISTRY.dispatch("propose.shard_plan", self.ctx, {}) + self.trace.record_step( + round_index=self.round_index, + actor="anatomy:propose.shard_plan", + action_type="anatomy", + result=result, + tool_name="propose.shard_plan", + tool_args={}, + ) + if result.status not in {"ok", "invalid"}: + raise RuntimeError(result.error or "propose.shard_plan failed") + self.round_index += 1 anatomy = build_anatomy_map(self.ctx) self._persist_ready_anatomy(anatomy) return anatomy diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index ce9fa1d0..b6d5d77f 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -29,6 +29,7 @@ def profile_document( *, job_id: str | None = None, output_dir: str | None = None, + skip_shard_plan: bool = False, ) -> ParserDocumentProfile: """ General document profiling entry point. @@ -38,6 +39,10 @@ def profile_document( filename: File name (used to infer type) job_id: Parse job id for profile trace artifacts output_dir: Parser output directory + skip_shard_plan: When True, the lightweight anatomy stage skips the + LLM shard decision (+ H2 refinement) and populates a single-shard + placeholder instead. Used by the page-memory track, which never + consumes the shard plan. Chunk-track keeps the default (False). Returns: ParserDocumentProfile @@ -47,7 +52,13 @@ def profile_document( ext = os.path.splitext(filename)[1].lower() if ext == ".pdf": - return _profile_pdf(file_path, filename, job_id=job_id, output_dir=output_dir) + return _profile_pdf( + file_path, + filename, + job_id=job_id, + output_dir=output_dir, + skip_shard_plan=skip_shard_plan, + ) return ParserDocumentProfile( file_type=ext.lstrip("."), @@ -63,6 +74,7 @@ def _profile_pdf( *, job_id: str | None, output_dir: str | None, + skip_shard_plan: bool = False, ) -> ParserDocumentProfile: with _profile_db_context(enabled=bool(job_id)) as db: return _profile_pdf_with_db( @@ -71,6 +83,7 @@ def _profile_pdf( job_id=job_id, output_dir=output_dir, db=db, + skip_shard_plan=skip_shard_plan, ) @@ -81,6 +94,7 @@ def _profile_pdf_with_db( job_id: str | None, output_dir: str | None, db: Any | None, + skip_shard_plan: bool = False, ) -> ParserDocumentProfile: profile_job_id = job_id or filename agent_output_dir = os.path.join(output_dir, "_doc_agent") if output_dir else None @@ -134,7 +148,9 @@ def _profile_pdf_with_db( profile.toc = _map_toc_profile(coordinator) else: if not profile.is_atlas: - profile.anatomy = coordinator.run_lightweight_anatomy() + profile.anatomy = coordinator.run_lightweight_anatomy( + skip_shard_plan=skip_shard_plan + ) profile.toc = _map_toc_profile(coordinator) if trace := getattr(coordinator, "trace", None): diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 02c0c37f..d7dd8609 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -49,6 +49,7 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: pdf_filename, job_id=request.job_id, output_dir=full_output_dir, + skip_shard_plan=True, ) verdict = _decide_granularity(profile) @@ -111,33 +112,27 @@ def _build_page_dataframe( C1 page_renderer → PageRenderResult[] C2 page_plan → PagePlan[] C3 page_tagger → PageTagResult[] - C3b title detection → observed_titles (native hierarchy only) - C4b fine_hierarchy → refined SectionSkeleton[] (native hierarchy only) - C6 page_section_mapper → PageSectionMapping[] - C7 assemble DataFrame + C3b title detection → observed_titles + C4b fine_hierarchy → refined SectionSkeleton[] + C7 assemble node-granularity DataFrame """ from app.services.document_agent.budget import BudgetTracker, StageEnvelope from app.services.page_memory.page_plan import derive_page_processing_plan from app.services.page_memory.page_renderer import render_document_pages - from app.services.page_memory.page_section_mapper import map_pages_to_sections from app.services.page_memory.page_tagger import tag_pages - from app.services.page_memory.skeleton_extractor import extract_section_skeletons + from app.services.page_memory.skeleton_extractor import ( + SectionSkeleton, + extract_section_skeletons, + ) from app.services.page_memory.page_assets import ( extract_page_assets_from_renders, get_asset_budget, get_asset_confidence_threshold, get_asset_max_pages, + get_asset_model, page_asset_extraction_enabled, ) - # ── Native hierarchy flag (Step 3) ──────────────────────────────── - # Page-memory handles PDF/PPT(→PDF) structure with its own page-native - # hierarchy. Chunk-track hierarchy detection remains isolated to non - # page-memory formats such as DOCX/MD. - native_hierarchy = os.environ.get( - "PAGE_MEMORY_NATIVE_HIERARCHY", "true" - ).strip().lower() in ("1", "true", "yes", "on") - anatomy = getattr(profile, "anatomy", None) page_count = max(int(profile.page_count or 0), 0) if page_count <= 0: @@ -150,10 +145,8 @@ def _build_page_dataframe( page_locate_budget = int( os.environ.get("PAGE_MEMORY_LOCATE_BUDGET", str(min(page_count * 1600, 2_000_000))) ) - title_detection_budget = ( - int(os.environ.get("PAGE_MEMORY_TITLE_BUDGET", str(page_count * 800))) - if native_hierarchy - else 0 + title_detection_budget = int( + os.environ.get("PAGE_MEMORY_TITLE_BUDGET", str(page_count * 800)) ) asset_extraction_enabled = page_asset_extraction_enabled() asset_extraction_budget = ( @@ -184,11 +177,10 @@ def _build_page_dataframe( cap=None, ), } - if native_hierarchy: - stage_envelopes["page_title_detection"] = StageEnvelope( - min_guarantee=title_detection_budget, - cap=None, - ) + stage_envelopes["page_title_detection"] = StageEnvelope( + min_guarantee=title_detection_budget, + cap=None, + ) if asset_extraction_enabled: stage_envelopes["page_asset_extraction"] = StageEnvelope( min_guarantee=asset_extraction_budget, @@ -239,14 +231,11 @@ def _build_page_dataframe( page_assets_by_page: dict[int, list[Any]] = {} if asset_extraction_enabled: - asset_model = os.environ.get("PAGE_MEMORY_ASSET_MODEL") or os.environ.get( - "IMAGE_MODEL" - ) page_assets_by_page = extract_page_assets_from_renders( pdf_path=pdf_path, rendered_pages=rendered, output_dir=output_dir, - model_name=asset_model, + model_name=get_asset_model(), budget=budget, max_pages=get_asset_max_pages(page_count), confidence_threshold=get_asset_confidence_threshold(), @@ -269,8 +258,8 @@ def _build_page_dataframe( vlm_model=vlm_model, ) - # ── C3b + C4b: Native hierarchy (Step 2 + Step 3) ──────────────── - if native_hierarchy and skeletons: + # ── C3b + C4b: page-native hierarchy refinement ───────────────── + if skeletons: from app.services.page_memory.fine_hierarchy import ( compute_fat_leaf_pages, refine_fat_leaf_skeletons, @@ -316,12 +305,17 @@ def _build_page_dataframe( len(skeletons), ) - # ── C6: page → section mapping ─────────────────────────────────── - mappings = map_pages_to_sections( - page_count=page_count, - skeletons=skeletons, - filename=filename, - ) + if not skeletons: + skeletons = [ + SectionSkeleton( + section_path=f"{filename}/Root", + level=1, + start_page=1, + end_page=page_count, + title="Root", + parent_path=filename, + ) + ] # ── C7: assemble DataFrame rows ────────────────────────────────── tag_map = {t.page_index: t for t in tags} @@ -333,22 +327,6 @@ def _build_page_dataframe( for lbl in page_labels: label_map[lbl.page] = lbl.kind - # Build page → observed_titles. Native hierarchy uses only true VLM - # observations; the skeleton-title fallback is kept only for legacy mode. - observed_titles_map: dict[int, list[str]] = {} - if native_hierarchy: - for tag in tags: - if tag.observed_titles: - observed_titles_map[tag.page_index] = [ - t["text"] for t in tag.observed_titles if t.get("text") - ] - else: - # Legacy: map skeleton titles to start pages - for skel in skeletons: - titles = observed_titles_map.setdefault(skel.start_page, []) - if skel.title and skel.title not in titles: - titles.append(skel.title) - # Shared per-page lookups for node-granularity assembly. raw_text_by_page: dict[int, str] = {} image_uri_by_page: dict[int, str] = {} @@ -362,100 +340,30 @@ def _build_page_dataframe( Path(rend.image_path).relative_to(output_dir) ) - if native_hierarchy and skeletons: - from app.services.page_memory.node_assembler import build_node_rows + from app.services.page_memory.node_assembler import build_node_rows - rows = build_node_rows( - skeletons=skeletons, - raw_text_by_page=raw_text_by_page, - image_uri_by_page=image_uri_by_page, - image_path_by_page=image_path_by_page, - kind_by_page=label_map, - tag_by_page=tag_map, - filename=filename, - verdict=verdict, - native_hierarchy=native_hierarchy, - budget=budget, - vlm_model=vlm_model, - page_assets_by_page=page_assets_by_page, - ) - logger.info( - "[page_memory] C7 assembled {} node rows (verdict={})", - len(rows), verdict, - ) - else: - rows = _build_legacy_page_rows( - mappings=mappings, - tag_map=tag_map, - raw_text_by_page=raw_text_by_page, - image_uri_by_page=image_uri_by_page, - label_map=label_map, - observed_titles_map=observed_titles_map, - filename=filename, - verdict=verdict, - native_hierarchy=native_hierarchy, - ) - logger.info( - "[page_memory] C7 assembled {} page rows (verdict={})", - len(rows), verdict, - ) + rows = build_node_rows( + skeletons=skeletons, + raw_text_by_page=raw_text_by_page, + image_uri_by_page=image_uri_by_page, + image_path_by_page=image_path_by_page, + kind_by_page=label_map, + tag_by_page=tag_map, + filename=filename, + verdict=verdict, + budget=budget, + vlm_model=vlm_model, + page_assets_by_page=page_assets_by_page, + ) + logger.info( + "[page_memory] C7 assembled {} node rows (verdict={})", + len(rows), verdict, + ) df = pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) df.attrs["page_memory_skeletons"] = [skel.to_dict() for skel in skeletons] - df.attrs["page_memory_native_hierarchy"] = native_hierarchy return df -def _build_legacy_page_rows( - *, - mappings: Any, - tag_map: dict[int, Any], - raw_text_by_page: dict[int, str], - image_uri_by_page: dict[int, str], - label_map: dict[int, str], - observed_titles_map: dict[int, list[str]], - filename: str, - verdict: str, - native_hierarchy: bool, -) -> list[dict[str, Any]]: - """Legacy per-page rows (one ``type=page`` chunk per physical page).""" - rows: list[dict[str, Any]] = [] - for mapping in mappings: - page = mapping.page_index - tag = tag_map.get(page) - - content = (raw_text_by_page.get(page, "") or "").strip() - summary = tag.summary if tag else "" - keywords_str = ";".join(tag.keywords) if tag else "" - - know_id = f"page_{gen_str_codes(f'{filename}::{page}')}" - - rows.append({ - "content": content, - "path": mapping.section_path, - "type": "page", - "length": len(content), - "keywords": keywords_str, - "summary": summary, - "know_id": know_id, - "tokens": "", - "connectto": "", - "addtime": get_str_time(), - "page_nums": str(page), - "extra_metadata": { - "granularity": "page", - "page_index": page, - "page_image_uri": image_uri_by_page.get(page, ""), - "strategy_used": tag.strategy_used if tag else "", - "kind": label_map.get(page, "normal"), - "observed_titles": observed_titles_map.get(page, []), - "section_roles": mapping.section_roles, - "source_verdict": verdict, - "native_hierarchy": native_hierarchy, - }, - }) - return rows - - def _build_page_ctx( *, pdf_path: str, diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index 35f61646..ae8399ce 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -33,7 +33,6 @@ from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time from app.services.page_memory.page_assets import ( PageAsset, - asset_reference, build_asset_rows, ) from app.services.page_memory.page_tagger import PageTagResult @@ -57,6 +56,7 @@ class LeafNode: start_page: int end_page: int parent_path: str | None = None + body_pages: tuple[int, ...] | None = None @dataclass @@ -72,20 +72,31 @@ class NodePageView: def identify_leaf_nodes(skeletons: list[SectionSkeleton]) -> list[LeafNode]: - """Return leaf skeletons in reading order. + """Return body-bearing skeleton nodes in reading order. - A skeleton is a leaf when no other skeleton declares it as ``parent_path``. - Reading order = ``(start_page, original index)`` so that siblings sharing a - start page keep the top-to-bottom order produced by fine hierarchy. + A skeleton is included when either: + - no other skeleton declares it as ``parent_path``; or + - it is an internal section with pages not covered by descendant sections. + + The second case preserves parent-section body pages, such as the first page + of a section before the first child heading. Internal nodes remain + structural unless they carry such own body pages. """ parent_paths = {skel.parent_path for skel in skeletons if skel.parent_path} leaves: list[tuple[int, int, LeafNode]] = [] for index, skel in enumerate(skeletons): - if skel.section_path in parent_paths: - continue + is_internal = skel.section_path in parent_paths + body_pages: tuple[int, ...] | None = None + if is_internal: + body_page_list = _internal_body_pages(skel, skeletons) + if not body_page_list: + continue + body_pages = tuple(body_page_list) + + sort_page = body_pages[0] if body_pages else skel.start_page leaves.append( ( - skel.start_page, + sort_page, index, LeafNode( section_path=skel.section_path, @@ -94,6 +105,7 @@ def identify_leaf_nodes(skeletons: list[SectionSkeleton]) -> list[LeafNode]: start_page=skel.start_page, end_page=skel.end_page, parent_path=skel.parent_path, + body_pages=body_pages, ), ) ) @@ -101,6 +113,22 @@ def identify_leaf_nodes(skeletons: list[SectionSkeleton]) -> list[LeafNode]: return [leaf for _, _, leaf in leaves] +def _internal_body_pages( + skel: SectionSkeleton, + skeletons: list[SectionSkeleton], +) -> list[int]: + """Pages owned by an internal section itself, excluding descendants.""" + pages = set(range(skel.start_page, skel.end_page + 1)) + descendant_prefix = f"{skel.section_path}/" + for other in skeletons: + if other.section_path == skel.section_path: + continue + if not other.section_path.startswith(descendant_prefix): + continue + pages.difference_update(range(other.start_page, other.end_page + 1)) + return sorted(pages) + + def assign_pages_to_leaves( leaves: list[LeafNode], *, @@ -126,11 +154,12 @@ def assign_pages_to_leaves( page_owner: dict[int, LeafNode] = {} views: list[NodePageView] = [] for leaf in leaves: - pages = [ - page - for page in range(leaf.start_page, leaf.end_page + 1) - if page in available_pages - ] + source_pages = ( + list(leaf.body_pages) + if leaf.body_pages is not None + else list(range(leaf.start_page, leaf.end_page + 1)) + ) + pages = [page for page in source_pages if page in available_pages] owned: list[int] = [] for page in pages: if page not in page_owner: @@ -455,7 +484,6 @@ def build_node_rows( tag_by_page: dict[int, PageTagResult], filename: str, verdict: str, - native_hierarchy: bool, budget: BudgetTracker | None = None, vlm_model: str | None = None, page_assets_by_page: dict[int, list[PageAsset]] | None = None, @@ -522,7 +550,6 @@ def build_node_rows( "page_image_uris": page_image_uris, "kind": kind_by_page.get(leaf.start_page, "normal"), "source_verdict": verdict, - "native_hierarchy": native_hierarchy, }, } rows.append(row) @@ -559,16 +586,23 @@ def _attach_asset_connections( owner_leaf = page_owner.get(page_index) leaves_on_page = page_to_leaves.get(page_index, []) for asset in assets: - ref = asset_reference(asset) - if not ref: + # target = bare URI path (resolvable via target_map → chunk_id) + # ref = bracketed display reference (matches chunk-track convention) + uri = ( + asset.html_uri + if asset.kind == "table" and asset.html_uri + else asset.image_uri + ) + if not uri: continue + ref = f"[{uri}]" if owner_leaf is not None: owner_row = rows_by_path.get(owner_leaf.section_path) if owner_row is not None: _append_connect_to( owner_row, { - "target": ref, + "target": uri, "relation": "embeds", "ref": ref, }, @@ -583,7 +617,7 @@ def _attach_asset_connections( if row is None: continue connection: dict[str, Any] = { - "target": ref, + "target": uri, "relation": "related", "ref": ref, } diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index b1f5ffd6..e349f733 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -14,7 +14,7 @@ import pandas as pd from loguru import logger -from PIL import Image +from PIL import Image, ImageDraw, ImageFont from app.services.document_agent.budget import BudgetTracker from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time @@ -25,6 +25,8 @@ _BUDGET_STAGE = "page_asset_extraction" _GRID_SIZE = 1000 _VALID_KINDS = {"table", "chart", "figure"} +_DEFAULT_ASSET_MODEL = "qwen3-vl-32b-instruct" +_ASSET_ANNOTATE_DIR = "asset_annotate" @dataclass @@ -63,11 +65,18 @@ def get_asset_confidence_threshold() -> float: return 0.3 +def get_asset_model() -> str | None: + model = os.environ.get("PAGE_MEMORY_ASSET_MODEL", "").strip() + if model: + return model + return _DEFAULT_ASSET_MODEL + + def get_asset_max_pages(page_count: int) -> int: try: - value = int(os.environ.get("PAGE_MEMORY_ASSET_MAX_PAGES", "50")) + value = int(os.environ.get("PAGE_MEMORY_ASSET_MAX_PAGES", str(page_count))) except ValueError: - value = 50 + value = page_count return max(0, min(page_count, value)) @@ -366,6 +375,11 @@ def extract_page_assets_from_renders( page_assets.append(asset) if page_assets: assets_by_page[page.page_index] = page_assets + annotate_page_assets( + page=page, + assets=page_assets, + output_dir=output_dir, + ) logger.info( "[page_assets] extracted {} assets across {} pages", sum(len(items) for items in assets_by_page.values()), @@ -374,6 +388,52 @@ def extract_page_assets_from_renders( return assets_by_page +def annotate_page_assets( + *, + page: PageRenderResult, + assets: list[PageAsset], + output_dir: str, +) -> str: + if not assets or not page.image_path or not os.path.exists(page.image_path): + return "" + annotate_dir = Path(output_dir) / _ASSET_ANNOTATE_DIR + annotate_dir.mkdir(parents=True, exist_ok=True) + output_path = annotate_dir / f"page_{page.page_index}.png" + try: + with Image.open(page.image_path) as image: + canvas = image.convert("RGB").copy() + draw = ImageDraw.Draw(canvas) + try: + font = ImageFont.truetype( + "/System/Library/Fonts/Supplemental/Arial.ttf", 22 + ) + except Exception: + font = ImageFont.load_default() + for asset in assets: + x1, y1, x2, y2 = asset.bbox_px + color = _asset_box_color(asset.kind) + draw.rectangle([x1, y1, x2, y2], outline=color, width=4) + label = f"{asset.kind} {asset.confidence:.2f}" + draw.text((x1 + 2, max(0, y1 - 24)), label, fill=color, font=font) + canvas.save(output_path) + except Exception as exc: + logger.warning( + "[page_assets] failed to write asset annotation for page {}: {}", + page.page_index, + exc, + ) + return "" + return str(output_path) + + +def _asset_box_color(kind: str) -> tuple[int, int, int]: + if kind == "table": + return (220, 0, 0) + if kind == "chart": + return (0, 90, 220) + return (0, 150, 0) + + def build_asset_rows( page_assets_by_page: dict[int, list[PageAsset]], ) -> list[dict[str, Any]]: @@ -486,11 +546,50 @@ def _bbox_px_to_area_pt( def _has_working_java() -> bool: - if shutil.which("java") is None: + java_path = _resolve_working_java() + if java_path is None: return False + java_bin = str(java_path.parent) + path_parts = os.environ.get("PATH", "").split(os.pathsep) + if java_bin not in path_parts: + os.environ["PATH"] = os.pathsep.join([java_bin, *path_parts]) + os.environ.setdefault("JAVA_HOME", str(java_path.parent.parent)) + return True + + +def _resolve_working_java() -> Path | None: + candidates: list[Path] = [] + java_home = os.environ.get("JAVA_HOME") + if java_home: + candidates.append(Path(java_home) / "bin" / "java") + + which_java = shutil.which("java") + if which_java: + candidates.append(Path(which_java)) + + candidates.extend( + [ + Path("/opt/homebrew/opt/java/bin/java"), + Path("/usr/local/opt/java/bin/java"), + ] + ) + + seen: set[Path] = set() + for candidate in candidates: + if candidate in seen: + continue + seen.add(candidate) + if not candidate.exists(): + continue + if _java_version_ok(candidate): + return candidate + return None + + +def _java_version_ok(java_path: Path) -> bool: try: result = subprocess.run( - ["java", "-version"], + [str(java_path), "-version"], check=False, capture_output=True, text=True, @@ -527,11 +626,13 @@ def _safe_float(value: object, *, default: float) -> float: __all__ = [ "PageAsset", + "annotate_page_assets", "asset_reference", "build_asset_rows", "extract_page_assets_from_renders", "get_asset_budget", "get_asset_confidence_threshold", "get_asset_max_pages", + "get_asset_model", "page_asset_extraction_enabled", ] diff --git a/apps/worker/app/services/page_memory/page_section_mapper.py b/apps/worker/app/services/page_memory/page_section_mapper.py deleted file mode 100644 index 776393fe..00000000 --- a/apps/worker/app/services/page_memory/page_section_mapper.py +++ /dev/null @@ -1,128 +0,0 @@ -"""Page-to-section mapper: assigns each page a section_path and role. - -Given ``SectionSkeleton`` (from C4), maps every page to one or more -section paths with roles: - -- ``primary`` — page is the *start page* of this section -- ``spans`` — page falls within the section range but is not the start -- ``inherited`` — page is not in any section range; inherits the nearest - preceding primary section's path - -First page with no section → ``/Root``. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from enum import Enum - -from loguru import logger - -from app.services.page_memory.skeleton_extractor import SectionSkeleton - - -class SectionRole(str, Enum): - PRIMARY = "primary" - SPANS = "spans" - INHERITED = "inherited" - - -@dataclass(frozen=True) -class PageSectionMapping: - """Section assignment for a single page.""" - - page_index: int - section_path: str - """Primary section_path written to the DataFrame ``path`` column.""" - - section_roles: list[dict[str, str]] = field(default_factory=list) - """All roles: ``[{"section_path": "...", "role": "primary|spans|inherited"}, ...]``.""" - - -def map_pages_to_sections( - *, - page_count: int, - skeletons: list[SectionSkeleton], - filename: str = "", -) -> list[PageSectionMapping]: - """Map every page (1..page_count) to a section_path. - - Parameters - ---------- - page_count: - Total number of pages. - skeletons: - Leaf skeletons from ``skeleton_extractor``. Each has - ``section_path``, ``start_page``, ``end_page``. - filename: - Source filename for the fallback root path. - - Returns - ------- - list[PageSectionMapping] - One entry per page, ordered by page_index. - """ - root_path = f"{filename}/Root" if filename else "Root" - - # Build a sorted list of section ranges - sorted_skeletons = sorted(skeletons, key=lambda s: (s.start_page, s.level)) - - # Pre-compute: for each page, which skeletons cover it - page_roles: dict[int, list[dict[str, str]]] = { - page: [] for page in range(1, page_count + 1) - } - - for skel in sorted_skeletons: - for page in range(skel.start_page, skel.end_page + 1): - if page < 1 or page > page_count: - continue - role = ( - SectionRole.PRIMARY - if page == skel.start_page - else SectionRole.SPANS - ) - page_roles[page].append( - {"section_path": skel.section_path, "role": role.value} - ) - - # Assign primary section_path per page - results: list[PageSectionMapping] = [] - last_primary_path = root_path - - for page in range(1, page_count + 1): - roles = page_roles.get(page, []) - - if roles: - # Pick the deepest-level primary as the main path, - # or fall back to the first spans entry - primaries = [r for r in roles if r["role"] == SectionRole.PRIMARY.value] - if primaries: - main_path = primaries[-1]["section_path"] # deepest - last_primary_path = main_path - else: - main_path = roles[0]["section_path"] - else: - # No skeleton covers this page → inherited - main_path = last_primary_path - roles = [ - {"section_path": last_primary_path, "role": SectionRole.INHERITED.value} - ] - - results.append( - PageSectionMapping( - page_index=page, - section_path=main_path, - section_roles=roles, - ) - ) - - assigned = sum(1 for r in results if any( - sr["role"] != SectionRole.INHERITED.value for sr in r.section_roles - )) - logger.info( - "[page_section_mapper] mapped {} pages: {} assigned, {} inherited", - len(results), - assigned, - len(results) - assigned, - ) - return results diff --git a/apps/worker/experiments/chart_asset_probe.py b/apps/worker/experiments/chart_asset_probe.py index 74175034..125c288e 100644 --- a/apps/worker/experiments/chart_asset_probe.py +++ b/apps/worker/experiments/chart_asset_probe.py @@ -22,7 +22,7 @@ pages/page-N.png full page render crops/page-N_vlm-K_.png VLM crops crops/page-N_base-K_.png baseline crops - annotated/page-N.png page with VLM(red) + baseline(green) boxes + asset_annotate/page_N.png page with VLM(red) + baseline(green) boxes results.json all regions + metadata report.md human-readable summary @@ -338,7 +338,7 @@ def main() -> int: ) (out_dir / "pages").mkdir(parents=True, exist_ok=True) (out_dir / "crops").mkdir(parents=True, exist_ok=True) - (out_dir / "annotated").mkdir(parents=True, exist_ok=True) + (out_dir / "asset_annotate").mkdir(parents=True, exist_ok=True) doc = fitz.open(str(pdf_path)) page_nums = parse_pages(args.pages, doc.page_count) @@ -395,7 +395,7 @@ def main() -> int: crop_region(im, r, args.margin, out_dir / "crops" / f"page-{pno}_base-{k}_{r.kind}.png") if pr.vlm_regions or pr.baseline_regions: - draw_overlay(im, pr, out_dir / "annotated" / f"page-{pno}.png") + draw_overlay(im, pr, out_dir / "asset_annotate" / f"page_{pno}.png") results.append(pr) doc.close() @@ -409,9 +409,9 @@ def main() -> int: # report.md _write_report(out_dir, pdf_path, args, results, total_tokens) print(f"\nDone. Output: {out_dir}") - print(f" - annotated overlays: {out_dir/'annotated'}") + print(f" - annotated overlays: {out_dir/'asset_annotate'}") print(f" - crops: {out_dir/'crops'}") - print(f" - results.json / report.md") + print(" - results.json / report.md") return 0 @@ -440,7 +440,7 @@ def _write_report(out_dir: Path, pdf_path: Path, args: Any, lines += [ "", "## How to read", - "- `annotated/page-N.png`: red = VLM boxes, green = PyMuPDF baseline.", + "- `asset_annotate/page_N.png`: red = VLM boxes, green = PyMuPDF baseline.", "- Judge VLM by: does the red box tightly enclose the table/chart " "(incl. caption, excl. body text)? Compare against green baseline.", "- `crops/`: the actual extracted assets to feed a table model next.", diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 19c5242d..85fadef9 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -592,7 +592,7 @@ def run_coarse(self) -> DocumentProfile: routing_category=PdfRoutingCategory.GENERIC.value, ) - def run_lightweight_anatomy(self): + def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): self.calls.append("run_lightweight_anatomy") return fake_anatomy @@ -644,7 +644,7 @@ def run_coarse(self) -> DocumentProfile: def run_toc(self) -> TocResult: raise AssertionError("kill switch should not call TOC profiling") - def run_lightweight_anatomy(self): + def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): self.calls.append("run_lightweight_anatomy") self.blackboard.toc_result = TocResult( method="none", @@ -716,7 +716,7 @@ def run_toc(self) -> TocResult: self.calls.append("run_toc") raise AssertionError("run_toc should be no-op after TOC-before-coarse") - def run_lightweight_anatomy(self): + def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): self.calls.append("run_lightweight_anatomy") return fake_anatomy diff --git a/apps/worker/tests/contract/test_document_agent_budget_contract.py b/apps/worker/tests/contract/test_document_agent_budget_contract.py index a3076b56..2dfc47e4 100644 --- a/apps/worker/tests/contract/test_document_agent_budget_contract.py +++ b/apps/worker/tests/contract/test_document_agent_budget_contract.py @@ -90,7 +90,7 @@ def test_dataframe_converter_accepts_page_chunks_with_extra_metadata() -> None: assert chunks[0]["metadata"]["page_nums"] == [1, 2] -def test_zip_chunk_schema_preserves_page_memory_section_roles() -> None: +def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: chunks = [ { "chunk_id": "page-231", @@ -100,21 +100,14 @@ def test_zip_chunk_schema_preserves_page_memory_section_roles() -> None: "metadata": { "summary": "page summary", "page_nums": [231], - "page_index": 231, - "section_roles": [ - { - "section_path": "demo.pdf/3 基本规定", - "role": "primary", - }, - { - "section_path": "demo.pdf/3 基本规定/3.1 职责", - "role": "primary", - }, - { - "section_path": "demo.pdf/3 基本规定/3.2 管理规定", - "role": "primary", - }, - ], + "granularity": "node", + "page_indices": [231, 232], + "owned_pages": [232], + "section_path": "demo.pdf/3 基本规定/3.2 管理规定", + "section_level": 2, + "page_image_uris": ["pages/page-231.png", "pages/page-232.png"], + "kind": "normal", + "source_verdict": "page", }, } ] @@ -127,8 +120,13 @@ def test_zip_chunk_schema_preserves_page_memory_section_roles() -> None: metadata = formatted[0]["metadata"] assert metadata["page_nums"] == [231] - assert metadata["page_index"] == 231 - assert metadata["section_roles"] == chunks[0]["metadata"]["section_roles"] + assert metadata["granularity"] == "node" + assert metadata["page_indices"] == [231, 232] + assert metadata["owned_pages"] == [232] + assert metadata["section_path"] == "demo.pdf/3 基本规定/3.2 管理规定" + assert metadata["section_level"] == 2 + assert metadata["page_image_uris"] == ["pages/page-231.png", "pages/page-232.png"] + assert metadata["source_verdict"] == "page" def test_page_memory_granularity_routes_supported_page_modes() -> None: diff --git a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py new file mode 100644 index 00000000..b5b4a8ce --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import json +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") + +from PIL import Image + +from app.services.page_memory import page_assets +from app.services.page_memory.page_renderer import PageRenderResult + + +def test_page_assets_default_to_qwen32b_and_full_page_scan(monkeypatch) -> None: + monkeypatch.delenv("PAGE_MEMORY_ASSET_MODEL", raising=False) + monkeypatch.delenv("PAGE_MEMORY_ASSET_MAX_PAGES", raising=False) + + assert page_assets.get_asset_model() == "qwen3-vl-32b-instruct" + assert page_assets.get_asset_max_pages(301) == 301 + + +def test_page_asset_detection_uses_lowest_temperature(monkeypatch, tmp_path) -> None: + captured: dict[str, object] = {} + + class _FakeClient: + def chat_completion_with_usage(self, **kwargs): + captured.update(kwargs) + return ( + json.dumps( + { + "regions": [ + { + "kind": "table", + "bbox": [100, 100, 300, 300], + "caption": "caption", + "title": "title", + "summary": "summary", + "keywords": ["table"], + "confidence": 1.0, + } + ] + } + ), + {"total_tokens": 10}, + ) + + import shared.services.ai.openai_compatible_client_sync as client_mod + + monkeypatch.setattr( + client_mod, "get_openai_client", lambda model=None: _FakeClient() + ) + + image_path = tmp_path / "page.png" + Image.new("RGB", (100, 120), "white").save(image_path) + page = PageRenderResult( + page_index=1, + image_path=str(image_path), + raw_text="", + width=50, + height=60, + is_landscape=False, + ) + + assets = page_assets.detect_page_assets( + page=page, + source_name="demo.pdf", + model_name=page_assets.get_asset_model(), + budget=None, + confidence_threshold=0.3, + ) + + assert captured["model"] == "qwen3-vl-32b-instruct" + assert captured["temperature"] == 0 + assert assets[0].bbox_px == [10, 12, 30, 36] + + +def test_page_asset_extraction_keeps_debug_output_minimal( + monkeypatch, tmp_path +) -> None: + page_images = [] + for page_index in [1, 2]: + image_path = tmp_path / f"page-{page_index}.png" + Image.new("RGB", (100, 120), "white").save(image_path) + page_images.append( + PageRenderResult( + page_index=page_index, + image_path=str(image_path), + raw_text="", + width=50, + height=60, + is_landscape=False, + ) + ) + + def _fake_detect(**kwargs): + page = kwargs["page"] + if page.page_index != 1: + return [] + return [ + page_assets.PageAsset( + asset_id="asset_1", + page_index=1, + asset_index=1, + kind="table", + bbox_px=[10, 12, 30, 36], + width_px=100, + height_px=120, + width_pt=50, + height_pt=60, + title="table", + summary="summary", + confidence=1.0, + ) + ] + + def _fake_crop(*, asset, page_image_path, output_dir, margin_px=4): + asset.image_uri = "images/image_page_1_table_1.png" + asset.image_path = str(tmp_path / asset.image_uri) + return asset + + def _fake_extract_table(*, asset, pdf_path, output_dir): + table_path = tmp_path / "tables" / "table_page_1_1.html" + table_path.parent.mkdir(parents=True, exist_ok=True) + table_path.write_text("
A
", encoding="utf-8") + asset.html_uri = "tables/table_page_1_1.html" + asset.html_path = str(table_path) + asset.extraction_status = "table_html_extracted" + return asset + + monkeypatch.setattr(page_assets, "detect_page_assets", _fake_detect) + monkeypatch.setattr(page_assets, "crop_page_asset", _fake_crop) + monkeypatch.setattr(page_assets, "extract_table_html", _fake_extract_table) + + assets_by_page = page_assets.extract_page_assets_from_renders( + pdf_path="demo.pdf", + rendered_pages=page_images, + output_dir=str(tmp_path), + model_name=page_assets.get_asset_model(), + budget=None, + max_pages=2, + confidence_threshold=0.3, + ) + + assert sorted(assets_by_page) == [1] + assert (tmp_path / "tables" / "table_page_1_1.html").exists() + assert (tmp_path / "asset_annotate" / "page_1.png").exists() + assert not (tmp_path / "page_asset_bboxes.json").exists() + + +def test_page_assets_adds_java_home_to_path(monkeypatch, tmp_path) -> None: + java_home = tmp_path / "jdk" + java_bin = java_home / "bin" + java_bin.mkdir(parents=True) + java = java_bin / "java" + java.write_text( + "#!/bin/sh\n" + "echo 'openjdk version \"25.0.2\"' >&2\n" + "exit 0\n", + encoding="utf-8", + ) + java.chmod(0o755) + + monkeypatch.setenv("JAVA_HOME", str(java_home)) + monkeypatch.setenv("PATH", os.defpath) + + assert page_assets._has_working_java() is True # noqa: SLF001 + assert os.environ["PATH"].split(os.pathsep)[0] == str(java_bin) 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 cb8fa7db..6b77632d 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 @@ -103,7 +103,6 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: }, filename="demo.pdf", verdict="page", - native_hierarchy=True, budget=None, vlm_model=None, ) @@ -131,6 +130,53 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: assert leaf_b["extra_metadata"]["owned_pages"] == [232] +def test_build_node_rows_keeps_internal_section_body_pages() -> None: + parent = SectionSkeleton( + section_path="demo.pdf/4 风险辨识与分级管控", + level=1, + start_page=233, + end_page=234, + title="4 风险辨识与分级管控", + parent_path="demo.pdf", + ) + child = SectionSkeleton( + section_path="demo.pdf/4 风险辨识与分级管控/4.1 风险评价方法", + level=2, + start_page=234, + end_page=234, + title="4.1 风险评价方法", + parent_path="demo.pdf/4 风险辨识与分级管控", + ) + + rows = build_node_rows( + skeletons=[parent, child], + raw_text_by_page={233: "parent body", 234: "child body"}, + image_uri_by_page={233: "pages/page-233.png", 234: "pages/page-234.png"}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={ + 233: PageTagResult(page_index=233, summary="s233", keywords=["parent"]), + 234: PageTagResult(page_index=234, summary="s234", keywords=["child"]), + }, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model=None, + ) + + by_path = {row["path"]: row for row in rows} + assert by_path["demo.pdf/4 风险辨识与分级管控"]["page_nums"] == "233" + assert by_path["demo.pdf/4 风险辨识与分级管控"]["content"] == "parent body" + assert ( + by_path["demo.pdf/4 风险辨识与分级管控/4.1 风险评价方法"]["page_nums"] + == "234" + ) + assert ( + by_path["demo.pdf/4 风险辨识与分级管控/4.1 风险评价方法"]["content"] + == "child body" + ) + + def test_build_node_rows_uses_vlm_node_summary_with_boundary( monkeypatch, tmp_path ) -> None: @@ -167,7 +213,6 @@ def chat_completion_with_usage(self, **kwargs): }, filename="demo.pdf", verdict="page", - native_hierarchy=True, budget=None, vlm_model="fake-vlm", ) @@ -210,7 +255,6 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: }, filename="demo.pdf", verdict="page", - native_hierarchy=True, budget=None, vlm_model=None, page_assets_by_page={231: [asset]}, @@ -224,7 +268,7 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: owner = by_path["demo.pdf/3 基本规定/3.1 职责"] shared = by_path["demo.pdf/3 基本规定/3.2 管理规定"] assert '"relation": "embeds"' in owner["connectto"] - assert '"target": "[tables/table_page_231_1.html]"' in owner["connectto"] + assert '"target": "tables/table_page_231_1.html"' in owner["connectto"] assert '"relation": "related"' in shared["connectto"] assert '"same_as_owner": "demo.pdf/3 基本规定/3.1 职责"' in shared["connectto"] @@ -254,7 +298,7 @@ def test_page_connectto_normalizes_to_asset_chunk_id() -> None: "summary": "", "know_id": "page_1", "tokens": "", - "connectto": '[{"target":"[tables/table_page_1_1.html]","relation":"embeds","ref":"[tables/table_page_1_1.html]"}]', + "connectto": '[{"target":"tables/table_page_1_1.html","relation":"embeds","ref":"[tables/table_page_1_1.html]"}]', "addtime": "", "page_nums": "1", "extra_metadata": {}, diff --git a/deploy/docker/Dockerfile.worker b/deploy/docker/Dockerfile.worker index ad12e0df..fd25e573 100644 --- a/deploy/docker/Dockerfile.worker +++ b/deploy/docker/Dockerfile.worker @@ -49,8 +49,12 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libreoffice-calc \ libreoffice-impress \ libreoffice-writer \ + openjdk-17-jre-headless \ && rm -rf /var/lib/apt/lists/* +# tabula-py shells out to Java for page-memory table HTML extraction. +RUN java -version + # Copy the virtualenv from the builder stage COPY --from=builder /app/venv /app/venv diff --git a/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py index 5f7b6a9a..962c7f65 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -106,8 +106,8 @@ def format_chunks( ) metadata["tokens"] = [] elif chunk_type == "page": - # Page chunks carry page_image_uri and section_roles in - # extra_metadata; preserve them for page-memory navigation. + # Page-memory page chunks preserve node/whole-doc navigation + # metadata from extra_metadata. _merge_page_chunk_metadata(metadata, existing_metadata) metadata["keywords"] = existing_metadata.get("keywords") or [] metadata["tokens"] = [] @@ -132,20 +132,14 @@ def _normalize_chunk_type(value: Any) -> str: _PAGE_CHUNK_METADATA_KEYS = ( "granularity", - "page_index", "page_indices", "owned_pages", "section_path", "section_level", - "page_image_uri", "page_image_uris", - "strategy_used", "kind", "status", - "observed_titles", - "section_roles", "source_verdict", - "native_hierarchy", ) From a0ab8ca5db63a14ed8acd390f13b078c4cc50491 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 24 Jun 2026 14:04:37 +0800 Subject: [PATCH 04/27] feat: enhance retrieval and document ingestion with new chunk types and effective parse track handling - Updated `data_type` in `RetrievalQueryRequest` to support new chunk types: page (7) and text+image+table (8). - Refactored document ingestion service to apply effective parse track based on file extension and user-defined settings. - Introduced new validation for parse track handling in oversized PDF processing. - Enhanced tests to cover new chunk types and parse track logic, ensuring robust validation and functionality. --- apps/api/app/api/v1/routes/retrieval.py | 7 +- .../services/document_ingestion/service.py | 61 ++- .../test_page_memory_parse_track_contract.py | 77 ++++ .../tests/contract/test_retrieval_contract.py | 8 +- .../connect_builder/summary_builder.py | 10 +- .../document_ingestion/artifact_refs.py | 59 +++ .../success_finalization.py | 11 +- .../document_parser/assets/inline_asset.py | 6 +- .../orchestration/oversized_pdf_policy.py | 22 ++ .../document_parser/profiling/doc_profiler.py | 19 +- .../tables/table_asset_writer.py | 5 +- .../services/page_memory/memory_service.py | 10 +- .../services/page_memory/node_assembler.py | 7 - .../app/services/page_memory/page_assets.py | 32 +- apps/worker/experiments/chart_asset_probe.py | 226 ++++++----- ...test_agentic_evidence_renderer_contract.py | 58 ++- .../test_doc_profile_anatomy_contract.py | 47 +++ .../test_document_agent_budget_contract.py | 21 +- .../contract/test_excel_parser_contract.py | 3 +- .../test_page_memory_asset_java_contract.py | 19 +- .../test_page_memory_navigation_contract.py | 102 +---- ...est_page_memory_node_assembler_contract.py | 9 +- .../test_page_memory_retrieval_contract.py | 306 +++++++++++++++ .../contract/test_parse_task_contract.py | 4 +- .../shared/models/schemas/job.py | 6 +- .../shared/services/ai/prompt_service.py | 369 ++++++++++-------- .../services/retrieval/agentic/core/types.py | 2 +- .../retrieval/agentic/evidence/builder.py | 9 +- .../retrieval/agentic/evidence/renderer.py | 106 ++++- .../retrieval/agentic/navigation/assets.py | 2 +- .../retrieval/agentic/orchestrator.py | 2 +- .../retrieval/execution/reference_resolver.py | 43 +- .../services/retrieval/hydration/assets.py | 76 +++- .../retrieval/hydration/result_assembly.py | 62 ++- .../services/retrieval/hydration/row_utils.py | 4 +- .../services/retrieval/search/lexical_text.py | 44 ++- .../shared/services/retrieval/settings.py | 2 + .../shared/services/storage/result_storage.py | 29 +- .../services/storage/zip_chunk_schema.py | 12 +- .../services/storage/zip_doc_navigation.py | 194 +-------- .../services/storage/zip_result_resources.py | 24 -- .../services/storage/zip_result_schema.py | 17 +- .../services/storage/zip_result_service.py | 25 +- 43 files changed, 1389 insertions(+), 768 deletions(-) create mode 100644 apps/worker/app/services/document_ingestion/artifact_refs.py create mode 100644 apps/worker/tests/contract/test_page_memory_retrieval_contract.py diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index eb6a5846..446e4695 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -36,8 +36,11 @@ class RetrievalQueryRequest(BaseModel): data_type: int = Field( 1, ge=1, - le=6, - description="Chunk type filter: 1=all, 2=text, 3=image, 4=table, 5=text+image, 6=text+table", + le=8, + description=( + "Chunk type filter: 1=all, 2=text, 3=image, 4=table, " + "5=text+image, 6=text+table, 7=page, 8=text+image+table" + ), ) signal_paths: list[str] = Field( default_factory=list, description="Path keywords for include/exclude filtering" diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index 0bc17bd3..64e1f92d 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -2,7 +2,7 @@ import os import uuid -from typing import cast +from typing import Literal, cast from app.services.document_ingestion.confirmation_service import ( DocumentIngestionConfirmationService, @@ -38,6 +38,8 @@ from shared.services.http.url_security import validate_http_url_and_resolve_ip_async JobMetadata = dict[str, object] +ParseTrack = Literal["chunk", "page_memory"] +_PAGE_MEMORY_PARSE_TRACK_EXTENSIONS = {".pdf", ".pptx"} class DocumentIngestionService: @@ -196,13 +198,13 @@ async def _validate_create_payload(self, payload: JobCreate) -> None: } ], ) - _validate_parse_track_for_extension( - parse_track=payload.parse_track, + _apply_effective_parse_track( + payload, file_extension=file_extension, ) elif payload.file_name: - _validate_parse_track_for_extension( - parse_track=payload.parse_track, + _apply_effective_parse_track( + payload, file_extension=os.path.splitext(payload.file_name)[1].lower(), ) @@ -292,7 +294,7 @@ def _validate_parse_track_for_extension(*, parse_track: str, file_extension: str } ], ) - if file_extension.lower() not in {".pdf", ".pptx"}: + if file_extension.lower() not in _PAGE_MEMORY_PARSE_TRACK_EXTENSIONS: raise ValidationException( user_message="page_memory parse track only supports PDF and PPTX", violations=[ @@ -302,3 +304,50 @@ def _validate_parse_track_for_extension(*, parse_track: str, file_extension: str } ], ) + + +def _apply_effective_parse_track( + payload: JobCreate, + *, + file_extension: str, +) -> ParseTrack: + explicit = "parse_track" in payload.model_fields_set + effective_parse_track = _resolve_effective_parse_track( + parse_track=payload.parse_track, + file_extension=file_extension, + explicit=explicit, + ) + payload.parse_track = effective_parse_track + return effective_parse_track + + +def _resolve_effective_parse_track( + *, + parse_track: ParseTrack, + file_extension: str, + explicit: bool, +) -> ParseTrack: + if parse_track == "chunk": + return "chunk" + if parse_track != "page_memory": + raise ValidationException( + user_message="Unsupported parse_track", + violations=[ + {"field": "parse_track", "description": "Must be chunk or page_memory"} + ], + ) + + if ( + not explicit + and ( + not settings.RETRIEVAL_PAGE_MEMORY_ENABLED + or file_extension.lower() not in _PAGE_MEMORY_PARSE_TRACK_EXTENSIONS + ) + ): + return "chunk" + + _validate_parse_track_for_extension( + parse_track=parse_track, + file_extension=file_extension, + ) + return parse_track diff --git a/apps/api/tests/contract/test_page_memory_parse_track_contract.py b/apps/api/tests/contract/test_page_memory_parse_track_contract.py index c670e6dd..b6a03fb8 100644 --- a/apps/api/tests/contract/test_page_memory_parse_track_contract.py +++ b/apps/api/tests/contract/test_page_memory_parse_track_contract.py @@ -12,6 +12,7 @@ os.environ.setdefault("S3_TEMP_PATH", "/tmp") from shared.core.exceptions.domain_exceptions import ValidationException +from shared.models.schemas.job import JobCreate def test_page_memory_parse_track_rejects_when_flag_disabled(monkeypatch) -> None: @@ -58,3 +59,79 @@ def test_page_memory_parse_track_allows_only_pdf_and_pptx(monkeypatch) -> None: parse_track="page_memory", file_extension=".docx", ) + + +def test_implicit_parse_track_defaults_to_page_memory_for_supported_files( + monkeypatch, +) -> None: + from app.services.document_ingestion.service import _apply_effective_parse_track + from shared.core.config import settings + + monkeypatch.setattr( + settings, + "RETRIEVAL_PAGE_MEMORY_ENABLED", + True, + ) + payload = JobCreate(source_type="file", file_name="policy.pdf") + + effective = _apply_effective_parse_track(payload, file_extension=".pdf") + + assert effective == "page_memory" + assert payload.parse_track == "page_memory" + + +def test_implicit_parse_track_falls_back_to_chunk_for_unsupported_files( + monkeypatch, +) -> None: + from app.services.document_ingestion.service import _apply_effective_parse_track + from shared.core.config import settings + + monkeypatch.setattr( + settings, + "RETRIEVAL_PAGE_MEMORY_ENABLED", + True, + ) + payload = JobCreate(source_type="file", file_name="policy.docx") + + effective = _apply_effective_parse_track(payload, file_extension=".docx") + + assert effective == "chunk" + assert payload.parse_track == "chunk" + + +def test_implicit_parse_track_falls_back_to_chunk_when_flag_disabled( + monkeypatch, +) -> None: + from app.services.document_ingestion.service import _apply_effective_parse_track + from shared.core.config import settings + + monkeypatch.setattr( + settings, + "RETRIEVAL_PAGE_MEMORY_ENABLED", + False, + ) + payload = JobCreate(source_type="file", file_name="policy.pdf") + + effective = _apply_effective_parse_track(payload, file_extension=".pdf") + + assert effective == "chunk" + assert payload.parse_track == "chunk" + + +def test_explicit_page_memory_still_rejects_unsupported_files(monkeypatch) -> None: + from app.services.document_ingestion.service import _apply_effective_parse_track + from shared.core.config import settings + + monkeypatch.setattr( + settings, + "RETRIEVAL_PAGE_MEMORY_ENABLED", + True, + ) + payload = JobCreate( + source_type="file", + file_name="policy.docx", + parse_track="page_memory", + ) + + with pytest.raises(ValidationException): + _apply_effective_parse_track(payload, file_extension=".docx") diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index ac5d7876..4c189c29 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -187,7 +187,7 @@ async def fake_retrieval_run( { "chunk_id": policy_document["chunk_id"], "document_id": policy_document["document_id"], - "chunk_type": "text", + "chunk_type": "page", "section_path": policy_document["section_path"], "file_path": "", "job_id": policy_document["job_id"], @@ -203,6 +203,7 @@ async def fake_retrieval_run( source_file_name="policy.pdf", section_path="Root/Policy", content="policy marker content", + chunk_type="page", ) await _seed_retrieval_document( user_id="local-dev-user", @@ -210,6 +211,7 @@ async def fake_retrieval_run( source_file_name="filler.pdf", section_path="Root/Filler", content="filler content", + chunk_type="page", ) monkeypatch.setattr( "shared.services.retrieval.workflow.planner.QueryPlanner.plan", @@ -226,7 +228,7 @@ async def fake_retrieval_run( "namespace": "contract-agentic-request-policy", "query": "policy marker", "top_k": 1, - "data_type": 2, + "data_type": 7, "signal_paths": ["Root"], "filter_mode": "keep", "channels": ["content"], @@ -248,7 +250,7 @@ async def fake_retrieval_run( assert request["top_k"] == 1 assert request["exclude_document_ids"] == [] assert request["exclude_sections"] == [] - assert request["data_type"] == 2 + assert request["data_type"] == 7 assert request["signal_paths"] == ["Root"] assert request["filter_mode"] == "keep" assert request["channels"] == ["content"] diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index a55c5520..f082782b 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -114,7 +114,6 @@ def ensure_doc_nav_json( source_file_name: str = "", *, overwrite: bool = False, - skeletons: List[Dict[str, Any]] | None = None, ) -> str: """Materialize ``doc_nav.json`` from chunks when the parser did not emit one.""" nav_path = os.path.join(file_dir, DOC_NAV_FILENAME) @@ -124,14 +123,7 @@ def ensure_doc_nav_json( from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder schema = ZipResultSchemaBuilder() - if skeletons: - doc_nav = schema.build_doc_nav_from_skeletons( - skeletons, - chunks, - source_file_name, - ) - else: - doc_nav = schema.build_doc_nav(chunks, source_file_name) + doc_nav = schema.build_doc_nav(chunks, source_file_name) _save_doc_nav(file_dir, doc_nav) return nav_path diff --git a/apps/worker/app/services/document_ingestion/artifact_refs.py b/apps/worker/app/services/document_ingestion/artifact_refs.py new file mode 100644 index 00000000..db554676 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/artifact_refs.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from typing import Any + + +def collect_referenced_artifact_refs(chunks: list[dict[str, Any]]) -> set[str]: + """Collect client-visible artifact refs referenced by persisted chunks.""" + refs: set[str] = set() + for chunk in chunks: + chunk_type = str(chunk.get("type") or chunk.get("chunk_type") or "").strip().lower() + metadata = chunk.get("metadata") or {} + if not isinstance(metadata, dict): + metadata = {} + + if chunk_type == "page": + raw_page_refs = metadata.get("page_image_uris") or [] + if isinstance(raw_page_refs, list): + for raw_ref in raw_page_refs: + _add_artifact_ref(refs, raw_ref, allowed_roots={"pages"}) + continue + + if chunk_type in {"image", "table"}: + allowed_root = f"{chunk_type}s" + _add_artifact_ref( + refs, + metadata.get("file_path") or chunk.get("file_path") or chunk.get("path"), + allowed_roots={allowed_root}, + ) + return refs + + +def _add_artifact_ref( + refs: set[str], + raw_ref: object, + *, + allowed_roots: set[str], +) -> None: + normalized = _normalize_client_artifact_ref(raw_ref) + if normalized is None: + return + root = normalized.split("/", 1)[0] + if root in allowed_roots: + refs.add(normalized) + + +def _normalize_client_artifact_ref(raw_ref: object) -> str | None: + if raw_ref is None: + return None + normalized = str(raw_ref).strip().replace("\\", "/").lstrip("/") + parts = [ + part + for part in normalized.split("/") + if part and part not in {".", ".."} + ] + if len(parts) < 2: + return None + if parts[0] not in {"pages", "images", "tables"}: + return None + return "/".join(parts) diff --git a/apps/worker/app/services/document_ingestion/success_finalization.py b/apps/worker/app/services/document_ingestion/success_finalization.py index 91f9691d..9af4e7e6 100644 --- a/apps/worker/app/services/document_ingestion/success_finalization.py +++ b/apps/worker/app/services/document_ingestion/success_finalization.py @@ -17,6 +17,7 @@ ParseResultPackage, build_generated_result_package, ) +from app.services.document_ingestion.artifact_refs import collect_referenced_artifact_refs from app.services.document_ingestion.processing_context import ParseJobContext from loguru import logger @@ -138,14 +139,12 @@ def _enrich_document_navigation( section_summaries: dict[str, str] = {} add_dir = artifact.add_dir parsed_contents_df = artifact.dataframe - skeletons = _get_page_memory_skeletons(parsed_contents_df) if add_dir and source_file_name: if "path" in parsed_contents_df.columns: ensure_doc_nav_json( str(add_dir), chunks, source_file_name=source_file_name, - skeletons=skeletons, ) try: document_root_for_enrich = os.path.dirname(str(add_dir)) @@ -196,13 +195,6 @@ def _attach_document_top_summary( metadata["document_top_summary"] = document_top_summary -def _get_page_memory_skeletons(parsed_df: Any) -> list[dict[str, Any]]: - raw_skeletons = getattr(parsed_df, "attrs", {}).get("page_memory_skeletons") - if not isinstance(raw_skeletons, list): - return [] - return [skel for skel in raw_skeletons if isinstance(skel, dict)] - - def _record_processing_completion( *, job_id: str, @@ -263,5 +255,6 @@ def _upload_result_package( if result_package.artifact.add_dir else "", zip_file_path=generated_package.zip_file_path, + artifact_refs=collect_referenced_artifact_refs(result_package.chunks), ) return result_bundle.zip_key diff --git a/apps/worker/app/services/document_parser/assets/inline_asset.py b/apps/worker/app/services/document_parser/assets/inline_asset.py index 1d67eae0..732aaa1a 100644 --- a/apps/worker/app/services/document_parser/assets/inline_asset.py +++ b/apps/worker/app/services/document_parser/assets/inline_asset.py @@ -33,8 +33,10 @@ def build_table_asset_row( know_id: str, addtime: str, ) -> ParsedRow: + del content + row_content = relative_path return ParsedRow( - content=content, + content=row_content, path=relative_path, type="table", keywords=keywords, @@ -43,5 +45,5 @@ def build_table_asset_row( tokens="", connectto="", addtime=addtime, + length=len(row_content), ) - diff --git a/apps/worker/app/services/document_parser/orchestration/oversized_pdf_policy.py b/apps/worker/app/services/document_parser/orchestration/oversized_pdf_policy.py index e763967b..da4cd810 100644 --- a/apps/worker/app/services/document_parser/orchestration/oversized_pdf_policy.py +++ b/apps/worker/app/services/document_parser/orchestration/oversized_pdf_policy.py @@ -59,6 +59,28 @@ def build_oversized_pdf_processing_failed_exception( ) +def build_oversized_pdf_profile_failed_exception( + *, + page_count: int, + original_exception: Exception, +) -> PDFParsingException: + reason = _format_original_error(original_exception) + return PDFParsingException( + user_message=( + "Oversized PDF page-memory profiling failed: " + f"{reason}. This document has {page_count} pages and requires a " + "successful structural profile before page-memory parsing can continue." + ), + reason="OVERSIZED_PAGE_MEMORY_PROFILE_FAILED", + internal_message=( + "Oversized PDF page-memory structural profile failed. " + f"page_count={page_count}, " + f"original={type(original_exception).__name__}: {original_exception}" + ), + original_exception=original_exception, + ) + + def _build_standard_page_limit_exception(page_count: int) -> ValidationException: page_limit = settings.MAX_PDF_PAGE_LIMIT return ValidationException( diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index b6d5d77f..a069b235 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -4,11 +4,12 @@ import os from contextlib import contextmanager -from typing import Any, Iterator +from typing import Any, Iterator, Literal from loguru import logger from app.services.document_agent.coordinator import ProfileCoordinator from app.services.document_parser.orchestration.oversized_pdf_policy import ( + build_oversized_pdf_profile_failed_exception, build_oversized_pdf_processing_failed_exception, raise_if_oversized_pdf_not_supported, ) @@ -30,6 +31,7 @@ def profile_document( job_id: str | None = None, output_dir: str | None = None, skip_shard_plan: bool = False, + oversized_policy: Literal["chunk", "page_memory"] = "chunk", ) -> ParserDocumentProfile: """ General document profiling entry point. @@ -43,6 +45,9 @@ def profile_document( LLM shard decision (+ H2 refinement) and populates a single-shard placeholder instead. Used by the page-memory track, which never consumes the shard plan. Chunk-track keeps the default (False). + oversized_policy: Controls oversized PDF admission. ``chunk`` applies + the legacy MinerU shard gate, while ``page_memory`` lets the + page-memory track continue to structural profiling. Returns: ParserDocumentProfile @@ -58,6 +63,7 @@ def profile_document( job_id=job_id, output_dir=output_dir, skip_shard_plan=skip_shard_plan, + oversized_policy=oversized_policy, ) return ParserDocumentProfile( @@ -75,6 +81,7 @@ def _profile_pdf( job_id: str | None, output_dir: str | None, skip_shard_plan: bool = False, + oversized_policy: Literal["chunk", "page_memory"] = "chunk", ) -> ParserDocumentProfile: with _profile_db_context(enabled=bool(job_id)) as db: return _profile_pdf_with_db( @@ -84,6 +91,7 @@ def _profile_pdf( output_dir=output_dir, db=db, skip_shard_plan=skip_shard_plan, + oversized_policy=oversized_policy, ) @@ -95,6 +103,7 @@ def _profile_pdf_with_db( output_dir: str | None, db: Any | None, skip_shard_plan: bool = False, + oversized_policy: Literal["chunk", "page_memory"] = "chunk", ) -> ParserDocumentProfile: profile_job_id = job_id or filename agent_output_dir = os.path.join(output_dir, "_doc_agent") if output_dir else None @@ -134,12 +143,18 @@ def _profile_pdf_with_db( }, ) if profile.page_count > settings.MAX_PDF_PAGE_LIMIT: - raise_if_oversized_pdf_not_supported(page_count=profile.page_count) + if oversized_policy != "page_memory": + raise_if_oversized_pdf_not_supported(page_count=profile.page_count) if not profile.is_atlas: try: profile.anatomy = coordinator.run_structural() profile.toc = _map_toc_profile(coordinator) except Exception as exc: + if oversized_policy == "page_memory": + raise build_oversized_pdf_profile_failed_exception( + page_count=profile.page_count, + original_exception=exc, + ) from exc raise build_oversized_pdf_processing_failed_exception( page_count=profile.page_count, original_exception=exc, diff --git a/apps/worker/app/services/document_parser/tables/table_asset_writer.py b/apps/worker/app/services/document_parser/tables/table_asset_writer.py index 6ddeb57e..920e0e97 100644 --- a/apps/worker/app/services/document_parser/tables/table_asset_writer.py +++ b/apps/worker/app/services/document_parser/tables/table_asset_writer.py @@ -30,7 +30,8 @@ def write_table_asset(table_input: TableAssetInput) -> ParsedRow: table_path = os.path.join(table_dir, table_filename) with open(table_path, "w", encoding="utf-8") as table_file: table_file.write(table_input.html) - row_content = table_input.content if table_input.content is not None else table_input.html + asset_ref = table_input.asset_path or f"tables/{table_filename}" + row_content = asset_ref row_type = "table" if table_input.asset_path: row_type = f"table\n{build_chunk_ref(table_input.asset_path)}" @@ -44,7 +45,7 @@ def write_table_asset(table_input: TableAssetInput) -> ParsedRow: tokens=table_input.tokens, connectto="", addtime=table_input.addtime, - length=table_input.length, + length=len(row_content), ) diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index d7dd8609..03a35791 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -50,6 +50,7 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: job_id=request.job_id, output_dir=full_output_dir, skip_shard_plan=True, + oversized_policy="page_memory", ) verdict = _decide_granularity(profile) @@ -359,9 +360,7 @@ def _build_page_dataframe( "[page_memory] C7 assembled {} node rows (verdict={})", len(rows), verdict, ) - df = pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) - df.attrs["page_memory_skeletons"] = [skel.to_dict() for skel in skeletons] - return df + return pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) def _build_page_ctx( @@ -429,12 +428,7 @@ def _build_whole_doc_dataframe( "addtime": get_str_time(), "page_nums": ",".join(str(page) for page in pages), "extra_metadata": { - "granularity": "whole_doc", - "strategy_used": "whole_doc" if verdict == "whole_doc" else "whole_doc_fallback", - "source_verdict": verdict, - "page_index": None, "page_image_uris": page_image_uris, - "status": "clear", }, } return pd.DataFrame([row], columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index ae8399ce..c739d6d4 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -542,14 +542,7 @@ def build_node_rows( "addtime": get_str_time(), "page_nums": ",".join(str(page) for page in view.pages), "extra_metadata": { - "granularity": "node", - "section_path": leaf.section_path, - "section_level": leaf.level, - "page_indices": list(view.pages), - "owned_pages": list(view.owned_pages), "page_image_uris": page_image_uris, - "kind": kind_by_page.get(leaf.start_page, "normal"), - "source_verdict": verdict, }, } rows.append(row) diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index e349f733..fb8e1545 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -1,4 +1,4 @@ -"""Page-memory asset extraction for charts, figures, and tables.""" +"""Page-memory asset extraction for figures and tables.""" from __future__ import annotations @@ -24,8 +24,8 @@ _BUDGET_STAGE = "page_asset_extraction" _GRID_SIZE = 1000 -_VALID_KINDS = {"table", "chart", "figure"} -_DEFAULT_ASSET_MODEL = "qwen3-vl-32b-instruct" +_VALID_KINDS = {"table", "figure"} +_DEFAULT_ASSET_MODEL = "qwen3.6-flash" _ASSET_ANNOTATE_DIR = "asset_annotate" @@ -192,11 +192,11 @@ def detect_page_assets( if bbox_px is None: continue asset_index = len(assets) + 1 + title = str(region.get("title") or region.get("caption") or "").strip() asset_id = "asset_" + gen_str_codes( f"{source_name}::{page.page_index}::{asset_index}::{kind}::{bbox_px}::" - f"{region.get('title') or ''}::{region.get('summary') or ''}" + f"{title}" ) - keywords = _normalize_keywords(region.get("keywords")) assets.append( PageAsset( asset_id=asset_id, @@ -208,11 +208,11 @@ def detect_page_assets( height_px=height_px, width_pt=float(page.width or 0.0), height_pt=float(page.height or 0.0), - caption=str(region.get("caption") or "").strip(), + caption="", confidence=confidence, - title=str(region.get("title") or "").strip(), - summary=str(region.get("summary") or "").strip(), - keywords=keywords, + title=title, + summary="", + keywords=[], extraction_status="detected", ) ) @@ -414,6 +414,8 @@ def annotate_page_assets( color = _asset_box_color(asset.kind) draw.rectangle([x1, y1, x2, y2], outline=color, width=4) label = f"{asset.kind} {asset.confidence:.2f}" + if asset.title: + label += f" | {asset.title[:24]}" draw.text((x1 + 2, max(0, y1 - 24)), label, fill=color, font=font) canvas.save(output_path) except Exception as exc: @@ -429,8 +431,6 @@ def annotate_page_assets( def _asset_box_color(kind: str) -> tuple[int, int, int]: if kind == "table": return (220, 0, 0) - if kind == "chart": - return (0, 90, 220) return (0, 150, 0) @@ -471,6 +471,7 @@ def build_asset_rows( "asset_kind": asset.kind, "bbox_px": asset.bbox_px, "confidence": asset.confidence, + "title": asset.title, "caption": asset.caption, "image_uri": asset.image_uri, "html_uri": asset.html_uri, @@ -490,13 +491,8 @@ def asset_reference(asset: PageAsset) -> str: def _asset_content(asset: PageAsset, *, row_type: str) -> str: - if row_type == "table" and asset.html_path: - try: - return Path(asset.html_path).read_text(encoding="utf-8") - except Exception as exc: - logger.warning( - "[page_assets] failed to read table HTML {}: {}", asset.html_path, exc - ) + if row_type == "table" and asset.html_uri: + return asset.html_uri parts = [asset.title, asset.summary, asset.caption] if asset.image_uri: parts.append(f"[{asset.image_uri}]") diff --git a/apps/worker/experiments/chart_asset_probe.py b/apps/worker/experiments/chart_asset_probe.py index 125c288e..08751c71 100644 --- a/apps/worker/experiments/chart_asset_probe.py +++ b/apps/worker/experiments/chart_asset_probe.py @@ -1,8 +1,8 @@ -"""Experimental: VLM-driven chart/table region detection + cropping. +"""Experimental: VLM-driven asset bbox detection + cropping. Goal of this experiment ----------------------- -Test whether a VLM can directly locate table/chart bounding boxes on a +Test whether a VLM can directly locate table/chart/figure bounding boxes on a *rendered page image* (the same way PAGE-TRACK renders pages), so we can **crop** those regions out and later hand the crop to a dedicated table model (e.g. tabular / table-transformer) instead of asking the VLM to @@ -12,17 +12,14 @@ 1. Renders selected PDF pages to PNG at a fixed DPI (mirrors PAGE-TRACK, default 144 DPI) so the pixel<->point mapping is fully controlled. - 2. (VLM) Asks the model for table/chart regions as normalized [0,1000] - boxes, maps them to pixels, crops, and draws an annotated overlay. - 3. (Baseline) Runs PyMuPDF's free, pixel-accurate detectors - (``find_tables`` + vector ``drawings`` rects + embedded ``images``) - so VLM output can be judged against a non-LLM reference. + 2. (VLM) Asks the model only for table/chart/figure regions as normalized + [0,1000] boxes, maps them to pixels, crops, and draws an annotated overlay. Outputs (under --out): pages/page-N.png full page render crops/page-N_vlm-K_.png VLM crops - crops/page-N_base-K_.png baseline crops - asset_annotate/page_N.png page with VLM(red) + baseline(green) boxes + crops/page-N_ref-K_.png reference crops from existing chunks + asset_annotate/page_N.png page with VLM(red) + reference(green) boxes results.json all regions + metadata report.md human-readable summary @@ -30,11 +27,11 @@ cd apps/worker uv run python experiments/chart_asset_probe.py \ --pdf "/path/to/doc.pdf" \ - --pages all --baseline + --pages all # Uses $IMAGE_MODEL (default qwen3.6-flash) unless --model is set. # Alternate: --model qwen3-vl-32b-instruct (open-weights, local-deployable). -VLM model guidance (chart/table bbox grounding) +VLM model guidance (bbox grounding) ------------------------------------------------ * **Default (cloud):** ``qwen3.6-flash`` via ``$IMAGE_MODEL`` — cheapest, strong bbox quality on our probe PDFs. @@ -67,34 +64,59 @@ # BOTH axes, with origin at the top-left of the page image. This is robust # to whatever internal resize the API performs. NORM = 1000 +_VALID_KINDS = {"table", "figure"} _PROMPT = ( - "You are a precise document layout detector. The attached image is a " - "single rendered page.\n\n" - "Task: locate every TABLE and every FIGURE (bar/line/pie charts, " - "plots, diagrams, schematic images). Do NOT transcribe their content. " - "Only return their bounding boxes.\n\n" - "Coordinate system: treat the page image as a {n}x{n} grid. The " - "top-left corner is (0,0) and the bottom-right is ({n},{n}). For each " - "region return [x1,y1,x2,y2] integers in that 0-{n} space, where " - "(x1,y1) is the top-left and (x2,y2) the bottom-right of a box that " - "TIGHTLY encloses the region INCLUDING its title/caption and axis " - "labels but EXCLUDING surrounding body paragraphs.\n\n" - "Return STRICT JSON only, no prose:\n" - '{{"regions":[{{"type":"table|chart|figure","bbox":[x1,y1,x2,y2],' - '"caption":"short caption text if visible else empty",' - '"confidence":0.0}}]}}\n' - "If there are no tables/charts, return {{\"regions\":[]}}." + "You are a precise document layout detector. The attached image is a single " + "rendered PDF page.\n\n" + "Find visually distinct tables and figures that should become reusable " + "document assets. Locate them only — do NOT summarize, transcribe full " + "content, extract keywords, or read data values. Return strict JSON:\n" + "{{\n" + ' "regions": [\n' + " {{\n" + ' "kind": "table|figure",\n' + ' "bbox": [x1, y1, x2, y2],\n' + ' "title": "",\n' + ' "confidence": 0.0\n' + " }}\n" + " ]\n" + "}}\n\n" + "Coordinate system:\n" + "- Treat the page image as a {n}x{n} grid.\n" + "- Origin is the top-left corner.\n" + "- bbox values must be integers in [0, {n}].\n" + "- bbox must tightly include the whole asset: its title, caption, legend, " + "axes, labels, table headers, and footnotes that belong to that asset.\n" + "- Exclude surrounding body paragraphs, page headers, page footers, and " + "page numbers.\n\n" + "Rules:\n" + "- \"table\": data arranged in clear rows and columns — grid lines, cell " + "borders, or strongly aligned cells (data tables, forms, financial tables, " + "appendix tables).\n" + "- \"figure\": any non-table visual asset — bar/line/pie/scatter charts, " + "plots, diagrams, flowcharts, architecture drawings, schematics, or " + "embedded images.\n" + "- Do not mark ordinary paragraphs, bullet lists, title blocks, or loose " + "multi-line text as tables.\n" + "- Do not split a single coherent table or figure into sub-parts.\n" + "- \"title\" is one short label only (single line). Do not duplicate it into " + "other fields and do not write a summary. Use an empty string when there is " + "no visible title or caption.\n" + "- Use confidence 0.0-1.0. Only include assets you can localize.\n" + "- If there are no qualifying assets, return {{\"regions\":[]}}.\n" + "- Return ONLY the JSON object, no markdown fences or explanations." ).format(n=NORM) @dataclass class Region: - source: str # "vlm" | "baseline" + source: str # "vlm" | "reference" page: int - kind: str # table | chart | figure | drawing | image + kind: str # table | figure bbox_px: list[int] # [x1,y1,x2,y2] in rendered-image pixels - caption: str = "" + title: str = "" # short asset title / caption (single field, no duplication) confidence: float = 0.0 crop_path: str = "" @@ -108,7 +130,7 @@ class PageResult: height_pt: float image_path: str vlm_regions: list[Region] = field(default_factory=list) - baseline_regions: list[Region] = field(default_factory=list) + reference_regions: list[Region] = field(default_factory=list) vlm_error: str = "" @@ -190,59 +212,42 @@ def norm_to_px(box: list[float], w: int, h: int) -> list[int]: return [x1, y1, x2, y2] -# ── PyMuPDF baseline (free, pixel-accurate) ─────────────────────────── +def load_reference_regions(chunks_path: Path) -> dict[int, list[Region]]: + if not chunks_path.exists(): + raise FileNotFoundError(f"reference chunks not found: {chunks_path}") + payload = json.loads(chunks_path.read_text(encoding="utf-8")) + chunks = payload.get("chunks") if isinstance(payload, dict) else None + if not isinstance(chunks, list): + raise ValueError(f"reference chunks must contain a chunks[] array: {chunks_path}") - -def baseline_regions(page: fitz.Page, dpi: int, w: int, h: int) -> list[Region]: - zoom = dpi / 72.0 - regions: list[Region] = [] - - # 1) Native table finder - try: - tf = page.find_tables() - for t in tf.tables: - r = t.bbox # (x0,y0,x1,y1) in points - regions.append( - Region( - source="baseline", - page=page.number + 1, - kind="table", - bbox_px=_pt_box_to_px(r, zoom, w, h), - confidence=1.0, - ) + by_page: dict[int, list[Region]] = {} + for chunk in chunks: + if not isinstance(chunk, dict): + continue + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + continue + bbox = metadata.get("bbox_px") + page_index = metadata.get("page_index") + if not isinstance(bbox, list) or len(bbox) != 4 or page_index is None: + continue + try: + page = int(page_index) + bbox_px = [int(round(float(item))) for item in bbox] + except (TypeError, ValueError): + continue + kind = str(metadata.get("asset_kind") or chunk.get("type") or "asset") + by_page.setdefault(page, []).append( + Region( + source="reference", + page=page, + kind=kind, + bbox_px=bbox_px, + title=str(metadata.get("title") or metadata.get("caption") or ""), + confidence=float(metadata.get("confidence") or 0.0), ) - except Exception as exc: # noqa: BLE001 - print(f" [baseline] find_tables failed: {exc}", file=sys.stderr) - - # 2) Embedded raster images (charts often exported as images) - try: - for info in page.get_image_info(): - bbox = info.get("bbox") - if bbox: - regions.append( - Region( - source="baseline", - page=page.number + 1, - kind="image", - bbox_px=_pt_box_to_px(bbox, zoom, w, h), - confidence=0.5, - ) - ) - except Exception as exc: # noqa: BLE001 - print(f" [baseline] get_image_info failed: {exc}", file=sys.stderr) - - return regions - - -def _pt_box_to_px(box: Any, zoom: float, w: int, h: int) -> list[int]: - x1, y1, x2, y2 = box - px = [int(round(x1 * zoom)), int(round(y1 * zoom)), - int(round(x2 * zoom)), int(round(y2 * zoom))] - px[0] = max(0, min(px[0], w)) - px[2] = max(0, min(px[2], w)) - px[1] = max(0, min(px[1], h)) - px[3] = max(0, min(px[3], h)) - return px + ) + return by_page # ── cropping + overlay ──────────────────────────────────────────────── @@ -268,15 +273,17 @@ def draw_overlay(img: Image.Image, page_res: PageResult, out_path: Path) -> None except Exception: # noqa: BLE001 font = ImageFont.load_default() - for r in page_res.baseline_regions: + for r in page_res.reference_regions: x1, y1, x2, y2 = r.bbox_px draw.rectangle([x1, y1, x2, y2], outline=(0, 170, 0), width=3) - draw.text((x1 + 2, y1 + 2), f"base:{r.kind}", fill=(0, 120, 0), font=font) + draw.text((x1 + 2, y1 + 2), f"ref:{r.kind}", fill=(0, 120, 0), font=font) for i, r in enumerate(page_res.vlm_regions): x1, y1, x2, y2 = r.bbox_px draw.rectangle([x1, y1, x2, y2], outline=(220, 0, 0), width=3) label = f"vlm:{r.kind} {r.confidence:.2f}" + if r.title: + label += f" | {r.title[:24]}" draw.text((x1 + 2, max(0, y1 - 24)), label, fill=(200, 0, 0), font=font) canvas.save(out_path) @@ -313,14 +320,18 @@ def main() -> int: "--model", default=os.environ.get("IMAGE_MODEL", ""), help=( - "VLM model for bbox detection. Defaults to $IMAGE_MODEL " + "VLM model for bbox-only detection. Defaults to $IMAGE_MODEL " "(qwen3.6-flash). Alternate: qwen3-vl-32b-instruct " "(open-weights, local-deployable; DashScope ¥2/¥8 per 1M in/out)." ), ) ap.add_argument("--margin", type=int, default=8, help="crop padding px") - ap.add_argument("--baseline", action="store_true", help="also run PyMuPDF baseline") - ap.add_argument("--no-vlm", action="store_true", help="skip VLM (baseline only)") + ap.add_argument("--no-vlm", action="store_true", help="skip VLM (reference only)") + ap.add_argument( + "--reference-chunks", + default="", + help="existing chunks.json with old asset bbox metadata to overlay in green", + ) ap.add_argument("--out", default="") args = ap.parse_args() @@ -342,8 +353,14 @@ def main() -> int: doc = fitz.open(str(pdf_path)) page_nums = parse_pages(args.pages, doc.page_count) + reference_by_page = ( + load_reference_regions(Path(args.reference_chunks)) + if args.reference_chunks + else {} + ) print(f"PDF: {pdf_path.name} | {doc.page_count} pages | probing {len(page_nums)} " - f"| dpi={args.dpi} | model={args.model or '(none)'} | baseline={args.baseline}") + f"| dpi={args.dpi} | model={args.model or '(none)'} " + f"| reference={sum(len(items) for items in reference_by_page.values())}") results: list[PageResult] = [] total_tokens = 0 @@ -357,6 +374,7 @@ def main() -> int: width_pt=page.rect.width, height_pt=page.rect.height, image_path=str(img_path), ) + pr.reference_regions = list(reference_by_page.get(pno, [])) print(f"\n[page {pno}] {w}x{h}px") if not args.no_vlm: @@ -368,11 +386,14 @@ def main() -> int: box = reg.get("bbox") or reg.get("bbox_norm") if not box or len(box) != 4: continue + kind = str(reg.get("kind") or reg.get("type") or "").strip().lower() + if kind not in _VALID_KINDS: + continue region = Region( source="vlm", page=pno, - kind=str(reg.get("type", "region")), + kind=kind, bbox_px=norm_to_px([float(v) for v in box], w, h), - caption=str(reg.get("caption", "")), + title=str(reg.get("title") or reg.get("caption") or ""), confidence=float(reg.get("confidence", 0.0) or 0.0), ) pr.vlm_regions.append(region) @@ -382,19 +403,15 @@ def main() -> int: pr.vlm_error = str(exc) print(f" [vlm] ERROR: {exc}", file=sys.stderr) - if args.baseline: - pr.baseline_regions = baseline_regions(page, args.dpi, w, h) - print(f" [baseline] {len(pr.baseline_regions)} region(s)") - # crops + overlay with Image.open(img_path) as im: for k, r in enumerate(pr.vlm_regions): crop_region(im, r, args.margin, out_dir / "crops" / f"page-{pno}_vlm-{k}_{r.kind}.png") - for k, r in enumerate(pr.baseline_regions): + for k, r in enumerate(pr.reference_regions): crop_region(im, r, args.margin, - out_dir / "crops" / f"page-{pno}_base-{k}_{r.kind}.png") - if pr.vlm_regions or pr.baseline_regions: + out_dir / "crops" / f"page-{pno}_ref-{k}_{r.kind}.png") + if pr.vlm_regions or pr.reference_regions: draw_overlay(im, pr, out_dir / "asset_annotate" / f"page_{pno}.png") results.append(pr) @@ -418,31 +435,32 @@ def main() -> int: def _write_report(out_dir: Path, pdf_path: Path, args: Any, results: list[PageResult], total_tokens: int) -> None: n_vlm = sum(len(r.vlm_regions) for r in results) - n_base = sum(len(r.baseline_regions) for r in results) + n_ref = sum(len(r.reference_regions) for r in results) lines = [ f"# Chart/Table Asset Probe — {pdf_path.name}", "", f"- pages probed: **{len(results)}**", f"- dpi: **{args.dpi}**, model: **{args.model or '(no-vlm)'}**", - f"- VLM regions: **{n_vlm}**, baseline regions: **{n_base}**", + f"- VLM regions: **{n_vlm}**, reference regions: **{n_ref}**", f"- total VLM tokens: **{total_tokens}**", "", - "| page | px | vlm | base | vlm kinds | vlm error |", - "|-----:|----|----:|-----:|-----------|-----------|", + "| page | px | vlm | ref | vlm kinds | vlm titles | vlm error |", + "|-----:|----|----:|----:|-----------|------------|-----------|", ] for r in results: kinds = ",".join(sorted({x.kind for x in r.vlm_regions})) or "-" + titles = "; ".join(x.title for x in r.vlm_regions if x.title) or "-" err = (r.vlm_error[:40] + "…") if r.vlm_error else "" lines.append( f"| {r.page} | {r.width_px}x{r.height_px} | {len(r.vlm_regions)} " - f"| {len(r.baseline_regions)} | {kinds} | {err} |" + f"| {len(r.reference_regions)} | {kinds} | {titles} | {err} |" ) lines += [ "", "## How to read", - "- `asset_annotate/page_N.png`: red = VLM boxes, green = PyMuPDF baseline.", + "- `asset_annotate/page_N.png`: red = new VLM boxes, green = reference boxes from existing chunks.", "- Judge VLM by: does the red box tightly enclose the table/chart " - "(incl. caption, excl. body text)? Compare against green baseline.", + "(incl. caption, excl. body text)? Compare against the green reference.", "- `crops/`: the actual extracted assets to feed a table model next.", ] (out_dir / "report.md").write_text("\n".join(lines), encoding="utf-8") diff --git a/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py b/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py index a3de92fa..fc05f663 100644 --- a/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py +++ b/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py @@ -1,7 +1,7 @@ from shared.services.retrieval.agentic.evidence.renderer import render_leaf_chunks -def test_render_direct_small_table_chunk_includes_asset_url() -> None: +def test_render_direct_table_chunk_uses_summary_and_asset_url() -> None: parts: list[str] = [] render_leaf_chunks( parts, @@ -9,8 +9,12 @@ def test_render_direct_small_table_chunk_includes_asset_url() -> None: { "chunk_id": "table-1", "chunk_type": "table", - "content": "
企业名称
", + "content": "tables/table-企业入驻信息表.html", "file_path": "tables/table-企业入驻信息表.html", + "chunk_metadata": { + "summary": "企业入驻信息登记模板", + "keywords": ["企业信息", "入驻管理"], + }, } ], " ", @@ -19,12 +23,12 @@ def test_render_direct_small_table_chunk_includes_asset_url() -> None: rendered = "\n".join(parts) assert "[Table: http://localhost:4566/table.html?signature=test]" in rendered - assert "
企业名称
" in rendered - + assert "企业入驻信息登记模板" in rendered + assert "企业信息;入驻管理" in rendered + assert " None: - monkeypatch.setenv("RETRIEVAL_AGENTIC_INLINE_TABLE_CHAR_LIMIT", "10") +def test_render_direct_table_chunk_does_not_inline_html() -> None: parts: list[str] = [] render_leaf_chunks( parts, @@ -32,7 +36,7 @@ def test_render_direct_large_table_chunk_uses_asset_stub(monkeypatch) -> None: { "chunk_id": "table-1", "chunk_type": "table", - "content": "
企业名称
", + "content": "tables/table-企业入驻信息表.html", "file_path": "tables/table-企业入驻信息表.html", "source_chunk_path": "企业信息汇总260509 (1).xlsx/企业批量录入", "chunk_metadata": { @@ -49,7 +53,6 @@ def test_render_direct_large_table_chunk_uses_asset_stub(monkeypatch) -> None: assert "[Table: http://localhost:4566/table.html?signature=test]" in rendered assert "Table path: 企业信息汇总260509 (1).xlsx/企业批量录入" in rendered assert "Table asset: tables/table-企业入驻信息表.html" in rendered - assert "Large table omitted from evidence_text" in rendered assert "table-企业批量录入" in rendered assert "Main columns:" in rendered assert "企业信息;入驻管理" in rendered @@ -77,8 +80,9 @@ def test_render_connected_table_chunk_includes_asset_url() -> None: { "chunk_id": "table-1", "chunk_type": "table", - "content": "
A
", + "content": "tables/table-1.html", "file_path": "tables/table-1.html", + "chunk_metadata": {"summary": "A 表摘要"}, }, ], " ", @@ -87,4 +91,38 @@ def test_render_connected_table_chunk_includes_asset_url() -> None: rendered = "\n".join(parts) assert "[Table: http://localhost:4566/table-1.html?signature=test]" in rendered - assert "
A
" in rendered + assert "A 表摘要" in rendered + assert " None: + parts: list[str] = [] + render_leaf_chunks( + parts, + [ + { + "chunk_id": "page-node-1", + "chunk_type": "page", + "content": "RAW OCR SHOULD NOT LEAK", + "chunk_metadata": { + "summary": "制度标准总则摘要", + "page_nums": [225, 226], + "page_image_uris": ["pages/page-225.png", "pages/page-226.png"], + }, + } + ], + " ", + asset_lookup={ + "page-node-1": [ + "http://localhost:4566/page-225.png?signature=test", + "http://localhost:4566/page-226.png?signature=test", + ] + }, + ) + + rendered = "\n".join(parts) + assert "Pages: 225, 226" in rendered + assert "制度标准总则摘要" in rendered + assert "Page image 1: http://localhost:4566/page-225.png?signature=test" in rendered + assert "Page image 2: http://localhost:4566/page-226.png?signature=test" in rendered + assert "RAW OCR SHOULD NOT LEAK" not in rendered diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 85fadef9..353d0653 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -672,6 +672,53 @@ def run_lightweight_anatomy(self, *, skip_shard_plan: bool = False): assert profile.anatomy is fake_anatomy +def test_page_memory_profile_bypasses_chunk_oversized_gate( + monkeypatch, + tmp_path: Path, +) -> None: + fake_anatomy = object() + fake_instances = [] + + class FakeCoordinator: + def __init__(self, **_kwargs) -> None: + self.calls: list[str] = [] + self.blackboard = SimpleNamespace( + page_count=201, + doc_stats={"page_count": 201}, + global_signals={}, + toc_result=None, + toc_hierarchies=None, + ) + fake_instances.append(self) + + def run_coarse(self) -> DocumentProfile: + self.calls.append("run_coarse") + return DocumentProfile( + is_scanned=False, + category="Research Report", + routing_category=PdfRoutingCategory.GENERIC.value, + ) + + def run_structural(self): + self.calls.append("run_structural") + return fake_anatomy + + monkeypatch.setattr(doc_profiler, "ProfileCoordinator", FakeCoordinator) + monkeypatch.setattr(doc_profiler.settings, "MAX_PDF_PAGE_LIMIT", 200) + monkeypatch.setattr(doc_profiler.settings, "OVERSIZED_PDF_SHARD_ENABLED", False) + + profile = profile_document( + str(tmp_path / "oversized.pdf"), + "oversized.pdf", + job_id="job-page-memory-oversized", + output_dir=str(tmp_path), + oversized_policy="page_memory", + ) + + assert profile.anatomy is fake_anatomy + assert fake_instances[0].calls == ["run_coarse", "run_structural"] + + def test_standard_pdf_profile_maps_page_toc_evidence( monkeypatch, tmp_path: Path, diff --git a/apps/worker/tests/contract/test_document_agent_budget_contract.py b/apps/worker/tests/contract/test_document_agent_budget_contract.py index 2dfc47e4..db111193 100644 --- a/apps/worker/tests/contract/test_document_agent_budget_contract.py +++ b/apps/worker/tests/contract/test_document_agent_budget_contract.py @@ -75,9 +75,7 @@ def test_dataframe_converter_accepts_page_chunks_with_extra_metadata() -> None: "addtime": "2026-06-11 00:00:00", "page_nums": "1,2", "extra_metadata": { - "granularity": "whole_doc", "page_image_uris": ["pages/page_page_1.png"], - "page_nums": [99], }, } ] @@ -86,8 +84,8 @@ def test_dataframe_converter_accepts_page_chunks_with_extra_metadata() -> None: chunks = dataframe_to_chunks(df) assert chunks[0]["type"] == "page" - assert chunks[0]["metadata"]["granularity"] == "whole_doc" assert chunks[0]["metadata"]["page_nums"] == [1, 2] + assert chunks[0]["metadata"]["page_image_uris"] == ["pages/page_page_1.png"] def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: @@ -100,14 +98,7 @@ def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: "metadata": { "summary": "page summary", "page_nums": [231], - "granularity": "node", - "page_indices": [231, 232], - "owned_pages": [232], - "section_path": "demo.pdf/3 基本规定/3.2 管理规定", - "section_level": 2, "page_image_uris": ["pages/page-231.png", "pages/page-232.png"], - "kind": "normal", - "source_verdict": "page", }, } ] @@ -120,13 +111,11 @@ def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: metadata = formatted[0]["metadata"] assert metadata["page_nums"] == [231] - assert metadata["granularity"] == "node" - assert metadata["page_indices"] == [231, 232] - assert metadata["owned_pages"] == [232] - assert metadata["section_path"] == "demo.pdf/3 基本规定/3.2 管理规定" - assert metadata["section_level"] == 2 assert metadata["page_image_uris"] == ["pages/page-231.png", "pages/page-232.png"] - assert metadata["source_verdict"] == "page" + assert "granularity" not in metadata + assert "page_indices" not in metadata + assert "owned_pages" not in metadata + assert "section_path" not in metadata def test_page_memory_granularity_routes_supported_page_modes() -> None: diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py index 694992d6..cdaf54d9 100644 --- a/apps/worker/tests/contract/test_excel_parser_contract.py +++ b/apps/worker/tests/contract/test_excel_parser_contract.py @@ -55,8 +55,7 @@ def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( assert "[tables/table-Visible.html]" in parsed_df["type"].iloc[0] assert parsed_df["path"].tolist() == ["budget.xlsx/Visible"] assert parsed_df["summary"].tolist() == ["table-Visible"] - assert parsed_df["content"].iloc[0].lstrip().startswith(" None: +def test_page_assets_default_to_qwen_flash_and_full_page_scan(monkeypatch) -> None: monkeypatch.delenv("PAGE_MEMORY_ASSET_MODEL", raising=False) monkeypatch.delenv("PAGE_MEMORY_ASSET_MAX_PAGES", raising=False) - assert page_assets.get_asset_model() == "qwen3-vl-32b-instruct" + assert page_assets.get_asset_model() == "qwen3.6-flash" assert page_assets.get_asset_max_pages(301) == 301 @@ -37,10 +37,7 @@ def chat_completion_with_usage(self, **kwargs): { "kind": "table", "bbox": [100, 100, 300, 300], - "caption": "caption", - "title": "title", - "summary": "summary", - "keywords": ["table"], + "title": "table title", "confidence": 1.0, } ] @@ -74,9 +71,17 @@ def chat_completion_with_usage(self, **kwargs): confidence_threshold=0.3, ) - assert captured["model"] == "qwen3-vl-32b-instruct" + prompt = captured["messages"][0]["content"][0]["text"] # type: ignore[index] + assert captured["model"] == "qwen3.6-flash" assert captured["temperature"] == 0 assert assets[0].bbox_px == [10, 12, 30, 36] + assert assets[0].title == "table title" + assert assets[0].caption == "" + assert '"kind": "table|figure"' in prompt + assert '"caption"' not in prompt + assert '"summary"' not in prompt + assert '"keywords"' not in prompt + assert '"table|chart|figure"' not in prompt def test_page_asset_extraction_keeps_debug_output_minimal( diff --git a/apps/worker/tests/contract/test_page_memory_navigation_contract.py b/apps/worker/tests/contract/test_page_memory_navigation_contract.py index 75160aa4..23e60c28 100644 --- a/apps/worker/tests/contract/test_page_memory_navigation_contract.py +++ b/apps/worker/tests/contract/test_page_memory_navigation_contract.py @@ -9,19 +9,13 @@ os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") os.environ.setdefault("S3_TEMP_PATH", "/tmp") -import pandas as pd - -from shared.services.storage.zip_doc_navigation import build_doc_nav_from_skeletons +from shared.services.storage.zip_doc_navigation import ZipDocNavigationBuilder from shared.services.storage.zip_result_service import ZipResultService -def test_doc_nav_from_skeletons_preserves_same_page_sibling_sections() -> None: - skeletons = _same_page_sibling_skeletons() - chunks = _page_chunks() - - doc_nav = build_doc_nav_from_skeletons( - skeletons, - chunks, +def test_page_doc_nav_uses_path_based_leaf_summaries_and_page_counts() -> None: + doc_nav = ZipDocNavigationBuilder().build_doc_nav( + _page_chunks(), "demo.pdf", ) @@ -29,7 +23,7 @@ def test_doc_nav_from_skeletons_preserves_same_page_sibling_sections() -> None: assert len(sections) == 1 parent = sections[0] assert parent["title"] == "3 基本规定" - assert parent["chunk_count"] == 2 + assert parent["chunk_count"] == 3 child_counts = { child["title"]: child["chunk_count"] for child in parent["children"] @@ -39,52 +33,19 @@ def test_doc_nav_from_skeletons_preserves_same_page_sibling_sections() -> None: "3.2 管理规定": 2, } + child_summaries = { + child["title"]: child["summary"] for child in parent["children"] + } + assert child_summaries == { + "3.1 职责": "page 231 summary", + "3.2 管理规定": "pages 232-233 summary", + } -def test_doc_nav_from_skeletons_synthesizes_missing_ancestor_sections() -> None: - skeletons = [ - { - "section_path": "demo.pdf/安全类/SJSYJ-SC103/3 基本规定/3.1 职责", - "level": 4, - "start_page": 231, - "end_page": 231, - "title": "3.1 职责", - "parent_path": "demo.pdf/安全类/SJSYJ-SC103/3 基本规定", - }, - { - "section_path": "demo.pdf/安全类/SJSYJ-SC103/3 基本规定/3.2 管理规定", - "level": 4, - "start_page": 231, - "end_page": 232, - "title": "3.2 管理规定", - "parent_path": "demo.pdf/安全类/SJSYJ-SC103/3 基本规定", - }, - ] - - doc_nav = build_doc_nav_from_skeletons( - skeletons, - _page_chunks(), - "demo.pdf", - ) - - assert doc_nav["sections"][0]["title"] == "安全类" - assert doc_nav["sections"][0]["chunk_count"] == 2 - standard = doc_nav["sections"][0]["children"][0] - basic = standard["children"][0] - assert basic["title"] == "3 基本规定" - assert [child["title"] for child in basic["children"]] == [ - "3.1 职责", - "3.2 管理规定", - ] - - -def test_zip_result_service_uses_page_memory_skeletons_for_manifest_hierarchy() -> None: - parsed_df = pd.DataFrame() - parsed_df.attrs["page_memory_skeletons"] = _same_page_sibling_skeletons() +def test_zip_result_service_builds_navigation_from_chunk_paths() -> None: doc_nav, hierarchy = ZipResultService()._build_navigation_outputs( # noqa: SLF001 formatted_chunks=_page_chunks(), source_file_name="demo.pdf", - parsed_df=parsed_df, ) assert doc_nav is not None @@ -96,49 +57,20 @@ def test_zip_result_service_uses_page_memory_skeletons_for_manifest_hierarchy() } -def _same_page_sibling_skeletons() -> list[dict[str, object]]: - return [ - { - "section_path": "demo.pdf/3 基本规定", - "level": 1, - "start_page": 231, - "end_page": 232, - "title": "3 基本规定", - "parent_path": None, - }, - { - "section_path": "demo.pdf/3 基本规定/3.1 职责", - "level": 2, - "start_page": 231, - "end_page": 231, - "title": "3.1 职责", - "parent_path": "demo.pdf/3 基本规定", - }, - { - "section_path": "demo.pdf/3 基本规定/3.2 管理规定", - "level": 2, - "start_page": 231, - "end_page": 232, - "title": "3.2 管理规定", - "parent_path": "demo.pdf/3 基本规定", - }, - ] - - def _page_chunks() -> list[dict[str, object]]: return [ { "chunk_id": "page-231", "type": "page", "content": "page 231", - "path": "demo.pdf/3 基本规定/3.2 管理规定", + "path": "demo.pdf/3 基本规定/3.1 职责", "metadata": {"page_nums": [231], "summary": "page 231 summary"}, }, { - "chunk_id": "page-232", + "chunk_id": "page-232-233", "type": "page", - "content": "page 232", + "content": "pages 232 and 233", "path": "demo.pdf/3 基本规定/3.2 管理规定", - "metadata": {"page_nums": [232], "summary": "page 232 summary"}, + "metadata": {"page_nums": [232, 233], "summary": "pages 232-233 summary"}, }, ] 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 6b77632d..d8f4eb2d 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 @@ -117,7 +117,7 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: assert leaf_a["page_nums"] == "231" assert leaf_a["content"] == "text-231" assert leaf_a["extra_metadata"]["page_image_uris"] == ["pages/page-231.png"] - assert leaf_a["extra_metadata"]["granularity"] == "node" + assert set(leaf_a["extra_metadata"]) == {"page_image_uris"} leaf_b = by_path["demo.pdf/3 基本规定/3.2 管理规定"] assert leaf_b["page_nums"] == "231,232" @@ -127,7 +127,7 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: "pages/page-231.png", "pages/page-232.png", ] - assert leaf_b["extra_metadata"]["owned_pages"] == [232] + assert set(leaf_b["extra_metadata"]) == {"page_image_uris"} def test_build_node_rows_keeps_internal_section_body_pages() -> None: @@ -262,6 +262,7 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: assert [row["type"] for row in rows] == ["table", "page", "page"] assert rows[0]["path"] == "tables/table_page_231_1.html" + assert rows[0]["content"] == "tables/table_page_231_1.html" assert rows[0]["know_id"] == "asset_table_1" by_path = {row["path"]: row for row in rows} @@ -276,10 +277,10 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: def test_page_connectto_normalizes_to_asset_chunk_id() -> None: rows = [ { - "content": "
A
", + "content": "tables/table_page_1_1.html", "path": "tables/table_page_1_1.html", "type": "table", - "length": 34, + "length": 26, "keywords": "asset", "summary": "asset summary", "know_id": "asset_table_1", diff --git a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py new file mode 100644 index 00000000..046ba891 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -0,0 +1,306 @@ +from __future__ import annotations + +import os + +import pytest + +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") + +from shared.services.retrieval.hydration.assets import ( # noqa: E402 + build_retrieval_asset_url_map, + enrich_rows_with_retrieval_asset_urls, +) +from shared.services.retrieval.hydration.result_assembly import ( # noqa: E402 + assemble_retrieval_results, +) +from shared.services.retrieval.search.lexical_text import ( # noqa: E402 + build_content_lexical_text, + build_content_search_text, + build_term_search_text, +) +from shared.services.retrieval.settings import resolve_allowed_chunk_types # noqa: E402 +from shared.services.retrieval.execution.reference_resolver import ( # noqa: E402 + resolve_workflow_references, +) +from shared.services.storage.result_storage import JobResultStorage # noqa: E402 + + +def test_data_type_one_allows_page_and_page_content_enters_search_text() -> None: + page_chunk = { + "type": "page", + "content": "安全风险分级管控 raw pymupdf content", + "metadata": {"summary": "node summary is not the primary content"}, + } + + assert resolve_allowed_chunk_types(1) is None + content_search_text = build_content_search_text(page_chunk) or "" + assert "安全" in content_search_text + assert "风险" in content_search_text + assert "安全风险" in (build_term_search_text(page_chunk, path_text="安全类") or "") + + +def test_data_type_seven_is_page_only() -> None: + assert resolve_allowed_chunk_types(7) == {"page"} + + +def test_table_search_text_uses_summary_keywords_and_caption_not_content() -> None: + table_chunk = { + "type": "table", + "content": "tables/table-1.html", + "metadata": { + "summary": "企业入驻信息登记模板", + "keywords": ["企业名称", "统一社会信用代码"], + "caption": "批量录入表", + }, + } + + content_search_text = build_content_search_text(table_chunk) or "" + content_lexical_text = build_content_lexical_text(table_chunk) or "" + term_search_text = build_term_search_text( + table_chunk, + path_text="企业批量录入", + ) or "" + + assert "企业" in content_search_text + assert "信用" in content_search_text + assert "批量录入表" in content_lexical_text + assert "企业批量录入" in term_search_text + assert "tables/table-1.html" not in content_search_text + + +@pytest.mark.asyncio +async def test_page_result_assembly_uses_summary_not_raw_content() -> None: + rows = [ + { + "chunk_id": "page-node-1", + "chunk_type": "page", + "content": "RAW OCR SHOULD NOT LEAK", + "chunk_metadata": { + "summary": "制度标准总则摘要", + "page_image_uris": ["pages/page-225.png"], + }, + } + ] + + assembled = await assemble_retrieval_results( + rows=rows, + exclude_document_ids=[], + exclude_sections=[], + ) + + assert assembled[0]["content"] == "制度标准总则摘要" + + +@pytest.mark.asyncio +async def test_table_result_assembly_uses_summary_not_html() -> None: + rows = [ + { + "chunk_id": "text-1", + "chunk_type": "text", + "content": "见表 [tables/table-1.html]", + "chunk_metadata": { + "connect_to": [ + { + "target": "table-1", + "relation": "embeds", + "ref": "[tables/table-1.html]", + } + ] + }, + }, + { + "chunk_id": "table-1", + "chunk_type": "table", + "content": "
SHOULD NOT LEAK
", + "file_path": "tables/table-1.html", + "asset_url": "https://assets.example.com/job-1/tables/table-1.html", + "chunk_metadata": { + "summary": "企业入驻信息登记模板", + "keywords": ["企业名称", "统一社会信用代码"], + }, + }, + ] + + assembled = await assemble_retrieval_results( + rows=rows, + exclude_document_ids=[], + exclude_sections=[], + ) + + assert len(assembled) == 1 + content = assembled[0]["content"] + assert "[Table: https://assets.example.com/job-1/tables/table-1.html]" in content + assert "企业入驻信息登记模板" in content + assert "企业名称;统一社会信用代码" in content + assert "SHOULD NOT LEAK" not in content + assert " None: + class FakeResultStorage: + def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: + if not artifact_ref: + return None + normalized = artifact_ref.strip().replace("\\", "/").lstrip("/") + root = normalized.split("/", 1)[0] + if root not in {"images", "tables", "pages"}: + return None + return normalized + + def generate_artifact_url( + self, + *, + job_id: str, + artifact_ref: str, + expires_in: int = 3600, + ) -> str | None: + del expires_in + return f"https://assets.example.com/{job_id}/{artifact_ref}" + + monkeypatch.setattr( + "shared.services.retrieval.hydration.assets.get_result_storage", + lambda: FakeResultStorage(), + ) + + rows = [ + { + "chunk_id": "page-node-1", + "chunk_type": "page", + "job_id": "job-1", + "chunk_metadata": { + "page_image_uris": [ + "pages/page-225.png", + "pages/page-226.png", + "../images/not-allowed.png", + "pages/page-225.png", + ] + }, + } + ] + + enriched = await enrich_rows_with_retrieval_asset_urls( + rows, + log_context="contract", + ) + url_map = await build_retrieval_asset_url_map(rows, log_context="contract") + + assert enriched[0]["asset_urls"] == [ + "https://assets.example.com/job-1/pages/page-225.png", + "https://assets.example.com/job-1/pages/page-226.png", + ] + assert url_map["page-node-1"] == enriched[0]["asset_urls"] + + +def test_result_storage_allows_pages_artifact_refs() -> None: + storage = JobResultStorage(results_bucket="test-results") + + assert storage.normalize_artifact_ref("pages/page-225.png") == "pages/page-225.png" + assert storage.normalize_artifact_ref("../pages/page-226.png") == "pages/page-226.png" + + +def test_result_storage_upload_filters_to_referenced_artifacts(tmp_path) -> None: + class FakeStorageAdapter: + def __init__(self) -> None: + self.uploaded_keys: list[str] = [] + + def upload_file(self, local_path: str, key: str, bucket: str | None = None): + del local_path, bucket + self.uploaded_keys.append(key) + return {"key": key} + + def generate_presigned_url(self, *args, **kwargs) -> str: + del args, kwargs + return "https://assets.example.test/file" + + result_dir = tmp_path / "result" + (result_dir / "pages").mkdir(parents=True) + (result_dir / "tables").mkdir() + (result_dir / "pages" / "page-225.png").write_bytes(b"anchored") + (result_dir / "pages" / "page-999.png").write_bytes(b"unanchored") + (result_dir / "tables" / "table-1.html").write_text("
") + (result_dir / "debug.csv").write_text("debug") + zip_path = tmp_path / "result.zip" + zip_path.write_bytes(b"zip") + + adapter = FakeStorageAdapter() + storage = JobResultStorage( + results_bucket="test-results", + storage_adapter=adapter, # type: ignore[arg-type] + ) + + bundle = storage.upload( + job_id="job-1", + result_dir=str(result_dir), + zip_file_path=str(zip_path), + artifact_refs={"pages/page-225.png", "tables/table-1.html"}, + ) + + assert set(bundle.raw_files) == {"pages/page-225.png", "tables/table-1.html"} + assert "results/job-1/pages/page-225.png" in adapter.uploaded_keys + assert "results/job-1/tables/table-1.html" in adapter.uploaded_keys + assert "results/job-1/pages/page-999.png" not in adapter.uploaded_keys + assert "results/job-1/debug.csv" not in adapter.uploaded_keys + + +@pytest.mark.asyncio +async def test_referenced_chunks_get_page_asset_urls_from_hydrated_rows( + monkeypatch, +) -> None: + async def fake_hydrate_referenced_chunk_rows(**_kwargs): + return [ + { + "document_id": "doc-1", + "chunk_id": "page-node-1", + "chunk_type": "page", + "section_path": "安全类 / 1 总则", + "file_path": None, + "chunk_metadata": {"page_image_uris": ["pages/page-225.png"]}, + "job_id": "job-1", + } + ] + + async def fake_enrich_referenced_chunks_with_asset_urls(rows): + enriched = [] + for row in rows: + enriched.append( + { + **row, + "asset_urls": [ + "https://assets.example.com/job-1/pages/page-225.png" + ], + } + ) + return enriched + + monkeypatch.setattr( + "shared.services.retrieval.execution.reference_resolver.hydrate_referenced_chunk_rows", + fake_hydrate_referenced_chunk_rows, + ) + monkeypatch.setattr( + "shared.services.retrieval.execution.reference_resolver.enrich_referenced_chunks_with_asset_urls", + fake_enrich_referenced_chunks_with_asset_urls, + ) + + resolved = await resolve_workflow_references( + db=None, # fake hydrate ignores db + user_id="user-1", + namespace="default", + refs=[ + { + "document_id": "doc-1", + "chunk_id": "page-node-1", + "chunk_type": "page", + "section_path": "安全类 / 1 总则", + } + ], + ) + + assert resolved.refs[0]["asset_urls"] == [ + "https://assets.example.com/job-1/pages/page-225.png" + ] diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index aca6e63d..ff0eacec 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -102,8 +102,8 @@ def test_parse_task_should_process_uploaded_file_through_real_contract_boundarie chunks_payload = result_zip["chunks"] assert len(chunks_payload["chunks"]) == len(job_chunks) assert all(chunk["type"] == "table" for chunk in chunks_payload["chunks"]) - assert any( - chunk["content"].lstrip().startswith("", - "keywords": ";;" -} - -Rules: -- "summary": describe the main content visible on the page in 1-3 sentences. - If the page contains tables, mention the table topic and key columns. - If the page contains figures or charts, describe what they depict. -- "keywords": extract the most important thematic keywords (up to 5), - separated by semicolons ";". Keywords must be in the same language as - the visible page content. -- Return ONLY the JSON object, no markdown fences or extra text. -""" + You are annotating a single PDF page screenshot for a document memory system. + Return strict JSON with exactly these keys: + + { + "summary": "<1-3 sentence summary of the page content>", + "keywords": ";;" + } + + Rules: + - "summary": describe the main content visible on the page in 1-3 sentences. + If the page contains tables, mention the table topic and key columns. + If the page contains figures or charts, describe what they depict. + - "keywords": extract the most important thematic keywords (up to 5), + separated by semicolons ";". Keywords must be in the same language as + the visible page content. + - Return ONLY the JSON object, no markdown fences or extra text. + """ elif task == "page-memory-vlm-title": temperature = 0 top_p = 0.01 max_tokens = kwargs.get("paras", {}).get("max_tokens", 300) prompt = """\ -You are extracting document-outline-level headings from a PDF page screenshot. -Your goal is to find ONLY the headings that would appear in a Table of Contents. -Most pages will have ZERO such headings — returning an empty list is expected -and correct for the majority of pages. - -Return strict JSON: -{ - "titles": [ - { - "text": "", - "prominence": <0.0-1.0>, - "is_in_table": , - "is_in_header_footer": - } - ] -} - -═══ MANDATORY BOOLEAN FLAGS (CRITICAL) ═══ -For EVERY extracted heading, you MUST accurately evaluate these two flags: -1. is_in_table (boolean): Set to `true` if the text is ANYWHERE inside a table. -2. is_in_header_footer (boolean): Set to `true` if the text is located in the top margin (header) or bottom margin (footer) of the page. - -═══ WHAT TO EXTRACT ═══ - -Only extract text that satisfies ALL three criteria: - -1. HEADING FUNCTION (primary — must be true): - The text serves as a TITLE for the body content that follows it. - It introduces or labels a block of subsequent paragraphs, clauses, - or sub-sections. If you removed this text, the following body content - would lose its topic label. - -2. STANDALONE LINE (must be true): - The text occupies its own line, clearly separated from surrounding - body paragraphs. It is NOT inside a table, NOT part of a list, - and NOT embedded within a sentence. - -3. VISUAL DISTINCTION (supporting): - The text is visually set apart from body text — larger font, bold, - centered, or has extra vertical spacing. - -"prominence": 1.0 = most prominent; 0.5 = medium; 0.1 = minor. -Return titles in TOP-TO-BOTTOM order. Text must be EXACT verbatim. - -═══ WHAT TO EXCLUDE (critical — read carefully) ═══ - -1. TABLE CONTENT — Any text that is part of a table. If the - text is surrounded by grid lines, borders, or cell boundaries, or if - its neighboring content is arranged in rows and columns, it is table - content and MUST BE EXCLUDED. This applies even when the text is bold, - large, or spans a merged cell. Specifically exclude: - - Column headers, row category labels, merged-cell group labels - - Any label inside a tabular layout, regardless of visual prominence - -2. PAGE PERIPHERY — Text in margins or corners of the page: - organization/document names repeated as running headers, page numbers, - book/volume titles used as running headers or footers. - -3. BODY TEXT — Numbered clauses, list items, paragraphs, or running - prose, even if bold or indented. - -4. CAPTIONS — Figure/table captions, footnotes. - -5. TOC ENTRIES — If the page is itself a Table of Contents or index, - do NOT extract its listed entries. A TOC page lists other sections - with page numbers — those entries are references, not headings. - -═══ IMPORTANT ═══ -Many pages consist entirely of tables, body text, or appendix forms. -These pages have NO qualifying headings. Return {"titles": []} for them. -Do NOT force-extract table labels or body text as headings. - -Return ONLY the JSON object, no markdown fences. -""" + You are extracting document-outline-level headings from a PDF page screenshot. + Your goal is to find ONLY the headings that would appear in a Table of Contents. + Most pages will have ZERO such headings — returning an empty list is expected + and correct for the majority of pages. + + Return strict JSON: + { + "titles": [ + { + "text": "", + "prominence": <0.0-1.0>, + "is_in_table": , + "is_in_header_footer": + } + ] + } + + ═══ MANDATORY BOOLEAN FLAGS (CRITICAL) ═══ + For EVERY extracted heading, you MUST accurately evaluate these two flags: + 1. is_in_table (boolean): Set to `true` if the text is ANYWHERE inside a table. + 2. is_in_header_footer (boolean): Set to `true` if the text is located in the top margin (header) or bottom margin (footer) of the page. + + ═══ WHAT TO EXTRACT ═══ + + Only extract text that satisfies ALL three criteria: + + 1. HEADING FUNCTION (primary — must be true): + The text serves as a TITLE for the body content that follows it. + It introduces or labels a block of subsequent paragraphs, clauses, + or sub-sections. If you removed this text, the following body content + would lose its topic label. + + 2. STANDALONE LINE (must be true): + The text occupies its own line, clearly separated from surrounding + body paragraphs. It is NOT inside a table, NOT part of a list, + and NOT embedded within a sentence. + + 3. VISUAL DISTINCTION (supporting): + The text is visually set apart from body text — larger font, bold, + centered, or has extra vertical spacing. + + "prominence": 1.0 = most prominent; 0.5 = medium; 0.1 = minor. + Return titles in TOP-TO-BOTTOM order. Text must be EXACT verbatim. + + ═══ WHAT TO EXCLUDE (critical — read carefully) ═══ + + 1. TABLE CONTENT — Any text that is part of a table. If the + text is surrounded by grid lines, borders, or cell boundaries, or if + its neighboring content is arranged in rows and columns, it is table + content and MUST BE EXCLUDED. This applies even when the text is bold, + large, or spans a merged cell. Specifically exclude: + - Column headers, row category labels, merged-cell group labels + - Any label inside a tabular layout, regardless of visual prominence + + 2. PAGE PERIPHERY — Text in margins or corners of the page: + organization/document names repeated as running headers, page numbers, + book/volume titles used as running headers or footers. + + 3. BODY TEXT — Numbered clauses, list items, paragraphs, or running + prose, even if bold or indented. + + 4. CAPTIONS — Figure/table captions, footnotes. + + 5. TOC ENTRIES — If the page is itself a Table of Contents or index, + do NOT extract its listed entries. A TOC page lists other sections + with page numbers — those entries are references, not headings. + + ═══ IMPORTANT ═══ + Many pages consist entirely of tables, body text, or appendix forms. + These pages have NO qualifying headings. Return {"titles": []} for them. + Do NOT force-extract table labels or body text as headings. + + Return ONLY the JSON object, no markdown fences. + """ elif task == "page-memory-hierarchy": temperature = 0 @@ -561,6 +561,7 @@ def build_prompt(task, texts, query, **kwargs): ''' {coarse_context} ''' + """ if coarse_context else "" prompt = f""" You are constructing a fine-grained document hierarchy for ONE already-bounded @@ -626,94 +627,137 @@ def build_prompt(task, texts, query, **kwargs): f"across the provided page image(s) as a single coherent section." ) prompt = f"""\ -You are summarizing one section of a document for a navigation/memory system. -You are given the page screenshot(s) that this section spans. + You are summarizing one section of a document for a navigation/memory system. + You are given the page screenshot(s) that this section spans. -{scope} + {scope} -Return strict JSON with exactly these keys: -{{ - "summary": "<1-4 sentence summary of THIS section's content>", - "keywords": ";;..." -}} + Return strict JSON with exactly these keys: + {{ + "summary": "", + "keywords": ";;..." + }} -Rules: -- "summary": describe what this section is about, in the same language as the - visible page content. If the section is mostly a table, describe the table's - topic and key columns. Do not summarize content that belongs to other - sections on the same page. -- "keywords": up to {kw_num} thematic keywords separated by ";". -- Return ONLY the JSON object, no markdown fences or extra text. -""" + Rules: + - "summary": describe what this section is about, in the same language as the + visible page content. If the section is mostly a table, describe the table's + topic and key columns. Do not summarize content that belongs to other + sections on the same page. + - "keywords": up to {kw_num} thematic keywords separated by ";". + - Return ONLY the JSON object, no markdown fences or extra text. + """ elif task == "page-memory-vlm-ocr": temperature = 0 top_p = 0.01 max_tokens = kwargs.get("paras", {}).get("max_tokens", 1500) prompt = """\ -You are transcribing a single scanned PDF page screenshot for a document -memory system. Extract the page's body text as faithfully as possible. - -Return strict JSON with exactly this key: -{ - "text": "" -} - -Rules: -- Preserve the reading order (top-to-bottom, left-to-right). -- Transcribe tables row by row using a simple readable layout. -- Do NOT add commentary, translation, or summary — transcription only. -- Omit pure decorative running headers/footers and page numbers. -- Return ONLY the JSON object, no markdown fences or extra text. -""" + You are transcribing a single scanned PDF page screenshot for a document + memory system. Extract the page's body text as faithfully as possible. + + Return strict JSON with exactly this key: + { + "text": "" + } + + Rules: + - Preserve the reading order (top-to-bottom, left-to-right). + - Transcribe tables row by row using a simple readable layout. + - Do NOT add commentary, translation, or summary — transcription only. + - Omit pure decorative running headers/footers and page numbers. + - Return ONLY the JSON object, no markdown fences or extra text. + """ elif task == "page-memory-asset-detect": temperature = 0 top_p = 0.01 max_tokens = kwargs.get("paras", {}).get("max_tokens", 1200) grid_size = kwargs.get("paras", {}).get("grid_size", 1000) + # Previous production prompt kept for comparison: + # prompt = f"""\ + # You are a precise document layout detector. The attached image is a single PDF + # page screenshot. + # + # Find visually distinct tables, charts, and figures that should become reusable + # document assets. Return strict JSON: + # {{ + # "regions": [ + # {{ + # "kind": "table|chart|figure", + # "bbox": [x1, y1, x2, y2], + # "caption": "", + # "title": "", + # "summary": "<1 sentence searchable summary>", + # "keywords": ["", ""], + # "confidence": 0.0 + # }} + # ] + # }} + # + # Coordinate system: + # - Treat the page image as a {grid_size}x{grid_size} grid. + # - Origin is the top-left corner. + # - bbox values must be integers in [0, {grid_size}]. + # - bbox must tightly include the whole asset: title, caption, legend, axes, + # labels, table headers, and footnotes that are part of the asset. + # - Exclude surrounding body paragraphs, page headers, page footers, and page + # numbers. + # + # Rules: + # - "table": rows/columns of data, forms, financial tables, appendix tables. + # - "chart": plotted data such as bar/line/pie/scatter charts. + # - "figure": distinct diagrams, flowcharts, architecture drawings, embedded images. + # + # - Do not transcribe entire tables. Summarize the topic and extreme values based on main columns or rows. + # - "keywords" must be an array of up to 5 strings in the same language as the visible asset text. + # - Use confidence 0.0-1.0. Only include assets you can localize. + # - If there are no assets, return {{"regions":[]}}. + # - Return ONLY the JSON object, no markdown fences or explanations. + # """ prompt = f"""\ -You are a precise document layout detector. The attached image is a single PDF -page screenshot. - -Find visually distinct tables, charts, and figures that should become reusable -document assets. Return strict JSON: -{{ - "regions": [ - {{ - "kind": "table|chart|figure", - "bbox": [x1, y1, x2, y2], - "caption": "", - "title": "", - "summary": "<1-3 sentence searchable summary>", - "keywords": ["", ""], - "confidence": 0.0 - }} - ] -}} - -Coordinate system: -- Treat the page image as a {grid_size}x{grid_size} grid. -- Origin is the top-left corner. -- bbox values must be integers in [0, {grid_size}]. -- bbox must tightly include the whole asset: title, caption, legend, axes, - labels, table headers, and footnotes that are part of the asset. -- Exclude surrounding body paragraphs, page headers, page footers, and page - numbers. + You are a precise document layout detector. The attached image is a single + rendered PDF page. -Rules: -- "table": rows/columns of data, forms, financial tables, appendix tables. -- "chart": plotted data such as bar/line/pie/scatter charts. -- "figure": diagrams, flowcharts, architecture drawings, embedded images, or - visual schematics that are not data charts. -- Do not transcribe entire tables. Summarize their topic and visible columns or - row groups. -- "keywords" must be an array of up to 8 strings in the same language as the - visible asset text. -- Use confidence 0.0-1.0. Only include assets you can localize. -- If there are no assets, return {{"regions":[]}}. -- Return ONLY the JSON object, no markdown fences or explanations. -""" + Find visually distinct tables and figures that should become reusable + document assets. Locate them only - do NOT summarize, transcribe full + content, extract keywords, or read data values. Return strict JSON: + {{ + "regions": [ + {{ + "kind": "table|figure", + "bbox": [x1, y1, x2, y2], + "title": "", + "confidence": 0.0 + }} + ] + }} + + Coordinate system: + - Treat the page image as a {grid_size}x{grid_size} grid. + - Origin is the top-left corner. + - bbox values must be integers in [0, {grid_size}]. + - bbox must tightly include the whole asset: its title, caption, legend, + axes, labels, table headers, and footnotes that belong to that asset. + - Exclude surrounding body paragraphs, page headers, page footers, and page + numbers. + + Rules: + - "table": data arranged in clear rows and columns - grid lines, cell + borders, or strongly aligned cells (data tables, forms, financial tables, + appendix tables). + - "figure": any non-table visual asset - bar/line/pie/scatter charts, + plots, diagrams, flowcharts, architecture drawings, schematics, or embedded + images. + - Do not mark ordinary paragraphs, bullet lists, title blocks, or loose + multi-line text as tables. + - Do not split a single coherent table or figure into sub-parts. + - "title" is one short label only (single line). Do not duplicate it into + other fields and do not write a summary. Use an empty string when there is + no visible title or caption. + - Use confidence 0.0-1.0. Only include assets you can localize. + - If there are no qualifying assets, return {{"regions":[]}}. + - Return ONLY the JSON object, no markdown fences or explanations. + """ # ==================== Image Processing Prompts ==================== @@ -937,7 +981,6 @@ def build_prompt(task, texts, query, **kwargs): else: from loguru import logger - logger.warning(f"Unknown task: {task}, returning empty prompt") prompt = "" diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/types.py b/packages/shared-python/shared/services/retrieval/agentic/core/types.py index 8d75618e..a3bb1293 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/core/types.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/types.py @@ -320,7 +320,7 @@ class AgenticResult: """ evidence_text: str answer_text: str = '' - referenced_chunks: list[dict[str, str]] = field(default_factory=list) + referenced_chunks: list[dict[str, Any]] = field(default_factory=list) router_used: str = 'agentic_discovery_only' budget_snapshot: dict[str, Any] | None = None stop_reason: str = '' diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py index f2a59c3c..80a94954 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py @@ -9,7 +9,10 @@ from shared.models.database.document import RetrievalHitStat from shared.services.retrieval.agentic.core.budget import BudgetLedger from shared.services.retrieval.agentic.core.types import DocTreeNode -from shared.services.retrieval.hydration.assets import build_retrieval_asset_url_map +from shared.services.retrieval.hydration.assets import ( + AssetUrlValue, + build_retrieval_asset_url_map, +) from shared.services.retrieval.stats.service import compute_importance_score from shared.utils.token_estimate import estimate_tokens @@ -32,7 +35,7 @@ def _collect_chunks_by_type( def collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]: - return _collect_chunks_by_type(node, {"image", "table"}) + return _collect_chunks_by_type(node, {"image", "table", "page"}) def collect_media_chunks_all( @@ -46,7 +49,7 @@ def collect_media_chunks_all( async def build_asset_url_map( media_chunks: list[dict[str, Any]], -) -> dict[str, str]: +) -> dict[str, AssetUrlValue]: return await build_retrieval_asset_url_map( media_chunks, log_context="agentic evidence", diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py index 6bb6188b..67ac6f6b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py @@ -1,20 +1,18 @@ """Render agentic document trees into evidence text.""" from __future__ import annotations -import os from typing import Any, cast from shared.services.retrieval.agentic.core.types import DocTreeNode -INLINE_TABLE_CHAR_LIMIT_ENV = "RETRIEVAL_AGENTIC_INLINE_TABLE_CHAR_LIMIT" -DEFAULT_INLINE_TABLE_CHAR_LIMIT = 10000 +AssetLookupValue = str | list[str] def render_unified_doc_tree( node: DocTreeNode, doc_name: str, depth: int = 0, - asset_lookup: dict[str, str] | None = None, + asset_lookup: dict[str, AssetLookupValue] | None = None, ) -> str: """Render a DocTreeNode as one coherent hierarchy.""" parts: list[str] = [] @@ -120,7 +118,7 @@ def render_leaf_chunks( parts: list[str], chunks: list[dict[str, Any]], indent: str, - asset_lookup: dict[str, str] | None = None, + asset_lookup: dict[str, AssetLookupValue] | None = None, ) -> None: chunk_by_id = { chunk.get("chunk_id", ""): chunk @@ -141,6 +139,10 @@ def render_leaf_chunks( if chunk_id: rendered_ids.add(chunk_id) + if chunk_type == "page": + render_page_chunk_lines(parts, chunk, indent, asset_lookup=asset_lookup) + continue + content = str(chunk.get("content", "")).strip() for connection in (chunk.get("chunk_metadata") or {}).get("connect_to") or []: target = chunk_by_id.get(connection.get("target", "")) @@ -158,7 +160,7 @@ def render_leaf_chunks( if target_type == "table": table_html = str(target.get("content", "")).strip() file_path = target.get("file_path") or "" - asset_url = (asset_lookup or {}).get(target_id, "") if target_id else "" + asset_url = _first_asset_url((asset_lookup or {}).get(target_id, "")) if target_id else "" display_ref = asset_url or file_path table_lines = render_table_chunk_lines( target, @@ -171,7 +173,7 @@ def render_leaf_chunks( image_description = str(target.get("content", "")).strip() if ref_str in image_description: image_description = image_description.replace(ref_str, "").strip() - asset_url = (asset_lookup or {}).get(target_id, "") if target_id else "" + asset_url = _first_asset_url((asset_lookup or {}).get(target_id, "")) if target_id else "" display_ref = asset_url or file_path if display_ref: content = content.replace(ref_str, f"\n[Image: {display_ref}]\n{image_description}\n") @@ -193,7 +195,7 @@ def render_leaf_chunks( if chunk_type == "image": file_path = chunk.get("file_path") or "" image_description = str(chunk.get("content", "")).strip() - asset_url = (asset_lookup or {}).get(chunk_id, "") if chunk_id else "" + asset_url = _first_asset_url((asset_lookup or {}).get(chunk_id, "")) if chunk_id else "" display_ref = asset_url or file_path if display_ref: parts.append(f"{indent}┈ [Image: {display_ref}]") @@ -204,7 +206,7 @@ def render_leaf_chunks( elif chunk_type == "table": table_html = str(chunk.get("content", "")).strip() file_path = chunk.get("file_path") or "" - asset_url = (asset_lookup or {}).get(chunk_id, "") if chunk_id else "" + asset_url = _first_asset_url((asset_lookup or {}).get(chunk_id, "")) if chunk_id else "" display_ref = asset_url or file_path for line in render_table_chunk_lines( chunk, @@ -215,16 +217,44 @@ def render_leaf_chunks( parts.append(f"{indent}┈ {line}") +def render_page_chunk_lines( + parts: list[str], + chunk: dict[str, Any], + indent: str, + asset_lookup: dict[str, AssetLookupValue] | None = None, +) -> None: + metadata = chunk.get("chunk_metadata") or chunk.get("metadata") or {} + summary = str(metadata.get("summary") or chunk.get("summary") or "").strip() + page_nums = _coerce_page_nums( + metadata.get("page_nums") or chunk.get("page_nums") + ) + if page_nums: + page_label = ", ".join(str(page) for page in page_nums) + parts.append(f"{indent}┈ Pages: {page_label}") + if summary: + for line in summary.split("\n"): + if line.strip(): + parts.append(f"{indent}┈ {line}") + elif not page_nums: + parts.append(f"{indent}┈ [Page]") + + urls = _asset_urls_for_chunk(chunk, asset_lookup=asset_lookup) + if urls: + for index, url in enumerate(urls, start=1): + parts.append(f"{indent}┈ [Page image {index}: {url}]") + elif page_nums: + for page in page_nums: + parts.append(f"{indent}┈ [Page image: page {page}]") + + def render_table_chunk_lines( chunk: dict[str, Any], *, table_html: str, display_ref: str, ) -> list[str]: + del table_html header = f"[Table: {display_ref}]" if display_ref else "[Table]" - if len(table_html) <= _inline_table_char_limit(): - return [header, *[line for line in table_html.split("\n") if line.strip()]] - lines = [header] table_path = chunk.get("source_chunk_path") or chunk.get("section_path") if table_path: @@ -232,10 +262,6 @@ def render_table_chunk_lines( file_path = chunk.get("file_path") if file_path: lines.append(f"Table asset: {file_path}") - lines.append( - f"Large table omitted from evidence_text: {len(table_html)} chars " - f"(inline limit {_inline_table_char_limit()} chars)." - ) metadata = chunk.get("chunk_metadata") or chunk.get("metadata") or {} summary = metadata.get("summary") if isinstance(metadata, dict) else "" @@ -249,11 +275,55 @@ def render_table_chunk_lines( if keyword_text: lines.append("Main columns:") lines.append(keyword_text) + elif isinstance(keywords, str) and keywords.strip(): + lines.append("Main columns:") + lines.append(keywords.strip()) + + caption = metadata.get("caption") if isinstance(metadata, dict) else "" + if caption: + lines.append("Caption:") + lines.append(str(caption).strip()) return lines -def _inline_table_char_limit() -> int: - return int(os.getenv(INLINE_TABLE_CHAR_LIMIT_ENV, str(DEFAULT_INLINE_TABLE_CHAR_LIMIT))) +def _asset_urls_for_chunk( + chunk: dict[str, Any], + *, + asset_lookup: dict[str, AssetLookupValue] | None, +) -> list[str]: + chunk_id = str(chunk.get("chunk_id") or "").strip() + value = (asset_lookup or {}).get(chunk_id, "") if chunk_id else "" + if isinstance(value, list): + return [str(url).strip() for url in value if str(url).strip()] + url = str(value or "").strip() + return [url] if url else [] + + +def _first_asset_url(value: object) -> str: + if isinstance(value, list): + for item in value: + url = str(item or "").strip() + if url: + return url + return "" + return str(value or "").strip() + + +def _coerce_page_nums(value: object) -> list[int]: + if isinstance(value, list): + raw_values = value + elif value is None: + raw_values = [] + else: + raw_values = str(value).split(",") + + pages: list[int] = [] + for item in raw_values: + try: + pages.append(int(str(item).strip())) + except (TypeError, ValueError): + continue + return pages def _infer_child_sort_order(child: DocTreeNode) -> float: diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py index ad6367c9..038f7000 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py @@ -631,7 +631,7 @@ async def _search_images_via_vlm( for index, asset in enumerate(assets, start=1): chunk_id = str(asset.get("chunk_id") or "") url = url_map.get(chunk_id) - if not url: + if not url or isinstance(url, list): continue row_id = f"I{index}" file_path = asset.get("file_path") or "" diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index 4b30012e..4fb8835a 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -361,7 +361,7 @@ async def run( ) # Collect referenced chunk IDs from all doc trees - all_refs: list[dict[str, str]] = [] + all_refs: list[dict[str, Any]] = [] seen_ref_ids: set[str] = set() for doc_id, doc_tree in state.doc_trees.items(): doc_name = state.doc_id_to_name.get(doc_id, doc_id) diff --git a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index fa42467f..534c016f 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -26,15 +26,19 @@ async def resolve_workflow_references( refs: list[dict[str, Any]], score_by_chunk_id: dict[str, float] | None = None, ) -> ResolvedWorkflowReferences: - enriched_refs = await enrich_referenced_chunks_with_asset_urls(refs) hydrated_rows = await hydrate_referenced_chunk_rows( db=db, user_id=user_id, namespace=namespace, - refs=enriched_refs, + refs=refs, score_by_chunk_id=score_by_chunk_id, ) - return _select_matching_references(enriched_refs, hydrated_rows) + resolved = _select_matching_references(refs, hydrated_rows) + enriched_rows = await enrich_referenced_chunks_with_asset_urls(resolved.rows) + return ResolvedWorkflowReferences( + refs=_merge_reference_asset_urls(resolved.refs, enriched_rows), + rows=resolved.rows, + ) def _select_matching_references( @@ -91,6 +95,39 @@ def _row_key(row: dict[str, Any]) -> tuple[str, str, str, str]: ) +def _merge_reference_asset_urls( + refs: list[dict[str, Any]], + rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + row_by_key = {_row_key(row): row for row in rows} + merged_refs: list[dict[str, Any]] = [] + for ref in refs: + merged = dict(ref) + row = row_by_key.get( + build_reference_lookup_key( + document_id=ref.get("document_id"), + chunk_id=ref.get("chunk_id"), + section_path=ref.get("section_path"), + file_path=ref.get("file_path"), + ) + ) + if row is None: + row = _find_matching_row_for_ref(ref, rows) + if row is not None: + for field in ("asset_url", "asset_urls"): + if row.get(field): + merged[field] = row[field] + merged_refs.append(merged) + return merged_refs + + +def _find_matching_row_for_ref( + ref: dict[str, Any], + rows: list[dict[str, Any]], +) -> dict[str, Any] | None: + return next((row for row in rows if _matches_reference(ref, row)), None) + + def _matches_root_alias(ref: dict[str, Any], row: dict[str, Any]) -> bool: ref_section_path = str(ref.get("section_path") or "").strip() row_section_path = str(row.get("section_path") or "").strip() diff --git a/packages/shared-python/shared/services/retrieval/hydration/assets.py b/packages/shared-python/shared/services/retrieval/hydration/assets.py index 99d1bb4a..352aa2dc 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/assets.py +++ b/packages/shared-python/shared/services/retrieval/hydration/assets.py @@ -7,6 +7,8 @@ from shared.services.retrieval.hydration.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type from shared.services.storage.result_storage import get_result_storage +AssetUrlValue = str | list[str] + def _normalize_artifact_ref(asset_ref: object) -> str | None: return get_result_storage().normalize_artifact_ref( @@ -19,6 +21,11 @@ def _is_retrieval_media_row(row: dict[str, Any]) -> bool: return normalize_chunk_type(raw_chunk_type) in MEDIA_CHUNK_TYPES +def _is_page_row(row: dict[str, Any]) -> bool: + raw_chunk_type = row.get("chunk_type") or row.get("type") + return normalize_chunk_type(raw_chunk_type) == "page" + + def _resolve_asset_request(row: dict[str, Any]) -> tuple[str, str] | None: job_id = str(row.get("job_id") or "").strip() if not job_id or not _is_retrieval_media_row(row): @@ -31,6 +38,33 @@ def _resolve_asset_request(row: dict[str, Any]) -> tuple[str, str] | None: return job_id, artifact_ref +def _resolve_page_asset_requests(row: dict[str, Any]) -> list[tuple[str, str]]: + job_id = str(row.get("job_id") or "").strip() + if not job_id or not _is_page_row(row): + return [] + + metadata = row.get("chunk_metadata") or row.get("metadata") or {} + if not isinstance(metadata, dict): + return [] + + requests: list[tuple[str, str]] = [] + seen_refs: set[str] = set() + raw_refs = metadata.get("page_image_uris") or [] + if not isinstance(raw_refs, list): + return [] + for raw_ref in raw_refs: + artifact_ref = _normalize_artifact_ref(raw_ref) + if ( + artifact_ref is None + or not artifact_ref.startswith("pages/") + or artifact_ref in seen_refs + ): + continue + seen_refs.add(artifact_ref) + requests.append((job_id, artifact_ref)) + return requests + + async def _generate_retrieval_asset_url( *, row: dict[str, Any], @@ -51,6 +85,30 @@ async def _generate_retrieval_asset_url( return None +async def _generate_retrieval_asset_urls( + *, + row: dict[str, Any], + log_context: str, +) -> list[str]: + requests = _resolve_page_asset_requests(row) + if not requests: + return [] + + urls: list[str] = [] + for job_id, artifact_ref in requests: + try: + url = get_result_storage().generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + ) + except Exception as exc: + logger.warning(f"Failed to generate {log_context} page asset URL (ignored): {exc}") + continue + if url: + urls.append(url) + return urls + + async def enrich_rows_with_retrieval_asset_urls( rows: list[dict[str, Any]], *, @@ -65,6 +123,12 @@ async def enrich_rows_with_retrieval_asset_urls( ) if asset_url: enriched["asset_url"] = asset_url + asset_urls = await _generate_retrieval_asset_urls( + row=row, + log_context=log_context, + ) + if asset_urls: + enriched["asset_urls"] = asset_urls enriched_rows.append(enriched) return enriched_rows @@ -73,13 +137,21 @@ async def build_retrieval_asset_url_map( rows: list[dict[str, Any]], *, log_context: str, -) -> dict[str, str]: - url_map: dict[str, str] = {} +) -> dict[str, AssetUrlValue]: + url_map: dict[str, AssetUrlValue] = {} for row in rows: chunk_id = str(row.get("chunk_id") or "").strip() if not chunk_id: continue + asset_urls = await _generate_retrieval_asset_urls( + row=row, + log_context=log_context, + ) + if asset_urls: + url_map[chunk_id] = asset_urls + continue + asset_url = await _generate_retrieval_asset_url( row=row, log_context=log_context, diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py index b2e3b649..90f74cf3 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -55,7 +55,12 @@ async def assemble_retrieval_results( continue assembled_row = dict(row) base_content = str(row.get('content') or '') - if normalize_chunk_type(row.get('chunk_type')) == 'text': + chunk_type = normalize_chunk_type(row.get('chunk_type')) + if chunk_type == 'page': + assembled_row['content'] = _page_summary(row) + elif chunk_type == 'table': + assembled_row['content'] = _table_summary_content(row) + elif chunk_type == 'text': connected_targets: list[tuple[int, str]] = [] for target_id in iter_connected_target_ids(row): target_row = rows_by_chunk_id.get(target_id) @@ -63,18 +68,13 @@ async def assemble_retrieval_results( continue if normalize_chunk_type(target_row.get('chunk_type')) != 'table': continue - target_content = str(target_row.get('content') or '').strip() + target_content = _table_summary_content(target_row) if target_content: sort_key = int(target_row.get('sort_order', 0) or 0) connected_targets.append((sort_key, target_content)) connected_targets.sort(key=lambda item: item[0]) related_parts = [content for _, content in connected_targets] - # TODO: Oversized Table Protection for results[].content - # Currently, if a text chunk connects to a large table (e.g., 50K chars), - # the full table HTML is appended here without any truncation. - # We should consider whether to truncate this and only leave the asset URL. - # # TODO: Dedicated Large Table Agent # For the Notebook/Agent environment, consider introducing a dedicated # "Large Table Agent" that can fetch and query oversized tables via URL. @@ -87,3 +87,51 @@ async def assemble_retrieval_results( assembled_row['content'] = clean_content(assembled_row['content']) assembled.append(assembled_row) return assembled + + +def _page_summary(row: dict[str, Any]) -> str: + metadata = row.get('chunk_metadata') or row.get('metadata') or {} + if not isinstance(metadata, dict): + return '' + return str(metadata.get('summary') or '').strip() + + +def _table_summary_content(row: dict[str, Any]) -> str: + metadata = row.get('chunk_metadata') or row.get('metadata') or {} + if not isinstance(metadata, dict): + metadata = {} + + display_ref = _table_display_ref(row) + lines = [f"[Table: {display_ref}]" if display_ref else "[Table]"] + + summary = str(metadata.get('summary') or row.get('summary') or '').strip() + if summary: + lines.extend(line for line in summary.split('\n') if line.strip()) + + keywords = metadata.get('keywords') or row.get('keywords') or [] + if isinstance(keywords, list): + keyword_text = ';'.join( + str(keyword).strip() for keyword in keywords if str(keyword).strip() + ) + else: + keyword_text = str(keywords or '').strip() + if keyword_text: + lines.append(keyword_text) + + caption = str(metadata.get('caption') or row.get('caption') or '').strip() + if caption: + lines.append(caption) + + return '\n'.join(lines) + + +def _table_display_ref(row: dict[str, Any]) -> str: + for key in ('asset_url', 'file_path', 'source_chunk_path'): + value = str(row.get(key) or '').strip() + if value: + return value + + content = str(row.get('content') or '').strip() + if content and not content.lstrip().lower().startswith(' str: diff --git a/packages/shared-python/shared/services/retrieval/search/lexical_text.py b/packages/shared-python/shared/services/retrieval/search/lexical_text.py index b6cbc9c8..25fba2f3 100644 --- a/packages/shared-python/shared/services/retrieval/search/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/search/lexical_text.py @@ -43,6 +43,12 @@ def build_lexical_text(value: str) -> str: def build_content_lexical_text(chunk: dict[str, Any]) -> Optional[str]: + if _is_table_chunk(chunk): + content = _table_search_source_text(chunk) + if not content: + return None + return build_lexical_text(content) + content = str(chunk.get("content") or chunk.get("text") or "").strip() if not content: return None @@ -100,7 +106,10 @@ def build_content_search_text( section_summary: Optional[str] = None, ) -> Optional[str]: """Pre-tokenized content field for retrieval BM25 scoring.""" - content = str(chunk.get("content") or chunk.get("text") or "").strip() + if _is_table_chunk(chunk): + content = _table_search_source_text(chunk) + else: + content = str(chunk.get("content") or chunk.get("text") or "").strip() if not content: return None parts = [content] @@ -137,7 +146,38 @@ def build_term_search_text( path_text: Optional[str] = None, ) -> Optional[str]: """Raw combined field for grep channel: content + path (not tokenized).""" - content = str(chunk.get("content") or chunk.get("text") or "").strip() + if _is_table_chunk(chunk): + content = _table_search_source_text(chunk) + else: + content = str(chunk.get("content") or chunk.get("text") or "").strip() path = str(path_text or "").strip() combined = f"{content} {path}".strip() return combined if combined else None + + +def _is_table_chunk(chunk: dict[str, Any]) -> bool: + raw_type = chunk.get("type") or chunk.get("chunk_type") or "" + return str(raw_type).strip().split("\n", 1)[0].lower() == "table" + + +def _table_search_source_text(chunk: dict[str, Any]) -> str: + metadata = chunk.get("metadata") or chunk.get("chunk_metadata") or {} + if not isinstance(metadata, dict): + metadata = {} + + keyword_value = metadata.get("keywords") or chunk.get("keywords") or [] + if isinstance(keyword_value, list): + keyword_text = " ".join( + str(keyword).strip() + for keyword in keyword_value + if str(keyword).strip() + ) + else: + keyword_text = str(keyword_value or "").replace(";", " ").strip() + + parts = [ + str(metadata.get("summary") or chunk.get("summary") or "").strip(), + keyword_text, + str(metadata.get("caption") or chunk.get("caption") or "").strip(), + ] + return " ".join(part for part in parts if part) diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py index dfc55bf0..8eed87e1 100644 --- a/packages/shared-python/shared/services/retrieval/settings.py +++ b/packages/shared-python/shared/services/retrieval/settings.py @@ -14,6 +14,8 @@ 4: {'table'}, 5: {'text', 'image'}, 6: {'text', 'table'}, + 7: {'page'}, + 8: {'text', 'image', 'table'}, } diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index e16b8dab..614e1dd8 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -13,7 +13,7 @@ _EXCLUDED_FILE_NAMES = {".DS_Store", "Thumbs.db"} _EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"} -_CLIENT_ARTIFACT_DIRS = {"images", "tables"} +_CLIENT_ARTIFACT_DIRS = {"images", "tables", "pages"} @dataclass(frozen=True) @@ -25,7 +25,12 @@ class UploadedResultBundle: class ResultStorage(Protocol): def upload( - self, *, job_id: str, result_dir: str, zip_file_path: str + self, + *, + job_id: str, + result_dir: str, + zip_file_path: str, + artifact_refs: set[str] | None = None, ) -> UploadedResultBundle: raise NotImplementedError @@ -73,7 +78,12 @@ def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: return normalized def upload( - self, *, job_id: str, result_dir: str, zip_file_path: str + self, + *, + job_id: str, + result_dir: str, + zip_file_path: str, + artifact_refs: set[str] | None = None, ) -> UploadedResultBundle: result_path = Path(result_dir) if not result_path.is_dir(): @@ -91,8 +101,11 @@ def upload( self._cleanup_file(zip_path) raw_files: dict[str, str] = {} + artifact_ref_filter = self._normalize_artifact_refs(artifact_refs) for file_path in self._iter_raw_files(result_path): relative_path = file_path.relative_to(result_path).as_posix() + if artifact_ref_filter is not None and relative_path not in artifact_ref_filter: + continue raw_key = self.build_raw_key(job_id=job_id, relative_path=relative_path) self._job_file_storage.upload_local_file( str(file_path), @@ -152,6 +165,16 @@ def _normalize_raw_relative_path(self, relative_path: str | None) -> str | None: return None return "/".join(parts) + def _normalize_artifact_refs(self, artifact_refs: set[str] | None) -> set[str] | None: + if artifact_refs is None: + return None + normalized_refs = { + normalized + for ref in artifact_refs + if (normalized := self.normalize_artifact_ref(ref)) is not None + } + return normalized_refs + def _is_excluded_file(self, file_name: str) -> bool: return file_name in _EXCLUDED_FILE_NAMES or file_name.startswith(".") diff --git a/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py index 962c7f65..37343884 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -130,17 +130,7 @@ def _normalize_chunk_type(value: Any) -> str: return raw_type.split("\n", 1)[0].lower() -_PAGE_CHUNK_METADATA_KEYS = ( - "granularity", - "page_indices", - "owned_pages", - "section_path", - "section_level", - "page_image_uris", - "kind", - "status", - "source_verdict", -) +_PAGE_CHUNK_METADATA_KEYS = ("page_image_uris",) def _merge_page_chunk_metadata( diff --git a/packages/shared-python/shared/services/storage/zip_doc_navigation.py b/packages/shared-python/shared/services/storage/zip_doc_navigation.py index 9ad5870a..ff5f52eb 100644 --- a/packages/shared-python/shared/services/storage/zip_doc_navigation.py +++ b/packages/shared-python/shared/services/storage/zip_doc_navigation.py @@ -86,11 +86,15 @@ def build_doc_nav( ) elif chunk_type == "page": stats["page_chunks"] += 1 - # Page chunks participate in section tree like text chunks + page_nums = metadata.get("page_nums") or [] + page_count = len(page_nums) if isinstance(page_nums, list) else 0 + # Page chunks participate in section tree like text chunks, while + # preserving page-count semantics for navigation counts. text_chunks.append( { "path": path, "summary": summary or content_preview, + "chunk_count": page_count or 1, } ) else: @@ -99,6 +103,7 @@ def build_doc_nav( { "path": path, "summary": summary or content_preview, + "chunk_count": 1, } ) @@ -144,7 +149,7 @@ def _build_section_tree( "chunk_count": 0, "_children_map": {}, } - root_children[key]["chunk_count"] += 1 + root_children[key]["chunk_count"] += int(chunk.get("chunk_count") or 1) if not root_children[key]["summary"]: root_children[key]["summary"] = chunk.get("summary", "") continue @@ -163,7 +168,7 @@ def _build_section_tree( } node = current_level[part] if index == len(section_parts) - 1: - node["chunk_count"] += 1 + node["chunk_count"] += int(chunk.get("chunk_count") or 1) if not node["summary"]: node["summary"] = chunk.get("summary", "") current_level = node["_children_map"] @@ -200,186 +205,3 @@ def _section_tree_to_output( ) return result - -def _coerce_int(value: Any) -> int | None: - try: - return int(value) - except (TypeError, ValueError): - return None - - -def build_doc_nav_from_skeletons( - skeletons: list[dict[str, Any]], - chunks: list[dict[str, Any]], - source_file_name: str, -) -> dict[str, Any]: - """Build doc_nav.json from SectionSkeleton dicts + page chunks. - - Unlike ``ZipDocNavigationBuilder.build_doc_nav`` which infers the tree - from chunk paths (losing sections whose start page is shared with a - sibling), this builder uses skeletons as the authoritative tree structure - and computes chunk_count by page-range overlap. - - Parameters - ---------- - skeletons: - List of ``SectionSkeleton.to_dict()`` dicts. Must have - ``section_path``, ``title``, ``level``, ``start_page``, ``end_page``, - and ``parent_path``. - chunks: - Formatted chunk dicts (from ``dataframe_to_chunks``). Only chunks - with ``type == "page"`` are counted; their ``metadata.page_nums[0]`` - identifies the page index. - source_file_name: - Original filename for the doc_nav envelope. - """ - page_indices: set[int] = set() - chunk_summaries: dict[int, str] = {} - for chunk in chunks: - meta = chunk.get("metadata") or {} - page_nums = meta.get("page_nums") or [] - if page_nums: - page = _coerce_int(page_nums[0]) - if page is None: - continue - page_indices.add(page) - if not chunk_summaries.get(page): - chunk_summaries[page] = ( - (meta.get("summary") or "").strip() - or (chunk.get("content") or "").strip()[:200] - ) - - sorted_skels = sorted( - skeletons, - key=lambda s: ( - _coerce_int(s.get("start_page")) or 0, - _coerce_int(s.get("level")) or 0, - str(s.get("section_path") or ""), - ), - ) - nodes: dict[str, dict[str, Any]] = {} - root_children_order: list[str] = [] - - def _link_child(parent_path: str, child_path: str) -> None: - if not parent_path: - if child_path not in root_children_order: - root_children_order.append(child_path) - return - parent = nodes.get(parent_path) - if parent is None: - return - children = parent["_children_paths"] - if child_path not in children: - children.append(child_path) - - def _ensure_path_node(path: str) -> str: - root_parts, section_parts = split_document_path( - path, - source_file_name=source_file_name, - ) - if not section_parts: - return "" - - parent_path = "" - path_parts = list(root_parts) - for index, title in enumerate(section_parts): - path_parts.append(title) - current_path = "/".join(path_parts) - if current_path not in nodes: - nodes[current_path] = { - "title": title, - "path": current_path, - "level": index + 1, - "start_page": None, - "end_page": None, - "summary": "", - "parent_path": parent_path, - "owned_pages": set(), - "_children_paths": [], - } - _link_child(parent_path, current_path) - parent_path = current_path - return parent_path - - for skel in sorted_skels: - sp = str(skel.get("section_path") or "").strip() - if not sp: - continue - current_path = _ensure_path_node(sp) - if not current_path: - continue - - node = nodes[current_path] - start_page = _coerce_int(skel.get("start_page")) - end_page = _coerce_int(skel.get("end_page")) - if start_page is not None and end_page is not None: - node["start_page"] = start_page - node["end_page"] = end_page - node["owned_pages"].update(page_indices & set(range(start_page, end_page + 1))) - node["summary"] = chunk_summaries.get(start_page, node["summary"]) - node["title"] = str(skel.get("title") or node["title"]) - node["level"] = _coerce_int(skel.get("level")) or node["level"] - - def _collect_owned_pages(path: str) -> set[int]: - node = nodes[path] - pages = set(node["owned_pages"]) - for child_path in node["_children_paths"]: - pages.update(_collect_owned_pages(child_path)) - node["owned_pages"] = pages - return pages - - for path in root_children_order: - if path in nodes: - _collect_owned_pages(path) - - def _to_output(node: dict[str, Any], level: int = 1) -> dict[str, Any]: - children_out = [ - _to_output(nodes[child_path], level + 1) - for child_path in node["_children_paths"] - if child_path in nodes - ] - return { - "title": node["title"], - "path": node["path"], - "level": level, - "summary": node["summary"], - "chunk_count": len(node["owned_pages"]), - "children": children_out, - } - - sections = [_to_output(nodes[sp]) for sp in root_children_order if sp in nodes] - - # Stats - stats = { - "total_chunks": len(chunks), - "text_chunks": sum(1 for c in chunks if c.get("type") == "text"), - "image_chunks": sum(1 for c in chunks if c.get("type") == "image"), - "table_chunks": sum(1 for c in chunks if c.get("type") == "table"), - "page_chunks": sum(1 for c in chunks if c.get("type") == "page"), - "max_depth": _max_depth(sections), - } - - # Resources - image_resources = [] - table_resources = [] - for chunk in chunks: - ct = chunk.get("type", "") - path = chunk.get("path", "") - meta = chunk.get("metadata") or {} - summary = (meta.get("summary") or "").strip() - if ct == "image": - image_resources.append({"path": path, "summary": summary}) - elif ct == "table": - table_resources.append({"path": path, "summary": summary}) - - return { - "version": "1.0", - "file_name": source_file_name or "", - "stats": stats, - "sections": sections, - "resources": { - "images": image_resources, - "tables": table_resources, - }, - } - diff --git a/packages/shared-python/shared/services/storage/zip_result_resources.py b/packages/shared-python/shared/services/storage/zip_result_resources.py index dd96c4e6..0a2a8847 100644 --- a/packages/shared-python/shared/services/storage/zip_result_resources.py +++ b/packages/shared-python/shared/services/storage/zip_result_resources.py @@ -18,7 +18,6 @@ class ZipPackageResources: image_files: tuple[ZipResourceFileInfo, ...] table_files: tuple[ZipResourceFileInfo, ...] - page_files: tuple[ZipResourceFileInfo, ...] = () @property def image_files_map(self) -> dict[str, ZipResourceFileInfo]: @@ -40,11 +39,9 @@ def collect( ) -> ZipPackageResources: images_dir = os.path.join(add_dir, "images") tables_dir = os.path.join(add_dir, "tables") - pages_dir = os.path.join(add_dir, "pages") return ZipPackageResources( image_files=tuple(self._collect_image_files(chunks, images_dir)), table_files=tuple(self._collect_table_files(chunks, tables_dir)), - page_files=tuple(self._collect_page_files(pages_dir)), ) def _collect_image_files( @@ -216,27 +213,6 @@ def _collect_table_files( return table_files - @staticmethod - def _collect_page_files(pages_dir: str) -> list[ZipResourceFileInfo]: - """Collect all page image/thumbnail files from pages/ directory.""" - if not os.path.exists(pages_dir): - return [] - page_files: list[ZipResourceFileInfo] = [] - for filename in sorted(os.listdir(pages_dir)): - file_path = os.path.join(pages_dir, filename) - if not os.path.isfile(file_path): - continue - page_files.append( - { - "file_path": f"pages/{filename}", - "original_name": filename, - "size_bytes": os.path.getsize(file_path), - "source_path": file_path, - "zip_path": f"pages/{filename}", - } - ) - return page_files - def _collect_files_by_name(directory_path: str) -> dict[str, str]: files: dict[str, str] = {} for filename in os.listdir(directory_path): diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py index 667cab70..e444475a 100644 --- a/packages/shared-python/shared/services/storage/zip_result_schema.py +++ b/packages/shared-python/shared/services/storage/zip_result_schema.py @@ -5,10 +5,7 @@ from typing import Any from shared.services.storage.zip_chunk_schema import ZipChunkSchemaBuilder -from shared.services.storage.zip_doc_navigation import ( - ZipDocNavigationBuilder, - build_doc_nav_from_skeletons, -) +from shared.services.storage.zip_doc_navigation import ZipDocNavigationBuilder from shared.services.storage.zip_manifest_schema import ZipManifestBuilder @@ -73,15 +70,3 @@ def build_doc_nav( formatted_chunks, source_file_name, ) - - def build_doc_nav_from_skeletons( - self, - skeletons: list[dict[str, Any]], - formatted_chunks: list[dict[str, Any]], - source_file_name: str, - ) -> dict[str, Any]: - return build_doc_nav_from_skeletons( - skeletons, - formatted_chunks, - source_file_name, - ) diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index 20e83e25..b2fee6e1 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -71,7 +71,6 @@ def generate_zip_package( doc_nav, hierarchy = self._build_navigation_outputs( formatted_chunks=formatted_chunks, source_file_name=source_file_name, - parsed_df=parsed_df, ) manifest = self._schema.generate_manifest( job_id=job_id, @@ -122,33 +121,11 @@ def _build_navigation_outputs( *, formatted_chunks: list[dict[str, Any]], source_file_name: str, - parsed_df: pd.DataFrame | None = None, ) -> tuple[dict[str, Any] | None, dict[str, Any]]: try: - skeletons = _get_page_memory_skeletons(parsed_df) - if skeletons: - doc_nav = self._schema.build_doc_nav_from_skeletons( - skeletons, - formatted_chunks, - source_file_name, - ) - else: - doc_nav = self._schema.build_doc_nav(formatted_chunks, source_file_name) + doc_nav = self._schema.build_doc_nav(formatted_chunks, source_file_name) hierarchy = self._schema.build_hierarchy_dict(doc_nav.get("sections", [])) return doc_nav, hierarchy except Exception as exc: logger.warning(f"generate doc_nav.json fail {exc}") return None, {} - - -def _get_page_memory_skeletons( - parsed_df: pd.DataFrame | None, -) -> list[dict[str, Any]]: - if parsed_df is None: - return [] - - raw_skeletons = parsed_df.attrs.get("page_memory_skeletons") - if not isinstance(raw_skeletons, list): - return [] - - return [skel for skel in raw_skeletons if isinstance(skel, dict)] From c8f88ce785282c60416eccd1fe1536b57c0965fb Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 25 Jun 2026 08:23:12 +0800 Subject: [PATCH 05/27] refactor: decouple page memory from image URIs and simplify node metadata structure --- .../document_agent/planner/planner.py | 1 + .../structure/page_locate_agent.py | 7 +- .../app/services/document_agent/trace.py | 26 + .../app/services/document_agent/visual.py | 39 +- .../document_ingestion/artifact_refs.py | 9 +- .../success_finalization.py | 10 +- .../document_parser/assets/inline_asset.py | 2 - .../document_parser/formats/docx/parser.py | 1 - .../formats/excel/table_parser.py | 2 - .../formats/markdown/table_asset.py | 1 - .../document_parser/profiling/doc_profiler.py | 22 +- .../tables/table_asset_writer.py | 2 - .../services/page_memory/fine_hierarchy.py | 70 +- .../services/page_memory/memory_service.py | 1108 ++++++++++++++--- .../services/page_memory/node_assembler.py | 47 +- .../app/services/page_memory/page_assets.py | 16 +- .../app/services/page_memory/page_renderer.py | 19 +- ...test_agentic_evidence_renderer_contract.py | 22 +- .../test_document_agent_budget_contract.py | 9 +- .../test_page_memory_asset_java_contract.py | 2 +- ...est_page_memory_node_assembler_contract.py | 57 +- .../test_page_memory_retrieval_contract.py | 207 ++- .../test_table_asset_schema_contract.py | 49 + packages/shared-python/pyproject.toml | 1 + .../chunks/dataframe_chunk_converter.py | 8 - .../retrieval/agentic/evidence/renderer.py | 67 +- .../retrieval/execution/reference_resolver.py | 13 +- .../execution/response_projection.py | 8 +- .../services/retrieval/hydration/assets.py | 137 +- .../services/retrieval/hydration/row_utils.py | 4 +- .../shared/services/storage/page_pdf_crop.py | 106 ++ .../shared/services/storage/result_storage.py | 46 +- .../services/storage/zip_chunk_schema.py | 15 - .../services/storage/zip_package_writer.py | 61 + page_memory_work_summary_20260624.md | 437 +++++++ uv.lock | 2 + 36 files changed, 2188 insertions(+), 445 deletions(-) create mode 100644 apps/worker/tests/contract/test_table_asset_schema_contract.py create mode 100644 packages/shared-python/shared/services/storage/page_pdf_crop.py create mode 100644 page_memory_work_summary_20260624.md diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py index 2efae784..b7df8671 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -317,6 +317,7 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: output_summary={"profile": profile.to_dict(), "decision": decision.to_dict()}, debug={ "prompt_text": prompt_text, + "sampled_pages": pages, "sampled_pngs": pngs, "raw_response": raw, }, diff --git a/apps/worker/app/services/document_agent/structure/page_locate_agent.py b/apps/worker/app/services/document_agent/structure/page_locate_agent.py index 15d06b8d..de881919 100644 --- a/apps/worker/app/services/document_agent/structure/page_locate_agent.py +++ b/apps/worker/app/services/document_agent/structure/page_locate_agent.py @@ -250,16 +250,17 @@ def verify_section_page_choice( model = None if ctx is not None: model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") + pages = [match.page for match in candidates] if ctx is None or not model or ctx.budget is None: best = candidates[0] return { "selected_page": best.page, + "candidate_pages": pages, "confidence": min(best.confidence, GREP_ONLY_CONFIDENCE_CAP), "source": "agent_heuristic", "reason": "VLM unavailable; selected top grep candidate", } - pages = [match.page for match in candidates] from app.services.document_agent.visual import render_pages rendered = render_pages( @@ -273,6 +274,7 @@ def verify_section_page_choice( best = candidates[0] return { "selected_page": best.page, + "candidate_pages": pages, "confidence": min(best.confidence, RENDER_FAILED_GREP_CONFIDENCE_CAP), "source": "agent_heuristic", "reason": "render failed; selected top grep candidate", @@ -285,6 +287,7 @@ def verify_section_page_choice( best = candidates[0] return { "selected_page": best.page, + "candidate_pages": pages, "confidence": min(best.confidence, BUDGET_EXHAUSTED_GREP_CONFIDENCE_CAP), "source": "agent_heuristic", "reason": "page_locate visual budget exhausted; selected top grep candidate", @@ -329,6 +332,7 @@ def verify_section_page_choice( selected_page = None return { "selected_page": selected_page, + "candidate_pages": pages, "confidence": float(payload.get("confidence") or VLM_CONFIRMED_DEFAULT_CONFIDENCE), "source": "agent_vlm", "reason": str(payload.get("reason") or ""), @@ -341,6 +345,7 @@ def verify_section_page_choice( logger.warning("[page_locate.agent] VLM failed for title={!r}: {}", title, exc) return { "selected_page": best.page, + "candidate_pages": pages, "confidence": min(best.confidence, VLM_FAILED_GREP_CONFIDENCE_CAP), "source": "agent_heuristic", "reason": f"VLM failed ({type(exc).__name__}); selected top grep candidate", diff --git a/apps/worker/app/services/document_agent/trace.py b/apps/worker/app/services/document_agent/trace.py index 5771e602..62f28df8 100644 --- a/apps/worker/app/services/document_agent/trace.py +++ b/apps/worker/app/services/document_agent/trace.py @@ -20,6 +20,7 @@ def __init__(self, *, job_id: str, db: Any | None = None) -> None: self._db = db self._started = time.monotonic() self._steps: list[dict[str, Any]] = [] + self._stages: list[dict[str, Any]] = [] self._anatomy: PageAnatomyMap | None = None self._doc_profile_source: Any | None = None self._artifact_path: str | None = None @@ -58,6 +59,22 @@ def record_step( } ) + def record_stage( + self, + stage: str, + *, + page_info: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, + ) -> None: + self._stages.append( + { + "stage": stage, + "page_info": page_info or {}, + "variables": variables or {}, + "created_at": datetime.utcnow(), + } + ) + def set_anatomy_map(self, anatomy: PageAnatomyMap, artifact_path: str) -> None: self._anatomy = anatomy self._artifact_path = artifact_path @@ -162,6 +179,13 @@ def write_trace_json( if created_at is not None and hasattr(created_at, "isoformat"): item["created_at"] = created_at.isoformat() serializable_steps.append(item) + serializable_stages = [] + for stage in self._stages: + item = dict(stage) + created_at = item.get("created_at") + if created_at is not None and hasattr(created_at, "isoformat"): + item["created_at"] = created_at.isoformat() + serializable_stages.append(item) Path(trace_path).write_text( json.dumps( { @@ -171,6 +195,7 @@ def write_trace_json( "summary": summary, "artifact_path": self._artifact_path, "steps": serializable_steps, + "stages": serializable_stages, }, ensure_ascii=False, indent=2, @@ -185,6 +210,7 @@ def summary(self) -> dict[str, Any]: return { "run_id": self.run_id, "step_count": len(self._steps), + "stage_count": len(self._stages), "artifact_path": self._artifact_path, "latency_ms": int((time.monotonic() - self._started) * 1000), } diff --git a/apps/worker/app/services/document_agent/visual.py b/apps/worker/app/services/document_agent/visual.py index 5a6e6da3..e1e8233f 100644 --- a/apps/worker/app/services/document_agent/visual.py +++ b/apps/worker/app/services/document_agent/visual.py @@ -4,6 +4,7 @@ import gc import os +import shutil from pathlib import Path from typing import Any @@ -14,6 +15,43 @@ ) +_DEBUG_VISUAL_DIRS = { + "planner_pages", + "page_locate_pages", + "toc_pages", + "inspect_pages", + "verify_pages", + "agent_visuals", +} +_PAGE_MEMORY_VISUAL_DIRS = {"pages", "asset_annotate"} + + +def visual_debug_enabled() -> bool: + return os.environ.get("DOC_AGENT_KEEP_PAGE_VISUALS", "false").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +def purge_debug_visual_dirs(output_dir: str | None) -> None: + if not output_dir: + return + root = Path(output_dir) + candidates = [root / "_doc_agent" / name for name in _DEBUG_VISUAL_DIRS] + candidates.extend(root / name for name in _PAGE_MEMORY_VISUAL_DIRS) + if root.name == "_doc_agent": + candidates.extend(root / name for name in _DEBUG_VISUAL_DIRS) + + for path in candidates: + try: + if path.exists(): + shutil.rmtree(path) + except Exception: + pass + + @worker def _render_pages_worker( queue, @@ -83,4 +121,3 @@ def render_pages( timeout=timeout, ) return list(result.get("results") or []) - diff --git a/apps/worker/app/services/document_ingestion/artifact_refs.py b/apps/worker/app/services/document_ingestion/artifact_refs.py index db554676..a72ab633 100644 --- a/apps/worker/app/services/document_ingestion/artifact_refs.py +++ b/apps/worker/app/services/document_ingestion/artifact_refs.py @@ -12,13 +12,6 @@ def collect_referenced_artifact_refs(chunks: list[dict[str, Any]]) -> set[str]: if not isinstance(metadata, dict): metadata = {} - if chunk_type == "page": - raw_page_refs = metadata.get("page_image_uris") or [] - if isinstance(raw_page_refs, list): - for raw_ref in raw_page_refs: - _add_artifact_ref(refs, raw_ref, allowed_roots={"pages"}) - continue - if chunk_type in {"image", "table"}: allowed_root = f"{chunk_type}s" _add_artifact_ref( @@ -54,6 +47,6 @@ def _normalize_client_artifact_ref(raw_ref: object) -> str | None: ] if len(parts) < 2: return None - if parts[0] not in {"pages", "images", "tables"}: + if parts[0] not in {"images", "tables"}: return None return "/".join(parts) diff --git a/apps/worker/app/services/document_ingestion/success_finalization.py b/apps/worker/app/services/document_ingestion/success_finalization.py index 9af4e7e6..d900acc3 100644 --- a/apps/worker/app/services/document_ingestion/success_finalization.py +++ b/apps/worker/app/services/document_ingestion/success_finalization.py @@ -249,12 +249,14 @@ def _upload_result_package( job_id: str, result_storage_factory: ResultStorageFactory, ) -> str: + artifact_refs = collect_referenced_artifact_refs(result_package.chunks) + add_dir = str(result_package.artifact.add_dir) if result_package.artifact.add_dir else "" + if add_dir and os.path.isfile(os.path.join(add_dir, "source.pdf")): + artifact_refs.add("source.pdf") result_bundle = result_storage_factory().upload( job_id=job_id, - result_dir=str(result_package.artifact.add_dir) - if result_package.artifact.add_dir - else "", + result_dir=add_dir, zip_file_path=generated_package.zip_file_path, - artifact_refs=collect_referenced_artifact_refs(result_package.chunks), + artifact_refs=artifact_refs, ) return result_bundle.zip_key diff --git a/apps/worker/app/services/document_parser/assets/inline_asset.py b/apps/worker/app/services/document_parser/assets/inline_asset.py index 732aaa1a..09827541 100644 --- a/apps/worker/app/services/document_parser/assets/inline_asset.py +++ b/apps/worker/app/services/document_parser/assets/inline_asset.py @@ -26,14 +26,12 @@ def build_image_asset_row( def build_table_asset_row( *, - content: str, relative_path: str, summary: str, keywords: str, know_id: str, addtime: str, ) -> ParsedRow: - del content row_content = relative_path return ParsedRow( content=row_content, diff --git a/apps/worker/app/services/document_parser/formats/docx/parser.py b/apps/worker/app/services/document_parser/formats/docx/parser.py index 59c7b297..fd05cb3f 100755 --- a/apps/worker/app/services/document_parser/formats/docx/parser.py +++ b/apps/worker/app/services/document_parser/formats/docx/parser.py @@ -423,7 +423,6 @@ def handle_table( headings_stack[-1]["content"].append(table_ref) df_list.append( build_table_asset_row( - content=tb_html_str, relative_path=table_asset.relative_path, summary=tb_summary, keywords=tb_keywords, diff --git a/apps/worker/app/services/document_parser/formats/excel/table_parser.py b/apps/worker/app/services/document_parser/formats/excel/table_parser.py index f581db96..fd57d2ee 100644 --- a/apps/worker/app/services/document_parser/formats/excel/table_parser.py +++ b/apps/worker/app/services/document_parser/formats/excel/table_parser.py @@ -249,9 +249,7 @@ def _write_excel_table_asset( keywords=keywords, know_id=know_id, addtime=time_stamp, - content=table_html_string, tokens=table_tokens, - length=len(table_html_string), path=_build_excel_table_path( request=request, sheet_name=sheet_name, diff --git a/apps/worker/app/services/document_parser/formats/markdown/table_asset.py b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py index a489e036..89558cb1 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/table_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py @@ -55,7 +55,6 @@ def build_markdown_table_asset( _write_table_html(table_path=table_path, table_html=request.table_html) table_row = build_table_asset_row( - content=request.table_html, relative_path=relative_table_path, summary=table_index, keywords="", diff --git a/apps/worker/app/services/document_parser/profiling/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py index a069b235..405bf2d6 100644 --- a/apps/worker/app/services/document_parser/profiling/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -8,6 +8,7 @@ from loguru import logger from app.services.document_agent.coordinator import ProfileCoordinator +from app.services.document_agent.visual import purge_debug_visual_dirs, visual_debug_enabled from app.services.document_parser.orchestration.oversized_pdf_policy import ( build_oversized_pdf_profile_failed_exception, build_oversized_pdf_processing_failed_exception, @@ -57,14 +58,18 @@ def profile_document( ext = os.path.splitext(filename)[1].lower() if ext == ".pdf": - return _profile_pdf( - file_path, - filename, - job_id=job_id, - output_dir=output_dir, - skip_shard_plan=skip_shard_plan, - oversized_policy=oversized_policy, - ) + try: + return _profile_pdf( + file_path, + filename, + job_id=job_id, + output_dir=output_dir, + skip_shard_plan=skip_shard_plan, + oversized_policy=oversized_policy, + ) + finally: + if not visual_debug_enabled(): + purge_debug_visual_dirs(output_dir) return ParserDocumentProfile( file_type=ext.lstrip("."), @@ -170,6 +175,7 @@ def _profile_pdf_with_db( if trace := getattr(coordinator, "trace", None): trace.persist_doc_profile(profile) + setattr(profile, "trace_recorder", trace) return profile diff --git a/apps/worker/app/services/document_parser/tables/table_asset_writer.py b/apps/worker/app/services/document_parser/tables/table_asset_writer.py index 920e0e97..6f3bf5d3 100644 --- a/apps/worker/app/services/document_parser/tables/table_asset_writer.py +++ b/apps/worker/app/services/document_parser/tables/table_asset_writer.py @@ -16,9 +16,7 @@ class TableAssetInput: keywords: str know_id: str addtime: str - content: str | None = None tokens: str = "" - length: int | None = None path: str | None = None asset_path: str | None = None diff --git a/apps/worker/app/services/page_memory/fine_hierarchy.py b/apps/worker/app/services/page_memory/fine_hierarchy.py index af2a05cf..0abc4ec7 100644 --- a/apps/worker/app/services/page_memory/fine_hierarchy.py +++ b/apps/worker/app/services/page_memory/fine_hierarchy.py @@ -12,7 +12,6 @@ import json import os import re -from pathlib import Path from typing import Any from loguru import logger @@ -30,7 +29,7 @@ def refine_fat_leaf_skeletons( tag_results: list[PageTagResult], fat_leaf_pages: set[int], model_name: str | None = None, - output_dir: str | None = None, + trace_recorder: Any | None = None, ) -> list[SectionSkeleton]: """Refine coarse TOC leaf skeletons using VLM-observed title candidates. @@ -73,7 +72,7 @@ def refine_fat_leaf_skeletons( candidates=candidates, skeleton=skeleton, model_name=model_name, - output_dir=output_dir, + trace_recorder=trace_recorder, ) if deeper: @@ -170,7 +169,7 @@ def _run_hierarchy_on_candidates( candidates: list[dict[str, Any]], skeleton: SectionSkeleton, model_name: str | None, - output_dir: str | None, + trace_recorder: Any | None, ) -> list[SectionSkeleton] | None: """Run the page-memory hierarchy prompt and rebuild a nested skeleton tree. @@ -316,8 +315,8 @@ def _run_hierarchy_on_candidates( }, )) - _save_debug_hierarchy( - output_dir=output_dir, + _record_trace_hierarchy( + trace_recorder=trace_recorder, skeleton=skeleton, candidates=candidates, ordered=ordered, @@ -365,28 +364,61 @@ def _parse_hierarchy_result(result: Any, *, max_depth: int) -> dict[int, int]: return parsed -def _save_debug_hierarchy( +def _record_trace_hierarchy( *, - output_dir: str | None, + trace_recorder: Any | None, skeleton: SectionSkeleton, candidates: list[dict[str, Any]], ordered: list[dict[str, Any]], ) -> None: - if not output_dir: + if trace_recorder is None or not hasattr(trace_recorder, "record_stage"): return - payload = { - "parent": skeleton.to_dict(), - "candidates": candidates, - "hierarchy": ordered, - } try: - debug_path = Path(output_dir) / "page_memory_fine_hierarchy.json" - debug_path.write_text( - json.dumps(payload, ensure_ascii=False, indent=2), - encoding="utf-8", + trace_recorder.record_stage( + "C4b.fine_hierarchy.scope", + page_info=_page_scope_info( + int(item.get("page") or 0) for item in candidates if item.get("page") + ), + variables={ + "parent": skeleton.to_dict(), + "candidate_count": len(candidates), + "candidates": candidates, + "hierarchy": ordered, + }, ) except Exception as exc: logger.debug( - "[page_memory.fine_hierarchy] failed to save debug hierarchy: {}", + "[page_memory.fine_hierarchy] failed to record trace hierarchy: {}", exc, ) + + +def _page_scope_info(pages: Any) -> dict[str, Any]: + normalized: list[int] = [] + for raw_page in pages or []: + try: + page = int(raw_page) + except (TypeError, ValueError): + continue + if page > 0: + normalized.append(page) + normalized = sorted(set(normalized)) + return { + "page_count": len(normalized), + "page_ranges": _collapse_page_ranges(normalized), + } + + +def _collapse_page_ranges(pages: list[int]) -> list[list[int]]: + if not pages: + return [] + ranges: list[list[int]] = [] + start = prev = pages[0] + for page in pages[1:]: + if page == prev + 1: + prev = page + continue + ranges.append([start, prev]) + start = prev = page + ranges.append([start, prev]) + return ranges diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 03a35791..e4f0bb76 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -1,14 +1,19 @@ from __future__ import annotations +import json import os -from dataclasses import dataclass +import shutil +from dataclasses import dataclass, replace from pathlib import Path from typing import Any import pandas as pd from app.services.document_agent.pdf_text import read_page_texts -from app.services.document_agent.visual import render_pages +from app.services.document_agent.visual import ( + purge_debug_visual_dirs, + visual_debug_enabled, +) from app.services.document_agent.manifest import ToolContext from app.services.document_agent.state import AgentBlackboard from app.services.document_parser.profiling.doc_profiler import profile_document @@ -29,6 +34,25 @@ class PageMemoryInput: base_url: str = "" +@dataclass(frozen=True) +class _HierarchyScope: + scope_id: str + skeletons: list[Any] + strategy: str + start_page: int + end_page: int + + +@dataclass(frozen=True) +class _ScopeRunResult: + scope_id: str + hierarchy: list[Any] + tags: list[Any] + assets_by_page: dict[int, list[Any]] + rendered: list[Any] + final_pages: list[int] + + def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: """Run the page-memory track. @@ -38,39 +62,66 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: """ full_output_dir = _resolve_output_dir(request) os.makedirs(full_output_dir, exist_ok=True) - pdf_path, pdf_filename = normalize_to_pdf( - file_path=request.file_path, - filename=request.filename, - output_dir=full_output_dir, - base_url=request.base_url, - ) - profile = profile_document( - pdf_path, - pdf_filename, - job_id=request.job_id, - output_dir=full_output_dir, - skip_shard_plan=True, - oversized_policy="page_memory", - ) - verdict = _decide_granularity(profile) - - if verdict == "whole_doc": - return full_output_dir, _build_whole_doc_dataframe( - pdf_path=pdf_path, + _cleanup_page_memory_artifacts(full_output_dir) + trace_recorder: Any | None = None + final_status = "failed" + trace_summary: dict[str, Any] = {} + try: + pdf_path, pdf_filename = normalize_to_pdf( + file_path=request.file_path, filename=request.filename, output_dir=full_output_dir, - page_count=max(int(profile.page_count or 0), 0), - verdict=verdict, + base_url=request.base_url, + ) + _persist_source_pdf(pdf_path=pdf_path, output_dir=full_output_dir) + profile = profile_document( + pdf_path, + pdf_filename, + job_id=request.job_id, + output_dir=full_output_dir, + skip_shard_plan=True, + oversized_policy="page_memory", ) + trace_recorder = getattr(profile, "trace_recorder", None) + verdict = _decide_granularity(profile) + trace_summary = { + "page_count": max(int(profile.page_count or 0), 0), + "verdict": verdict, + } - # page → per-page pipeline - return full_output_dir, _build_page_dataframe( - pdf_path=pdf_path, - filename=request.filename, - output_dir=full_output_dir, - profile=profile, - verdict=verdict, - ) + if verdict == "whole_doc": + parsed_df = _build_whole_doc_dataframe( + pdf_path=pdf_path, + filename=request.filename, + page_count=max(int(profile.page_count or 0), 0), + verdict=verdict, + trace_recorder=trace_recorder, + ) + else: + parsed_df = _build_page_dataframe( + pdf_path=pdf_path, + filename=request.filename, + output_dir=full_output_dir, + profile=profile, + verdict=verdict, + trace_recorder=trace_recorder, + ) + final_status = "success" + trace_summary["rows_count"] = len(parsed_df) + return full_output_dir, parsed_df + except Exception as exc: + trace_summary["error"] = str(exc) + raise + finally: + if trace_recorder is not None: + trace_recorder.write_trace_artifact( + full_output_dir, + final_status=final_status, + summary=trace_summary | trace_recorder.summary(), + ) + _remove_nested_doc_agent_trace(full_output_dir) + if not visual_debug_enabled(): + purge_debug_visual_dirs(full_output_dir) def _resolve_output_dir(request: PageMemoryInput) -> str: @@ -78,6 +129,13 @@ def _resolve_output_dir(request: PageMemoryInput) -> str: return os.path.join(request.output_dir, Path(output_name).stem) +def _persist_source_pdf(*, pdf_path: str, output_dir: str) -> str: + source_pdf_path = os.path.join(output_dir, "source.pdf") + if Path(pdf_path).resolve() != Path(source_pdf_path).resolve(): + shutil.copyfile(pdf_path, source_pdf_path) + return source_pdf_path + + def _decide_granularity(profile: Any) -> str: """Decide granularity for a page-memory document. @@ -105,6 +163,7 @@ def _build_page_dataframe( output_dir: str, profile: Any, verdict: str, + trace_recorder: Any | None = None, ) -> pd.DataFrame: """Build per-page DataFrame via the full C1-C7 pipeline. @@ -115,22 +174,17 @@ def _build_page_dataframe( C3 page_tagger → PageTagResult[] C3b title detection → observed_titles C4b fine_hierarchy → refined SectionSkeleton[] + C5 page_assets → assets anchored to refined hierarchy pages C7 assemble node-granularity DataFrame """ from app.services.document_agent.budget import BudgetTracker, StageEnvelope - from app.services.page_memory.page_plan import derive_page_processing_plan - from app.services.page_memory.page_renderer import render_document_pages - from app.services.page_memory.page_tagger import tag_pages from app.services.page_memory.skeleton_extractor import ( SectionSkeleton, extract_section_skeletons, ) from app.services.page_memory.page_assets import ( - extract_page_assets_from_renders, get_asset_budget, - get_asset_confidence_threshold, get_asset_max_pages, - get_asset_model, page_asset_extraction_enabled, ) @@ -200,6 +254,7 @@ def _build_page_dataframe( output_dir=output_dir, page_count=page_count, budget=budget, + trace_recorder=trace_recorder, ) # ── C4: skeleton (from profile anatomy) ─────────────────────────── @@ -220,92 +275,6 @@ def _build_page_dataframe( len(skeletons), ) - # ── C1: render pages ────────────────────────────────────────────── - page_features = anatomy.page_features if anatomy else [] - rendered = render_document_pages( - pdf_path=pdf_path, - page_count=page_count, - output_dir=output_dir, - page_features=page_features, - page_texts=page_texts, - ) - - page_assets_by_page: dict[int, list[Any]] = {} - if asset_extraction_enabled: - page_assets_by_page = extract_page_assets_from_renders( - pdf_path=pdf_path, - rendered_pages=rendered, - output_dir=output_dir, - model_name=get_asset_model(), - budget=budget, - max_pages=get_asset_max_pages(page_count), - confidence_threshold=get_asset_confidence_threshold(), - ) - - # ── C2: page plan ───────────────────────────────────────────────── - page_labels = anatomy.page_labels if anatomy else [] - plans = derive_page_processing_plan( - page_count=page_count, - page_labels=page_labels, - page_features=page_features, - ) - - # ── C3: page tagger ────────────────────────────────────────────── - vlm_model = os.environ.get("IMAGE_MODEL") - tags = tag_pages( - pages=rendered, - plans=plans, - budget=budget, - vlm_model=vlm_model, - ) - - # ── C3b + C4b: page-native hierarchy refinement ───────────────── - if skeletons: - from app.services.page_memory.fine_hierarchy import ( - compute_fat_leaf_pages, - refine_fat_leaf_skeletons, - ) - from app.services.page_memory.page_tagger import ( - get_fine_min_pages, - tag_page_titles, - ) - - fine_min = get_fine_min_pages() - fat_leaf_pages = compute_fat_leaf_pages(skeletons, min_pages=fine_min) - logger.info( - "[page_memory] native hierarchy: {} fat-leaf pages (min={})", - len(fat_leaf_pages), fine_min, - ) - - # Step 2: VLM title detection on fat-leaf pages - if fat_leaf_pages: - tags = tag_page_titles( - pages=rendered, - tag_results=tags, - fat_leaf_pages=fat_leaf_pages, - budget=budget, - vlm_model=vlm_model, - ) - - # Step 3: Refine coarse skeletons with LLM hierarchy - skeletons = refine_fat_leaf_skeletons( - coarse_skeletons=skeletons, - tag_results=tags, - fat_leaf_pages=fat_leaf_pages, - model_name=os.environ.get( - "PAGE_MEMORY_HIERARCHY_MODEL", - os.environ.get( - "HIERARCHY_LLM_MODEL", - os.environ.get("NORMOL_MODEL"), - ), - ), - output_dir=output_dir, - ) - logger.info( - "[page_memory] C4b refined: {} sections after fine hierarchy", - len(skeletons), - ) - if not skeletons: skeletons = [ SectionSkeleton( @@ -315,12 +284,104 @@ def _build_page_dataframe( end_page=page_count, title="Root", parent_path=filename, + evidence={"source": "fallback_root", "reason": "no_hierarchy"}, + ) + ] + + coarse_scopes = _build_hierarchy_scopes( + skeletons=skeletons, + filename=filename, + page_count=page_count, + ) + coarse_pages_scope = _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ) + _record_trace_stage( + trace_recorder, + "C4.skeleton", + page_info={ + "document_page_count": page_count, + "coarse_scope": _page_scope_info(coarse_pages_scope), + }, + variables={ + "sections": _summarize_skeletons(skeletons), + "scope_count": len(coarse_scopes), + "scopes": [ + _scope_manifest( + scope_id=scope.scope_id, + skeletons=scope.skeletons, + page_count=page_count, + strategy=scope.strategy, + ) + for scope in coarse_scopes + ], + }, + ) + + page_features = anatomy.page_features if anatomy else [] + page_labels = anatomy.page_labels if anatomy else [] + vlm_model = os.environ.get("IMAGE_MODEL") + + logger.info( + "[page_memory] processing {} hierarchy scopes (production-shaped)", + len(coarse_scopes), + ) + scope_results: list[_ScopeRunResult] = [] + asset_pages_remaining = ( + get_asset_max_pages(page_count) if asset_extraction_enabled else 0 + ) + for index, scope in enumerate(coarse_scopes, start=1): + result = _run_hierarchy_scope( + scope=scope, + scope_index=index, + scope_count=len(coarse_scopes), + pdf_path=pdf_path, + filename=filename, + output_dir=output_dir, + page_count=page_count, + page_texts=page_texts, + page_features=page_features, + page_labels=page_labels, + budget=budget, + vlm_model=vlm_model, + asset_extraction_enabled=asset_extraction_enabled, + asset_max_pages=asset_pages_remaining, + trace_recorder=trace_recorder, + ) + scope_results.append(result) + if asset_extraction_enabled: + asset_pages_remaining = max( + asset_pages_remaining - len(result.final_pages), + 0, ) + + skeletons = _sort_skeletons( + [ + skeleton + for result in scope_results + for skeleton in result.hierarchy ] + ) + tags = _merge_page_tags(result.tags for result in scope_results) + page_assets_by_page = _merge_assets_by_page( + result.assets_by_page for result in scope_results + ) + rendered_map: dict[int, Any] = {} + for result in scope_results: + for rendered_page in result.rendered: + rendered_map.setdefault(rendered_page.page_index, rendered_page) + + _write_top_level_artifacts( + output_dir=output_dir, + hierarchy=skeletons, + tags=tags, + assets_by_page=page_assets_by_page, + ) # ── C7: assemble DataFrame rows ────────────────────────────────── tag_map = {t.page_index: t for t in tags} - render_map = {r.page_index: r for r in rendered} + render_map = rendered_map # Build page → PageLabel.kind lookup label_map: dict[int, str] = {} @@ -330,23 +391,22 @@ def _build_page_dataframe( # Shared per-page lookups for node-granularity assembly. raw_text_by_page: dict[int, str] = {} - image_uri_by_page: dict[int, str] = {} image_path_by_page: dict[int, str] = {} - for page in range(1, page_count + 1): + final_pages_scope = _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ) + for page in final_pages_scope: rend = render_map.get(page) raw_text_by_page[page] = (rend.raw_text if rend else page_texts.get(page, "")) or "" if rend and rend.image_path and os.path.exists(rend.image_path): image_path_by_page[page] = rend.image_path - image_uri_by_page[page] = str( - Path(rend.image_path).relative_to(output_dir) - ) from app.services.page_memory.node_assembler import build_node_rows rows = build_node_rows( skeletons=skeletons, raw_text_by_page=raw_text_by_page, - image_uri_by_page=image_uri_by_page, image_path_by_page=image_path_by_page, kind_by_page=label_map, tag_by_page=tag_map, @@ -360,6 +420,16 @@ def _build_page_dataframe( "[page_memory] C7 assembled {} node rows (verdict={})", len(rows), verdict, ) + _record_trace_stage( + trace_recorder, + "C7.node_assembly", + page_info=_page_scope_info(final_pages_scope), + variables={ + "row_count": len(rows), + "scope_count": len(scope_results), + "page_to_node": _page_to_node_map(rows), + }, + ) return pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) @@ -370,6 +440,7 @@ def _build_page_ctx( output_dir: str, page_count: int, budget: Any, + trace_recorder: Any | None = None, ) -> ToolContext: """Construct a ToolContext for C4 sub-agent and C3 tagger VLM calls.""" blackboard = AgentBlackboard() @@ -384,7 +455,7 @@ def _build_page_ctx( job_id=job_id, blackboard=blackboard, budget=budget, - trace=None, + trace=trace_recorder, output_dir=output_dir, settings={ "vlm_model": vlm_model, @@ -392,6 +463,364 @@ def _build_page_ctx( }, ) + +def _build_hierarchy_scopes( + *, + skeletons: list[Any], + filename: str, + page_count: int, +) -> list[_HierarchyScope]: + if not skeletons: + return [] + + root_fallback = all( + getattr(item, "title", "") == "Root" + or (getattr(item, "evidence", {}) or {}).get("source") == "fallback_root" + for item in skeletons + ) + if root_fallback: + ordered = _sort_skeletons(skeletons) + return [ + _HierarchyScope( + scope_id=_scope_id_for_pages(1, page_count), + skeletons=ordered, + strategy="fallback_root", + start_page=1, + end_page=page_count, + ) + ] + + min_level = min(int(getattr(item, "level", 0) or 0) for item in skeletons) + top_nodes = [ + item + for item in skeletons + if int(getattr(item, "level", 0) or 0) == min_level + or not str(getattr(item, "parent_path", "") or "").startswith(f"{filename}/") + ] + seen_top_paths: set[str] = set() + unique_top_nodes: list[Any] = [] + for item in _sort_skeletons(top_nodes): + if item.section_path in seen_top_paths: + continue + seen_top_paths.add(item.section_path) + unique_top_nodes.append(item) + + scopes: list[_HierarchyScope] = [] + for index, top in enumerate(unique_top_nodes): + prefix = f"{top.section_path}/" + members = [ + item + for item in skeletons + if item.section_path == top.section_path + or str(item.section_path).startswith(prefix) + ] + if not members: + continue + start_page = max(1, int(getattr(top, "start_page", 1) or 1)) + next_top_start = ( + int(getattr(unique_top_nodes[index + 1], "start_page", page_count + 1) or page_count + 1) + if index + 1 < len(unique_top_nodes) + else page_count + 1 + ) + end_page = min( + page_count, + max( + start_page, + min( + int(getattr(top, "end_page", page_count) or page_count), + next_top_start - 1, + ), + ), + ) + bounded_members = [ + replace( + item, + start_page=max( + start_page, + int(getattr(item, "start_page", start_page) or start_page), + ), + end_page=min( + end_page, + int(getattr(item, "end_page", end_page) or end_page), + ), + ) + for item in members + if int(getattr(item, "start_page", 1) or 1) <= end_page + and int(getattr(item, "end_page", end_page) or end_page) >= start_page + ] + bounded_members = _sort_skeletons(bounded_members) + scopes.append( + _HierarchyScope( + scope_id=_scope_id_for_pages(start_page, end_page), + skeletons=bounded_members, + strategy=f"coarse_scope_{index + 1}", + start_page=start_page, + end_page=end_page, + ) + ) + if scopes: + return scopes + + ordered = _sort_skeletons(skeletons) + pages = _derive_hierarchy_page_scope(skeletons=ordered, page_count=page_count) + return [ + _HierarchyScope( + scope_id=_scope_id_for_pages( + pages[0] if pages else 1, + pages[-1] if pages else page_count, + ), + skeletons=ordered, + strategy="full_coarse_hierarchy", + start_page=pages[0] if pages else 1, + end_page=pages[-1] if pages else page_count, + ) + ] + + +def _run_hierarchy_scope( + *, + scope: _HierarchyScope, + scope_index: int, + scope_count: int, + pdf_path: str, + filename: str, + output_dir: str, + page_count: int, + page_texts: dict[int, str], + page_features: list[Any], + page_labels: list[Any], + budget: Any, + vlm_model: str | None, + asset_extraction_enabled: bool, + asset_max_pages: int, + trace_recorder: Any | None, +) -> _ScopeRunResult: + from app.services.page_memory.fine_hierarchy import ( + compute_fat_leaf_pages, + refine_fat_leaf_skeletons, + ) + from app.services.page_memory.page_assets import ( + extract_page_assets_from_renders, + get_asset_confidence_threshold, + get_asset_model, + ) + from app.services.page_memory.page_plan import derive_page_processing_plan + from app.services.page_memory.page_renderer import render_document_pages + from app.services.page_memory.page_tagger import ( + PageTagResult, + get_fine_min_pages, + tag_page_titles, + tag_pages, + ) + + scope_skeletons = _sort_skeletons(scope.skeletons) + scope_manifest = _scope_manifest( + scope_id=scope.scope_id, + skeletons=scope_skeletons, + page_count=page_count, + strategy=scope.strategy, + ) + coarse_pages = _derive_hierarchy_page_scope( + skeletons=scope_skeletons, + page_count=page_count, + ) + logger.info( + "[page_memory] scope {}/{} {} coarse ranges={}", + scope_index, + scope_count, + scope.scope_id, + _collapse_page_ranges(coarse_pages), + ) + _record_trace_stage( + trace_recorder, + "C4.coarse_scope", + page_info=_page_scope_info(coarse_pages), + variables={ + "scope_index": scope_index, + "scope_count": scope_count, + "scope": scope_manifest, + }, + ) + + fine_min = get_fine_min_pages() + fat_leaf_pages = compute_fat_leaf_pages(scope_skeletons, min_pages=fine_min) + if fat_leaf_pages: + title_pages = sorted(fat_leaf_pages) + title_rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=output_dir, + pages=title_pages, + page_features=page_features, + page_texts=page_texts, + ) + title_tags = [ + PageTagResult( + page_index=page, + summary="", + keywords=[], + strategy_used="title_detection_only", + ) + for page in title_pages + ] + title_tags = tag_page_titles( + pages=title_rendered, + tag_results=title_tags, + fat_leaf_pages=fat_leaf_pages, + budget=budget, + vlm_model=vlm_model, + ) + _record_trace_stage( + trace_recorder, + "C3b.title_detection", + page_info=_page_scope_info(title_pages), + variables={ + "scope_id": scope.scope_id, + "tags": _summarize_tags(title_tags), + }, + ) + scope_skeletons = refine_fat_leaf_skeletons( + coarse_skeletons=scope_skeletons, + tag_results=title_tags, + fat_leaf_pages=fat_leaf_pages, + model_name=os.environ.get( + "PAGE_MEMORY_HIERARCHY_MODEL", + os.environ.get( + "HIERARCHY_LLM_MODEL", + os.environ.get("NORMOL_MODEL"), + ), + ), + trace_recorder=trace_recorder, + ) + else: + logger.info( + "[page_memory] scope {} no fat leaves (min={}); skip fine hierarchy", + scope.scope_id, + fine_min, + ) + + scope_manifest = _scope_manifest( + scope_id=scope.scope_id, + skeletons=scope_skeletons, + page_count=page_count, + strategy=f"{scope.strategy}:refined", + ) + _record_trace_stage( + trace_recorder, + "C4b.fine_hierarchy", + page_info={"fat_leaf": _page_scope_info(sorted(fat_leaf_pages))}, + variables={ + "scope_id": scope.scope_id, + "sections": _summarize_skeletons(scope_skeletons), + "scope": scope_manifest, + }, + ) + + final_pages = _derive_hierarchy_page_scope( + skeletons=scope_skeletons, + page_count=page_count, + ) + final_scope_summary = _summarize_tag_scope( + skeletons=scope_skeletons, + page_count=page_count, + pages=final_pages, + ) + rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=output_dir, + pages=final_pages, + page_features=page_features, + page_texts=page_texts, + ) + _record_trace_stage( + trace_recorder, + "C1.render_pages", + page_info=_page_scope_info([item.page_index for item in rendered]), + variables={ + "scope_id": scope.scope_id, + "rendered_count": len(rendered), + "tag_scope": final_scope_summary, + }, + ) + + plans = derive_page_processing_plan( + page_count=page_count, + page_labels=page_labels, + page_features=page_features, + ) + final_page_set = set(final_pages) + plans = [plan for plan in plans if plan.page_index in final_page_set] + _record_trace_stage( + trace_recorder, + "C2.page_plan", + page_info=_page_scope_info([getattr(plan, "page_index", None) for plan in plans]), + variables={"scope_id": scope.scope_id, "plan_count": len(plans)}, + ) + + tags = tag_pages( + pages=rendered, + plans=plans, + budget=budget, + vlm_model=vlm_model, + ) + _record_trace_stage( + trace_recorder, + "C3.page_tagger", + page_info=_page_scope_info([tag.page_index for tag in tags]), + variables={"scope_id": scope.scope_id, "tags": _summarize_tags(tags)}, + ) + _write_scope_artifacts( + output_dir=output_dir, + scope_id=scope.scope_id, + scope_manifest=scope_manifest, + hierarchy=scope_skeletons, + tags=tags, + ) + + assets_by_page: dict[int, list[Any]] = {} + if asset_extraction_enabled and asset_max_pages > 0: + assets_by_page = extract_page_assets_from_renders( + pdf_path=pdf_path, + rendered_pages=rendered, + output_dir=output_dir, + model_name=get_asset_model(), + budget=budget, + max_pages=asset_max_pages, + confidence_threshold=get_asset_confidence_threshold(), + ) + _record_trace_stage( + trace_recorder, + "C5.page_assets", + page_info=_page_scope_info(sorted(assets_by_page)), + variables={ + "scope_id": scope.scope_id, + "asset_count": sum(len(items) for items in assets_by_page.values()), + "assets_by_page": { + page: [asset.asset_id for asset in assets] + for page, assets in assets_by_page.items() + }, + }, + ) + _write_scope_artifacts( + output_dir=output_dir, + scope_id=scope.scope_id, + scope_manifest=scope_manifest, + hierarchy=scope_skeletons, + tags=tags, + assets_by_page=assets_by_page, + ) + + return _ScopeRunResult( + scope_id=scope.scope_id, + hierarchy=scope_skeletons, + tags=tags, + assets_by_page=assets_by_page, + rendered=rendered, + final_pages=final_pages, + ) + + # ── whole_doc builder (PR3, unchanged) ──────────────────────────────── @@ -399,20 +828,14 @@ def _build_whole_doc_dataframe( *, pdf_path: str, filename: str, - output_dir: str, page_count: int, verdict: str, + trace_recorder: Any | None = None, ) -> pd.DataFrame: pages = list(range(1, page_count + 1)) if page_count > 0 else [1] page_texts = read_page_texts(pdf_path, pages) raw_text = "\n\n".join(page_texts.get(page, "") for page in pages).strip() summary = _build_summary(filename=filename, page_count=page_count, raw_text=raw_text) - page_image_uris = _render_page_images( - pdf_path=pdf_path, - output_dir=output_dir, - page_count=page_count, - pages=pages, - ) content = f"[SUMMARY]\n{summary}\n\n[RAW]\n{raw_text}".strip() know_id = gen_str_codes(f"wholedoc::{filename}::{content}") row = { @@ -427,10 +850,14 @@ def _build_whole_doc_dataframe( "connectto": "", "addtime": get_str_time(), "page_nums": ",".join(str(page) for page in pages), - "extra_metadata": { - "page_image_uris": page_image_uris, - }, + "extra_metadata": {}, } + _record_trace_stage( + trace_recorder, + "whole_doc", + page_info=_page_scope_info(pages), + variables={"summary": summary, "verdict": verdict}, + ) return pd.DataFrame([row], columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) @@ -440,29 +867,392 @@ def _build_summary(*, filename: str, page_count: int, raw_text: str) -> str: return f"{prefix}: {preview}" if preview else prefix -def _render_page_images( +def _record_trace_stage( + trace_recorder: Any | None, + stage: str, + *, + page_info: dict[str, Any], + variables: dict[str, Any], +) -> None: + if trace_recorder is None or not hasattr(trace_recorder, "record_stage"): + return + trace_recorder.record_stage(stage, page_info=page_info, variables=variables) + + +def _cleanup_page_memory_artifacts(output_dir: str) -> None: + root = Path(output_dir) + legacy_files = { + "assets.json", + "chunks.json", + "coarse_tag_scope.json", + "doc_nav.json", + "hierarchy.json", + "manifest.json", + "node_rows.csv", + "node_rows.json", + "page_memory_fine_hierarchy.json", + "page_plans.json", + "page_rendered.json", + "page_tags.json", + "page_tags_after_titles.json", + "page_tags_pre_hierarchy.json", + "report.md", + "skeletons.json", + "tag_scope.json", + "trace.json", + } + for name in legacy_files: + path = root / name + try: + if path.is_file(): + path.unlink() + except Exception: + logger.debug("[page_memory] failed to cleanup artifact {}", path) + for name in ("asset_annotate", "debug", "fine_hierarchy", "images", "pages", "scopes", "tables"): + path = root / name + try: + if path.is_dir(): + shutil.rmtree(path) + except Exception: + logger.debug("[page_memory] failed to cleanup artifact dir {}", path) + + +def _remove_nested_doc_agent_trace(output_dir: str) -> None: + try: + (Path(output_dir) / "_doc_agent" / "trace.json").unlink() + except FileNotFoundError: + pass + except Exception as exc: + logger.debug("[page_memory] failed to remove nested doc-agent trace: {}", exc) + + +def _sort_skeletons(skeletons: list[Any]) -> list[Any]: + return sorted( + skeletons, + key=lambda item: ( + int(getattr(item, "start_page", 0) or 0), + int(getattr(item, "level", 0) or 0), + str(getattr(item, "section_path", "") or ""), + ), + ) + + +def _exclusive_end_for_skeleton(skeletons: list[Any], target: Any) -> int: + later_starts = [ + int(getattr(item, "start_page", 0) or 0) + for item in skeletons + if int(getattr(item, "start_page", 0) or 0) + > int(getattr(target, "start_page", 0) or 0) + ] + end_page = int(getattr(target, "end_page", 1) or 1) + if not later_starts: + return end_page + return min(end_page, min(later_starts) - 1) + + +def _merge_page_tags(tag_groups: Any) -> list[Any]: + by_page: dict[int, Any] = {} + for tags in tag_groups: + for tag in tags: + by_page[int(getattr(tag, "page_index", 0) or 0)] = tag + return [by_page[page] for page in sorted(by_page) if page > 0] + + +def _merge_assets_by_page(asset_groups: Any) -> dict[int, list[Any]]: + merged: dict[int, list[Any]] = {} + seen_ids: set[str] = set() + for assets_by_page in asset_groups: + for page, assets in assets_by_page.items(): + for asset in assets: + asset_id = str(getattr(asset, "asset_id", "") or "") + if asset_id and asset_id in seen_ids: + continue + if asset_id: + seen_ids.add(asset_id) + merged.setdefault(int(page), []).append(asset) + return {page: merged[page] for page in sorted(merged)} + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, default=str), + encoding="utf-8", + ) + + +def _scope_id_for_pages(start_page: int, end_page: int) -> str: + return f"p{max(1, int(start_page))}-{max(1, int(end_page))}" + + +def _scope_manifest( + *, + scope_id: str, + skeletons: list[Any], + page_count: int, + strategy: str, +) -> dict[str, Any]: + pages = _derive_hierarchy_page_scope(skeletons=skeletons, page_count=page_count) + parent_paths = sorted({str(getattr(item, "parent_path", "") or "") for item in skeletons}) + return { + "scope_id": scope_id, + "strategy": strategy, + "document_page_count": page_count, + "page_count": len(pages), + "page_ranges": _collapse_page_ranges(pages), + "skeleton_count": len(skeletons), + "parent_paths": parent_paths, + } + + +def _write_scope_artifacts( *, - pdf_path: str, output_dir: str, + scope_id: str, + scope_manifest: dict[str, Any], + hierarchy: list[Any], + tags: list[Any], + assets_by_page: dict[int, list[Any]] | None = None, +) -> None: + scope_dir = Path(output_dir) / "scopes" / scope_id + _write_json(scope_dir / "scope.json", scope_manifest) + _write_json( + scope_dir / "fine_hierarchy.json", + _serialize_hierarchy_artifact(hierarchy, scope_manifest=scope_manifest), + ) + _write_json(scope_dir / "page_tags.json", _serialize_page_tags(tags)) + if assets_by_page: + _write_json(scope_dir / "assets.json", _serialize_assets(assets_by_page)) + + +def _write_top_level_artifacts( + *, + output_dir: str, + hierarchy: list[Any], + tags: list[Any], + assets_by_page: dict[int, list[Any]] | None = None, +) -> None: + root = Path(output_dir) + _write_json(root / "hierarchy.json", _serialize_hierarchy_artifact(hierarchy)) + _write_json(root / "page_tags.json", _serialize_page_tags(tags)) + if assets_by_page: + _write_json(root / "assets.json", _serialize_assets(assets_by_page)) + else: + try: + (root / "assets.json").unlink() + except FileNotFoundError: + pass + + +def _serialize_skeletons(skeletons: list[Any]) -> list[dict[str, Any]]: + return [ + { + "section_path": item.section_path, + "title": item.title, + "level": item.level, + "start_page": item.start_page, + "end_page": item.end_page, + "parent_path": item.parent_path, + } + for item in skeletons + ] + + +def _serialize_hierarchy_artifact( + skeletons: list[Any], + *, + scope_manifest: dict[str, Any] | None = None, +) -> dict[str, Any]: + nodes = _serialize_skeletons(skeletons) + artifact: dict[str, Any] = { + "HIERARCHY": _build_hierarchy_tree(skeletons), + "nodes": nodes, + "stats": { + "node_count": len(nodes), + **_page_scope_info( + _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=_max_skeleton_end_page(skeletons), + ), + ), + "max_depth": max( + (int(getattr(item, "level", 0) or 0) for item in skeletons), + default=0, + ), + }, + } + if scope_manifest is not None: + artifact["scope"] = scope_manifest + return artifact + + +def _build_hierarchy_tree(skeletons: list[Any]) -> dict[str, Any]: + hierarchy: dict[str, Any] = {} + for skel in _sort_skeletons(skeletons): + parts = str(getattr(skel, "section_path", "") or "").split("/") + section_parts = parts[1:] if len(parts) > 1 else parts + current = hierarchy + for part in section_parts: + if not part: + continue + current = current.setdefault(part, {}) + return hierarchy + + +def _serialize_page_tags(tags: list[Any]) -> list[dict[str, Any]]: + return [ + { + "page_index": item.page_index, + "summary": item.summary, + "keywords": list(item.keywords), + "strategy_used": item.strategy_used, + } + for item in tags + ] + + +def _serialize_assets(assets_by_page: dict[int, list[Any]]) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for page_index in sorted(assets_by_page): + for asset in assets_by_page[page_index]: + rows.append( + { + "asset_id": asset.asset_id, + "page_index": asset.page_index, + "asset_index": asset.asset_index, + "kind": asset.kind, + "bbox_px": asset.bbox_px, + "confidence": asset.confidence, + "title": asset.title, + "summary": asset.summary, + "keywords": list(asset.keywords), + "image_uri": asset.image_uri, + "html_uri": asset.html_uri, + "extraction_status": asset.extraction_status, + } + ) + return rows + + +def _summarize_skeletons(skeletons: list[Any]) -> list[dict[str, Any]]: + return [ + { + "path": item.section_path, + "level": item.level, + "start_page": item.start_page, + "end_page": item.end_page, + } + for item in skeletons + ] + + +def _derive_hierarchy_page_scope( + *, + skeletons: list[Any], page_count: int, - pages: list[int], -) -> list[str]: +) -> list[int]: + """Pages that should receive PAGE-TAG after hierarchy anchoring. + + Page-memory first anchors a hierarchy. PAGE-TAG should then run only on + pages covered by anchored sections. If hierarchy anchoring collapses to the + Root fallback, retain the legacy full-document behavior. + """ if page_count <= 0: return [] - blackboard = AgentBlackboard() - blackboard.page_count = page_count - ctx = ToolContext( - pdf_path=pdf_path, - job_id="page_memory_render", - blackboard=blackboard, - budget=None, - trace=None, - output_dir=output_dir, - settings={}, + if not skeletons: + return list(range(1, page_count + 1)) + + has_only_root_fallback = all( + getattr(item, "title", "") == "Root" + or (getattr(item, "evidence", {}) or {}).get("source") == "fallback_root" + for item in skeletons ) - rendered = render_pages(ctx, pages, folder_name="pages", prefix="page", timeout=180) + if has_only_root_fallback: + return list(range(1, page_count + 1)) + + pages: set[int] = set() + for item in skeletons: + start_page = max(1, int(getattr(item, "start_page", 1) or 1)) + end_page = min(page_count, int(getattr(item, "end_page", start_page) or start_page)) + if end_page < start_page: + continue + pages.update(range(start_page, end_page + 1)) + return sorted(pages) or list(range(1, page_count + 1)) + + +def _summarize_tag_scope( + *, + skeletons: list[Any], + page_count: int, + pages: list[int], +) -> dict[str, Any]: + return { + "strategy": "hierarchy_anchored_pages", + "document_page_count": page_count, + "skeleton_count": len(skeletons), + **_page_scope_info(pages), + } + + +def _page_scope_info(pages: Any) -> dict[str, Any]: + normalized: list[int] = [] + for raw_page in pages or []: + try: + page = int(raw_page) + except (TypeError, ValueError): + continue + if page > 0: + normalized.append(page) + normalized = sorted(set(normalized)) + return { + "page_count": len(normalized), + "page_ranges": _collapse_page_ranges(normalized), + } + + +def _max_skeleton_end_page(skeletons: list[Any]) -> int: + return max( + (int(getattr(item, "end_page", 0) or 0) for item in skeletons), + default=0, + ) + + +def _collapse_page_ranges(pages: list[int]) -> list[list[int]]: + if not pages: + return [] + ranges: list[list[int]] = [] + start = prev = pages[0] + for page in pages[1:]: + if page == prev + 1: + prev = page + continue + ranges.append([start, prev]) + start = prev = page + ranges.append([start, prev]) + return ranges + + +def _summarize_tags(tags: list[Any]) -> list[dict[str, Any]]: return [ - str(Path(item["png_path"]).relative_to(output_dir)) - for item in rendered - if item.get("png_path") + { + "page": item.page_index, + "title": getattr(item, "title", ""), + "summary": getattr(item, "summary", ""), + "kind": getattr(item, "kind", ""), + } + for item in tags ] + + +def _page_to_node_map(rows: list[dict[str, Any]]) -> dict[int, str]: + mapping: dict[int, str] = {} + for row in rows: + if row.get("type") != "page": + continue + path = str(row.get("path") or "") + for raw_page in str(row.get("page_nums") or "").split(","): + try: + mapping[int(raw_page.strip())] = path + except ValueError: + continue + return mapping diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index c739d6d4..881c6cd3 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -6,13 +6,12 @@ - One chunk per leaf node (path). Internal/structural nodes live in the navigation tree only (their summaries aggregate bottom-up). -- A page that belongs to multiple nodes is *referenced* by each of them; the - page image is never duplicated (``page_image_uris`` is a list of shared - references). +- A page that belongs to multiple nodes is *referenced* by each of them via + page numbers; retrieval resolves those pages to a cropped PDF on demand. - A page's body text is stored **once**, under the first leaf (in reading order) that covers it. Other nodes that share the page emit a ``SAME-AS `` marker instead of repeating the text. Body text is - not a core asset in page-track — the page image is. + not duplicated across page-track nodes. Summary/keywords are settled per node (see ``node_summary``): a node covering multiple pages is summarized as a whole, and a page hosting multiple nodes is @@ -85,10 +84,11 @@ def identify_leaf_nodes(skeletons: list[SectionSkeleton]) -> list[LeafNode]: parent_paths = {skel.parent_path for skel in skeletons if skel.parent_path} leaves: list[tuple[int, int, LeafNode]] = [] for index, skel in enumerate(skeletons): + effective_end_page = _exclusive_end(skeletons, index) is_internal = skel.section_path in parent_paths body_pages: tuple[int, ...] | None = None if is_internal: - body_page_list = _internal_body_pages(skel, skeletons) + body_page_list = _internal_body_pages(skel, skeletons, index=index) if not body_page_list: continue body_pages = tuple(body_page_list) @@ -103,7 +103,7 @@ def identify_leaf_nodes(skeletons: list[SectionSkeleton]) -> list[LeafNode]: title=skel.title, level=skel.level, start_page=skel.start_page, - end_page=skel.end_page, + end_page=effective_end_page, parent_path=skel.parent_path, body_pages=body_pages, ), @@ -116,19 +116,38 @@ def identify_leaf_nodes(skeletons: list[SectionSkeleton]) -> list[LeafNode]: def _internal_body_pages( skel: SectionSkeleton, skeletons: list[SectionSkeleton], + *, + index: int, ) -> list[int]: """Pages owned by an internal section itself, excluding descendants.""" - pages = set(range(skel.start_page, skel.end_page + 1)) + pages = set(range(skel.start_page, _exclusive_end(skeletons, index) + 1)) descendant_prefix = f"{skel.section_path}/" - for other in skeletons: + for other_index, other in enumerate(skeletons): if other.section_path == skel.section_path: continue if not other.section_path.startswith(descendant_prefix): continue - pages.difference_update(range(other.start_page, other.end_page + 1)) + pages.difference_update( + range(other.start_page, _exclusive_end(skeletons, other_index) + 1) + ) return sorted(pages) +def _exclusive_end(skeletons: list[SectionSkeleton], index: int) -> int: + """Return the last page owned by a skeleton before the next sibling starts. + + The locator emits closed-closed section ranges, so adjacent sections can + overlap at the boundary page. Ownership is exclusive at the next start page. + """ + skeleton = skeletons[index] + later_starts = [ + skel.start_page for skel in skeletons if skel.start_page > skeleton.start_page + ] + if not later_starts: + return skeleton.end_page + return min(skeleton.end_page, min(later_starts) - 1) + + def assign_pages_to_leaves( leaves: list[LeafNode], *, @@ -478,7 +497,6 @@ def build_node_rows( *, skeletons: list[SectionSkeleton], raw_text_by_page: dict[int, str], - image_uri_by_page: dict[int, str], image_path_by_page: dict[int, str], kind_by_page: dict[int, str], tag_by_page: dict[int, PageTagResult], @@ -523,11 +541,6 @@ def build_node_rows( vlm_model=vlm_model, budget=budget, ) - page_image_uris = [ - image_uri_by_page[page] - for page in view.pages - if image_uri_by_page.get(page) - ] know_id = f"node_{gen_str_codes(f'{filename}::{leaf.section_path}')}" row = { "content": content, @@ -541,9 +554,7 @@ def build_node_rows( "connectto": "", "addtime": get_str_time(), "page_nums": ",".join(str(page) for page in view.pages), - "extra_metadata": { - "page_image_uris": page_image_uris, - }, + "extra_metadata": {}, } rows.append(row) rows_by_path[leaf.section_path] = row diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index fb8e1545..b04c71b5 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -17,6 +17,7 @@ from PIL import Image, ImageDraw, ImageFont from app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.visual import visual_debug_enabled from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time from app.services.page_memory.page_renderer import PageRenderResult from shared.services.ai.prompt_service import build_prompt @@ -375,11 +376,12 @@ def extract_page_assets_from_renders( page_assets.append(asset) if page_assets: assets_by_page[page.page_index] = page_assets - annotate_page_assets( - page=page, - assets=page_assets, - output_dir=output_dir, - ) + if visual_debug_enabled(): + annotate_page_assets( + page=page, + assets=page_assets, + output_dir=output_dir, + ) logger.info( "[page_assets] extracted {} assets across {} pages", sum(len(items) for items in assets_by_page.values()), @@ -466,15 +468,11 @@ def build_asset_rows( "addtime": get_str_time(), "page_nums": str(asset.page_index), "extra_metadata": { - "page_index": asset.page_index, "asset_index": asset.asset_index, - "asset_kind": asset.kind, "bbox_px": asset.bbox_px, "confidence": asset.confidence, - "title": asset.title, "caption": asset.caption, "image_uri": asset.image_uri, - "html_uri": asset.html_uri, "extraction_status": asset.extraction_status, "table_engine": "tabula" if asset.kind == "table" else "", }, diff --git a/apps/worker/app/services/page_memory/page_renderer.py b/apps/worker/app/services/page_memory/page_renderer.py index fdada1ac..89a0a756 100644 --- a/apps/worker/app/services/page_memory/page_renderer.py +++ b/apps/worker/app/services/page_memory/page_renderer.py @@ -44,6 +44,7 @@ def render_document_pages( pdf_path: str, page_count: int, output_dir: str, + pages: list[int] | None = None, page_features: list[PageFeature] | None = None, page_texts: dict[int, str] | None = None, ctx: ToolContext | None = None, @@ -58,6 +59,8 @@ def render_document_pages( Local filesystem path to the PDF. page_count: Total page count (avoids re-opening the PDF). + pages: + Optional 1-based page subset to render. Defaults to every page. output_dir: Root output directory; pages are written to ``output_dir/pages/``. page_features: @@ -80,19 +83,23 @@ def render_document_pages( list[PageRenderResult] One entry per page, ordered by page_index. """ - pages = list(range(1, page_count + 1)) - if not pages: + requested_pages = ( + sorted({page for page in pages if 1 <= page <= page_count}) + if pages is not None + else list(range(1, page_count + 1)) + ) + if not requested_pages: return [] # ── raw text (reuse caller's data if provided) ──────────────────── if page_texts is None: - page_texts = read_page_texts(pdf_path, pages, timeout=timeout) + page_texts = read_page_texts(pdf_path, requested_pages, timeout=timeout) # ── full-resolution PNGs ────────────────────────────────────────── if ctx is not None: pngs = render_pages( ctx, - pages, + requested_pages, folder_name="pages", prefix="page", dpi=dpi, @@ -114,7 +121,7 @@ def render_document_pages( ) pngs = render_pages( tmp_ctx, - pages, + requested_pages, folder_name="pages", prefix="page", dpi=dpi, @@ -133,7 +140,7 @@ def render_document_pages( # ── assemble results ────────────────────────────────────────────── results: list[PageRenderResult] = [] - for page in pages: + for page in requested_pages: feat = feature_map.get(page) results.append( PageRenderResult( diff --git a/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py b/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py index fc05f663..5cf19845 100644 --- a/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py +++ b/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py @@ -9,7 +9,7 @@ def test_render_direct_table_chunk_uses_summary_and_asset_url() -> None: { "chunk_id": "table-1", "chunk_type": "table", - "content": "tables/table-企业入驻信息表.html", + "content": "
SHOULD NOT LEAK
", "file_path": "tables/table-企业入驻信息表.html", "chunk_metadata": { "summary": "企业入驻信息登记模板", @@ -25,6 +25,7 @@ def test_render_direct_table_chunk_uses_summary_and_asset_url() -> None: assert "[Table: http://localhost:4566/table.html?signature=test]" in rendered assert "企业入驻信息登记模板" in rendered assert "企业信息;入驻管理" in rendered + assert "SHOULD NOT LEAK" not in rendered assert " None: { "chunk_id": "table-1", "chunk_type": "table", - "content": "tables/table-企业入驻信息表.html", + "content": "
企业名称
", "file_path": "tables/table-企业入驻信息表.html", "source_chunk_path": "企业信息汇总260509 (1).xlsx/企业批量录入", "chunk_metadata": { @@ -80,7 +81,7 @@ def test_render_connected_table_chunk_includes_asset_url() -> None: { "chunk_id": "table-1", "chunk_type": "table", - "content": "tables/table-1.html", + "content": "
SHOULD NOT LEAK
", "file_path": "tables/table-1.html", "chunk_metadata": {"summary": "A 表摘要"}, }, @@ -92,10 +93,11 @@ def test_render_connected_table_chunk_includes_asset_url() -> None: rendered = "\n".join(parts) assert "[Table: http://localhost:4566/table-1.html?signature=test]" in rendered assert "A 表摘要" in rendered + assert "SHOULD NOT LEAK" not in rendered assert " None: +def test_render_page_chunk_uses_summary_page_nums_and_page_pdf() -> None: parts: list[str] = [] render_leaf_chunks( parts, @@ -107,22 +109,20 @@ def test_render_page_chunk_uses_summary_page_nums_and_page_images() -> None: "chunk_metadata": { "summary": "制度标准总则摘要", "page_nums": [225, 226], - "page_image_uris": ["pages/page-225.png", "pages/page-226.png"], }, } ], " ", asset_lookup={ - "page-node-1": [ - "http://localhost:4566/page-225.png?signature=test", - "http://localhost:4566/page-226.png?signature=test", - ] + "page-node-1": "http://localhost:4566/page_pdfs/225-226.pdf?signature=test" }, ) rendered = "\n".join(parts) assert "Pages: 225, 226" in rendered assert "制度标准总则摘要" in rendered - assert "Page image 1: http://localhost:4566/page-225.png?signature=test" in rendered - assert "Page image 2: http://localhost:4566/page-226.png?signature=test" in rendered + assert ( + "Page PDF (pages 225-226): " + "http://localhost:4566/page_pdfs/225-226.pdf?signature=test" + ) in rendered assert "RAW OCR SHOULD NOT LEAK" not in rendered diff --git a/apps/worker/tests/contract/test_document_agent_budget_contract.py b/apps/worker/tests/contract/test_document_agent_budget_contract.py index db111193..036f7ae8 100644 --- a/apps/worker/tests/contract/test_document_agent_budget_contract.py +++ b/apps/worker/tests/contract/test_document_agent_budget_contract.py @@ -74,9 +74,7 @@ def test_dataframe_converter_accepts_page_chunks_with_extra_metadata() -> None: "connectto": "", "addtime": "2026-06-11 00:00:00", "page_nums": "1,2", - "extra_metadata": { - "page_image_uris": ["pages/page_page_1.png"], - }, + "extra_metadata": {}, } ] ) @@ -85,7 +83,7 @@ def test_dataframe_converter_accepts_page_chunks_with_extra_metadata() -> None: assert chunks[0]["type"] == "page" assert chunks[0]["metadata"]["page_nums"] == [1, 2] - assert chunks[0]["metadata"]["page_image_uris"] == ["pages/page_page_1.png"] + assert "page_image_uris" not in chunks[0]["metadata"] def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: @@ -98,7 +96,6 @@ def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: "metadata": { "summary": "page summary", "page_nums": [231], - "page_image_uris": ["pages/page-231.png", "pages/page-232.png"], }, } ] @@ -111,7 +108,7 @@ def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: metadata = formatted[0]["metadata"] assert metadata["page_nums"] == [231] - assert metadata["page_image_uris"] == ["pages/page-231.png", "pages/page-232.png"] + assert "page_image_uris" not in metadata assert "granularity" not in metadata assert "page_indices" not in metadata assert "owned_pages" not in metadata diff --git a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py index 01b0ab5f..20efb3ae 100644 --- a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py +++ b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py @@ -153,7 +153,7 @@ def _fake_extract_table(*, asset, pdf_path, output_dir): assert sorted(assets_by_page) == [1] assert (tmp_path / "tables" / "table_page_1_1.html").exists() - assert (tmp_path / "asset_annotate" / "page_1.png").exists() + assert not (tmp_path / "asset_annotate" / "page_1.png").exists() assert not (tmp_path / "page_asset_bboxes.json").exists() 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 d8f4eb2d..76521963 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 @@ -94,7 +94,6 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: rows = build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, - image_uri_by_page={231: "pages/page-231.png", 232: "pages/page-232.png"}, image_path_by_page={}, kind_by_page={}, tag_by_page={ @@ -116,18 +115,13 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: leaf_a = by_path["demo.pdf/3 基本规定/3.1 职责"] assert leaf_a["page_nums"] == "231" assert leaf_a["content"] == "text-231" - assert leaf_a["extra_metadata"]["page_image_uris"] == ["pages/page-231.png"] - assert set(leaf_a["extra_metadata"]) == {"page_image_uris"} + assert leaf_a["extra_metadata"] == {} leaf_b = by_path["demo.pdf/3 基本规定/3.2 管理规定"] assert leaf_b["page_nums"] == "231,232" assert SAME_AS_PREFIX in leaf_b["content"] assert "text-232" in leaf_b["content"] - assert leaf_b["extra_metadata"]["page_image_uris"] == [ - "pages/page-231.png", - "pages/page-232.png", - ] - assert set(leaf_b["extra_metadata"]) == {"page_image_uris"} + assert leaf_b["extra_metadata"] == {} def test_build_node_rows_keeps_internal_section_body_pages() -> None: @@ -151,7 +145,6 @@ def test_build_node_rows_keeps_internal_section_body_pages() -> None: rows = build_node_rows( skeletons=[parent, child], raw_text_by_page={233: "parent body", 234: "child body"}, - image_uri_by_page={233: "pages/page-233.png", 234: "pages/page-234.png"}, image_path_by_page={}, kind_by_page={}, tag_by_page={ @@ -177,6 +170,37 @@ def test_build_node_rows_keeps_internal_section_body_pages() -> None: ) +def test_boundary_page_belongs_to_next_sibling_start() -> None: + coarse = SectionSkeleton( + section_path="demo.pdf/安全类/风险标准", + level=3, + start_page=225, + end_page=302, + title="风险标准", + parent_path="demo.pdf/安全类", + ) + next_sibling = SectionSkeleton( + section_path="demo.pdf/安全类/项目分类标准", + level=3, + start_page=302, + end_page=304, + title="项目分类标准", + parent_path="demo.pdf/安全类", + ) + + leaves = identify_leaf_nodes([coarse, next_sibling]) + views, page_owner = assign_pages_to_leaves( + leaves, + available_pages={301, 302, 303}, + ) + + by_title = {view.leaf.title: view for view in views} + assert by_title["风险标准"].pages == [301] + assert by_title["项目分类标准"].pages == [302, 303] + assert page_owner[301].title == "风险标准" + assert page_owner[302].title == "项目分类标准" + + def test_build_node_rows_uses_vlm_node_summary_with_boundary( monkeypatch, tmp_path ) -> None: @@ -204,7 +228,6 @@ def chat_completion_with_usage(self, **kwargs): rows = build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, - image_uri_by_page={231: "pages/page-231.png", 232: "pages/page-232.png"}, image_path_by_page={231: str(img), 232: str(img)}, kind_by_page={}, tag_by_page={ @@ -246,7 +269,6 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: rows = build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, - image_uri_by_page={231: "pages/page-231.png", 232: "pages/page-232.png"}, image_path_by_page={}, kind_by_page={}, tag_by_page={ @@ -264,6 +286,12 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: assert rows[0]["path"] == "tables/table_page_231_1.html" assert rows[0]["content"] == "tables/table_page_231_1.html" assert rows[0]["know_id"] == "asset_table_1" + assert "owner_hierarchy_path" not in rows[0]["extra_metadata"] + assert "related_hierarchy_paths" not in rows[0]["extra_metadata"] + assert "page_index" not in rows[0]["extra_metadata"] + assert "asset_kind" not in rows[0]["extra_metadata"] + assert "title" not in rows[0]["extra_metadata"] + assert "html_uri" not in rows[0]["extra_metadata"] by_path = {row["path"]: row for row in rows} owner = by_path["demo.pdf/3 基本规定/3.1 职责"] @@ -309,6 +337,13 @@ def test_page_connectto_normalizes_to_asset_chunk_id() -> None: chunks = dataframe_to_chunks(df) + asset_chunk = next(chunk for chunk in chunks if chunk["type"] == "table") + assert asset_chunk["chunk_id"] == "asset_table_1" + assert "know_id" not in asset_chunk + assert "keywords" not in asset_chunk + assert "summary" not in asset_chunk + assert "tokens" not in asset_chunk + page_chunk = next(chunk for chunk in chunks if chunk["type"] == "page") assert page_chunk["metadata"]["connect_to"] == [ { diff --git a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py index 046ba891..acd2b5a4 100644 --- a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -1,6 +1,8 @@ from __future__ import annotations import os +import shutil +from pathlib import Path import pytest @@ -13,7 +15,7 @@ from shared.services.retrieval.hydration.assets import ( # noqa: E402 build_retrieval_asset_url_map, - enrich_rows_with_retrieval_asset_urls, + enrich_rows_with_retrieval_asset_url, ) from shared.services.retrieval.hydration.result_assembly import ( # noqa: E402 assemble_retrieval_results, @@ -28,6 +30,7 @@ resolve_workflow_references, ) from shared.services.storage.result_storage import JobResultStorage # noqa: E402 +from shared.services.storage.page_pdf_crop import crop_source_pdf_pages # noqa: E402 def test_data_type_one_allows_page_and_page_content_enters_search_text() -> None: @@ -71,6 +74,8 @@ def test_table_search_text_uses_summary_keywords_and_caption_not_content() -> No assert "批量录入表" in content_lexical_text assert "企业批量录入" in term_search_text assert "tables/table-1.html" not in content_search_text + assert "tables/table-1.html" not in content_lexical_text + assert "tables/table-1.html" not in term_search_text @pytest.mark.asyncio @@ -82,7 +87,7 @@ async def test_page_result_assembly_uses_summary_not_raw_content() -> None: "content": "RAW OCR SHOULD NOT LEAK", "chunk_metadata": { "summary": "制度标准总则摘要", - "page_image_uris": ["pages/page-225.png"], + "page_nums": [225], }, } ] @@ -142,30 +147,12 @@ async def test_table_result_assembly_uses_summary_not_html() -> None: @pytest.mark.asyncio -async def test_page_asset_urls_are_generated_from_page_image_uris(monkeypatch) -> None: - class FakeResultStorage: - def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: - if not artifact_ref: - return None - normalized = artifact_ref.strip().replace("\\", "/").lstrip("/") - root = normalized.split("/", 1)[0] - if root not in {"images", "tables", "pages"}: - return None - return normalized - - def generate_artifact_url( - self, - *, - job_id: str, - artifact_ref: str, - expires_in: int = 3600, - ) -> str | None: - del expires_in - return f"https://assets.example.com/{job_id}/{artifact_ref}" - +async def test_page_asset_url_is_generated_from_page_nums(monkeypatch) -> None: monkeypatch.setattr( - "shared.services.retrieval.hydration.assets.get_result_storage", - lambda: FakeResultStorage(), + "shared.services.retrieval.hydration.assets.crop_source_pdf_pages", + lambda *, job_id, pages: ( + f"https://assets.example.com/{job_id}/page_pdfs/{'-'.join(map(str, pages))}.pdf" + ), ) rows = [ @@ -174,34 +161,32 @@ def generate_artifact_url( "chunk_type": "page", "job_id": "job-1", "chunk_metadata": { - "page_image_uris": [ - "pages/page-225.png", - "pages/page-226.png", - "../images/not-allowed.png", - "pages/page-225.png", - ] + "page_nums": [225, 226, 225], }, } ] - enriched = await enrich_rows_with_retrieval_asset_urls( + enriched = await enrich_rows_with_retrieval_asset_url( rows, log_context="contract", ) url_map = await build_retrieval_asset_url_map(rows, log_context="contract") - assert enriched[0]["asset_urls"] == [ - "https://assets.example.com/job-1/pages/page-225.png", - "https://assets.example.com/job-1/pages/page-226.png", - ] - assert url_map["page-node-1"] == enriched[0]["asset_urls"] + assert enriched[0]["asset_url"] == ( + "https://assets.example.com/job-1/page_pdfs/225-226.pdf" + ) + assert "asset_urls" not in enriched[0] + assert url_map["page-node-1"] == enriched[0]["asset_url"] -def test_result_storage_allows_pages_artifact_refs() -> None: +def test_result_storage_allows_page_pdf_artifact_refs_not_page_pngs() -> None: storage = JobResultStorage(results_bucket="test-results") - assert storage.normalize_artifact_ref("pages/page-225.png") == "pages/page-225.png" - assert storage.normalize_artifact_ref("../pages/page-226.png") == "pages/page-226.png" + assert storage.normalize_artifact_ref("pages/page-225.png") is None + assert ( + storage.normalize_artifact_ref("page_pdfs/page-225.pdf") + == "page_pdfs/page-225.pdf" + ) def test_result_storage_upload_filters_to_referenced_artifacts(tmp_path) -> None: @@ -221,6 +206,7 @@ def generate_presigned_url(self, *args, **kwargs) -> str: result_dir = tmp_path / "result" (result_dir / "pages").mkdir(parents=True) (result_dir / "tables").mkdir() + (result_dir / "source.pdf").write_bytes(b"source") (result_dir / "pages" / "page-225.png").write_bytes(b"anchored") (result_dir / "pages" / "page-999.png").write_bytes(b"unanchored") (result_dir / "tables" / "table-1.html").write_text("
") @@ -238,18 +224,130 @@ def generate_presigned_url(self, *args, **kwargs) -> str: job_id="job-1", result_dir=str(result_dir), zip_file_path=str(zip_path), - artifact_refs={"pages/page-225.png", "tables/table-1.html"}, + artifact_refs={"source.pdf", "pages/page-225.png", "tables/table-1.html"}, ) - assert set(bundle.raw_files) == {"pages/page-225.png", "tables/table-1.html"} - assert "results/job-1/pages/page-225.png" in adapter.uploaded_keys + assert set(bundle.raw_files) == {"source.pdf", "tables/table-1.html"} + assert "results/job-1/source.pdf" in adapter.uploaded_keys assert "results/job-1/tables/table-1.html" in adapter.uploaded_keys + assert "results/job-1/pages/page-225.png" not in adapter.uploaded_keys assert "results/job-1/pages/page-999.png" not in adapter.uploaded_keys assert "results/job-1/debug.csv" not in adapter.uploaded_keys +def test_crop_source_pdf_pages_uploads_and_reuses_page_pdf_cache(tmp_path) -> None: + from pypdf import PdfReader, PdfWriter + + source_pdf = tmp_path / "source.pdf" + writer = PdfWriter() + writer.add_blank_page(width=72, height=72) + writer.add_blank_page(width=72, height=72) + writer.add_blank_page(width=72, height=72) + with source_pdf.open("wb") as file_obj: + writer.write(file_obj) + + class FakeResultStorage: + def __init__(self) -> None: + self.files: dict[str, Path] = {"source.pdf": source_pdf} + self.upload_count = 0 + + def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: + del job_id + return relative_path in self.files + + def download_raw_to_temp( + self, + *, + job_id: str, + relative_path: str, + suffix: str, + temp_dir: str, + ) -> str: + del job_id, suffix + target = Path(temp_dir) / Path(relative_path).name + shutil.copyfile(self.files[relative_path], target) + return str(target) + + def upload_raw_file( + self, + *, + job_id: str, + relative_path: str, + local_file_path: str, + ) -> None: + del job_id + target = tmp_path / relative_path.replace("/", "_") + shutil.copyfile(local_file_path, target) + self.files[relative_path] = target + self.upload_count += 1 + + def generate_artifact_url( + self, + *, + job_id: str, + artifact_ref: str, + expires_in: int = 3600, + ) -> str: + del expires_in + return f"https://assets.example.com/{job_id}/{artifact_ref}" + + storage = FakeResultStorage() + + url = crop_source_pdf_pages( + job_id="job-1", + pages=[2, 1, 2], + storage=storage, # type: ignore[arg-type] + temp_dir=str(tmp_path), + ) + assert url and "/page_pdfs/" in url + page_pdf_refs = [ref for ref in storage.files if ref.startswith("page_pdfs/")] + assert len(page_pdf_refs) == 1 + assert len(PdfReader(str(storage.files[page_pdf_refs[0]])).pages) == 2 + + cached_url = crop_source_pdf_pages( + job_id="job-1", + pages=[1, 2], + storage=storage, # type: ignore[arg-type] + temp_dir=str(tmp_path), + ) + assert cached_url == url + assert storage.upload_count == 1 + + +def test_crop_source_pdf_pages_returns_none_when_source_pdf_is_missing(tmp_path) -> None: + class FakeResultStorage: + def __init__(self) -> None: + self.upload_count = 0 + + def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: + del job_id, relative_path + return False + + def download_raw_to_temp(self, **_kwargs) -> str: + raise AssertionError("source download should not be attempted") + + def upload_raw_file(self, **_kwargs) -> None: + self.upload_count += 1 + + def generate_artifact_url(self, **_kwargs) -> str: + raise AssertionError("URL generation should not be attempted") + + storage = FakeResultStorage() + + assert ( + crop_source_pdf_pages( + job_id="job-missing-source", + pages=[1], + storage=storage, # type: ignore[arg-type] + temp_dir=str(tmp_path), + ) + is None + ) + assert storage.upload_count == 0 + + @pytest.mark.asyncio -async def test_referenced_chunks_get_page_asset_urls_from_hydrated_rows( +async def test_referenced_chunks_get_page_asset_url_from_hydrated_rows( monkeypatch, ) -> None: async def fake_hydrate_referenced_chunk_rows(**_kwargs): @@ -260,20 +358,20 @@ async def fake_hydrate_referenced_chunk_rows(**_kwargs): "chunk_type": "page", "section_path": "安全类 / 1 总则", "file_path": None, - "chunk_metadata": {"page_image_uris": ["pages/page-225.png"]}, + "chunk_metadata": {"page_nums": [225]}, "job_id": "job-1", } ] - async def fake_enrich_referenced_chunks_with_asset_urls(rows): + async def fake_enrich_referenced_chunks_with_asset_url(rows): enriched = [] for row in rows: enriched.append( { **row, - "asset_urls": [ - "https://assets.example.com/job-1/pages/page-225.png" - ], + "asset_url": ( + "https://assets.example.com/job-1/page_pdfs/page-225.pdf" + ), } ) return enriched @@ -283,8 +381,8 @@ async def fake_enrich_referenced_chunks_with_asset_urls(rows): fake_hydrate_referenced_chunk_rows, ) monkeypatch.setattr( - "shared.services.retrieval.execution.reference_resolver.enrich_referenced_chunks_with_asset_urls", - fake_enrich_referenced_chunks_with_asset_urls, + "shared.services.retrieval.execution.reference_resolver.enrich_referenced_chunks_with_asset_url", + fake_enrich_referenced_chunks_with_asset_url, ) resolved = await resolve_workflow_references( @@ -301,6 +399,7 @@ async def fake_enrich_referenced_chunks_with_asset_urls(rows): ], ) - assert resolved.refs[0]["asset_urls"] == [ - "https://assets.example.com/job-1/pages/page-225.png" - ] + assert resolved.refs[0]["asset_url"] == ( + "https://assets.example.com/job-1/page_pdfs/page-225.pdf" + ) + assert "asset_urls" not in resolved.refs[0] diff --git a/apps/worker/tests/contract/test_table_asset_schema_contract.py b/apps/worker/tests/contract/test_table_asset_schema_contract.py new file mode 100644 index 00000000..41962377 --- /dev/null +++ b/apps/worker/tests/contract/test_table_asset_schema_contract.py @@ -0,0 +1,49 @@ +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") + +from app.services.document_parser.assets.inline_asset import build_table_asset_row # noqa: E402 +from app.services.document_parser.tables.table_asset_writer import ( # noqa: E402 + TableAssetInput, + write_table_asset, +) + + +def test_table_asset_writer_stores_reference_not_html_content(tmp_path) -> None: + html = "
SHOULD NOT LEAK
" + + row = write_table_asset( + TableAssetInput( + html=html, + output_dir=str(tmp_path), + table_name="table-1.html", + summary="table summary", + keywords="name;value", + know_id="table-1", + addtime="2026-06-24 00:00:00", + ) + ) + + assert row.content == "tables/table-1.html" + assert " None: + row = build_table_asset_row( + relative_path="tables/table-2.html", + summary="table summary", + keywords="name;value", + know_id="table-2", + addtime="2026-06-24 00:00:00", + ) + + assert row.content == "tables/table-2.html" + assert "=24.11.1", diff --git a/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py b/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py index c2e5c9e7..49a7a119 100644 --- a/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py +++ b/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py @@ -60,10 +60,6 @@ class ChunkPayload(TypedDict): path: str metadata: ChunkMetadata order: int - know_id: str - keywords: list[str] - summary: str - tokens: list[str] def _is_missing(value: object) -> bool: @@ -331,10 +327,6 @@ def dataframe_to_chunks(df: _ParserDataFrame | None) -> list[Dict[str, JsonValue "path": path, "metadata": metadata, "order": index, - "know_id": str(know_id), - "keywords": metadata["keywords"], - "summary": metadata["summary"], - "tokens": metadata["tokens"], } ) diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py index 67ac6f6b..5d31cf86 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py @@ -5,7 +5,7 @@ from shared.services.retrieval.agentic.core.types import DocTreeNode -AssetLookupValue = str | list[str] +AssetLookupValue = str def render_unified_doc_tree( @@ -158,13 +158,11 @@ def render_leaf_chunks( rendered_ids.add(target_id) if target_type == "table": - table_html = str(target.get("content", "")).strip() file_path = target.get("file_path") or "" - asset_url = _first_asset_url((asset_lookup or {}).get(target_id, "")) if target_id else "" + asset_url = _lookup_asset_url(asset_lookup, target_id) display_ref = asset_url or file_path table_lines = render_table_chunk_lines( target, - table_html=table_html, display_ref=display_ref, ) content = content.replace(ref_str, "\n" + "\n".join(table_lines) + "\n") @@ -173,7 +171,7 @@ def render_leaf_chunks( image_description = str(target.get("content", "")).strip() if ref_str in image_description: image_description = image_description.replace(ref_str, "").strip() - asset_url = _first_asset_url((asset_lookup or {}).get(target_id, "")) if target_id else "" + asset_url = _lookup_asset_url(asset_lookup, target_id) display_ref = asset_url or file_path if display_ref: content = content.replace(ref_str, f"\n[Image: {display_ref}]\n{image_description}\n") @@ -195,7 +193,7 @@ def render_leaf_chunks( if chunk_type == "image": file_path = chunk.get("file_path") or "" image_description = str(chunk.get("content", "")).strip() - asset_url = _first_asset_url((asset_lookup or {}).get(chunk_id, "")) if chunk_id else "" + asset_url = _lookup_asset_url(asset_lookup, chunk_id) display_ref = asset_url or file_path if display_ref: parts.append(f"{indent}┈ [Image: {display_ref}]") @@ -204,13 +202,11 @@ def render_leaf_chunks( if line.strip(): parts.append(f"{indent}┈ {line}") elif chunk_type == "table": - table_html = str(chunk.get("content", "")).strip() file_path = chunk.get("file_path") or "" - asset_url = _first_asset_url((asset_lookup or {}).get(chunk_id, "")) if chunk_id else "" + asset_url = _lookup_asset_url(asset_lookup, chunk_id) display_ref = asset_url or file_path for line in render_table_chunk_lines( chunk, - table_html=table_html, display_ref=display_ref, ): if line.strip(): @@ -238,22 +234,16 @@ def render_page_chunk_lines( elif not page_nums: parts.append(f"{indent}┈ [Page]") - urls = _asset_urls_for_chunk(chunk, asset_lookup=asset_lookup) - if urls: - for index, url in enumerate(urls, start=1): - parts.append(f"{indent}┈ [Page image {index}: {url}]") - elif page_nums: - for page in page_nums: - parts.append(f"{indent}┈ [Page image: page {page}]") + url = _asset_url_for_chunk(chunk, asset_lookup=asset_lookup) + if url: + parts.append(f"{indent}┈ [Page PDF ({_format_page_range(page_nums)}): {url}]") def render_table_chunk_lines( chunk: dict[str, Any], *, - table_html: str, display_ref: str, ) -> list[str]: - del table_html header = f"[Table: {display_ref}]" if display_ref else "[Table]" lines = [header] table_path = chunk.get("source_chunk_path") or chunk.get("section_path") @@ -286,27 +276,42 @@ def render_table_chunk_lines( return lines -def _asset_urls_for_chunk( +def _asset_url_for_chunk( chunk: dict[str, Any], *, asset_lookup: dict[str, AssetLookupValue] | None, -) -> list[str]: +) -> str: chunk_id = str(chunk.get("chunk_id") or "").strip() value = (asset_lookup or {}).get(chunk_id, "") if chunk_id else "" - if isinstance(value, list): - return [str(url).strip() for url in value if str(url).strip()] - url = str(value or "").strip() - return [url] if url else [] + return str(value or "").strip() -def _first_asset_url(value: object) -> str: - if isinstance(value, list): - for item in value: - url = str(item or "").strip() - if url: - return url +def _lookup_asset_url( + asset_lookup: dict[str, AssetLookupValue] | None, + chunk_id: str, +) -> str: + if not chunk_id: return "" - return str(value or "").strip() + return str((asset_lookup or {}).get(chunk_id, "") or "").strip() + + +def _format_page_range(page_nums: list[int]) -> str: + pages = sorted(set(page_nums)) + if not pages: + return "pages unknown" + if len(pages) == 1: + return f"page {pages[0]}" + ranges: list[str] = [] + start = pages[0] + prev = pages[0] + for page in pages[1:]: + if page == prev + 1: + prev = page + continue + ranges.append(str(start) if start == prev else f"{start}-{prev}") + start = prev = page + ranges.append(str(start) if start == prev else f"{start}-{prev}") + return f"pages {', '.join(ranges)}" def _coerce_page_nums(value: object) -> list[int]: diff --git a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index 534c016f..d5e2c358 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -7,7 +7,7 @@ from shared.services.retrieval.hydration.reference import hydrate_referenced_chunk_rows from shared.services.retrieval.execution.response_projection import ( - enrich_referenced_chunks_with_asset_urls, + enrich_referenced_chunks_with_asset_url, ) from shared.services.retrieval.hydration.row_utils import build_reference_lookup_key @@ -34,9 +34,9 @@ async def resolve_workflow_references( score_by_chunk_id=score_by_chunk_id, ) resolved = _select_matching_references(refs, hydrated_rows) - enriched_rows = await enrich_referenced_chunks_with_asset_urls(resolved.rows) + enriched_rows = await enrich_referenced_chunks_with_asset_url(resolved.rows) return ResolvedWorkflowReferences( - refs=_merge_reference_asset_urls(resolved.refs, enriched_rows), + refs=_merge_reference_asset_url(resolved.refs, enriched_rows), rows=resolved.rows, ) @@ -95,7 +95,7 @@ def _row_key(row: dict[str, Any]) -> tuple[str, str, str, str]: ) -def _merge_reference_asset_urls( +def _merge_reference_asset_url( refs: list[dict[str, Any]], rows: list[dict[str, Any]], ) -> list[dict[str, Any]]: @@ -114,9 +114,8 @@ def _merge_reference_asset_urls( if row is None: row = _find_matching_row_for_ref(ref, rows) if row is not None: - for field in ("asset_url", "asset_urls"): - if row.get(field): - merged[field] = row[field] + if row.get("asset_url"): + merged["asset_url"] = row["asset_url"] merged_refs.append(merged) return merged_refs diff --git a/packages/shared-python/shared/services/retrieval/execution/response_projection.py b/packages/shared-python/shared/services/retrieval/execution/response_projection.py index 80cf15aa..00136340 100644 --- a/packages/shared-python/shared/services/retrieval/execution/response_projection.py +++ b/packages/shared-python/shared/services/retrieval/execution/response_projection.py @@ -2,7 +2,7 @@ from typing import Any -from shared.services.retrieval.hydration.assets import enrich_rows_with_retrieval_asset_urls +from shared.services.retrieval.hydration.assets import enrich_rows_with_retrieval_asset_url from shared.services.retrieval.hydration.row_utils import ( PUBLIC_RESULT_FIELDS, PUBLIC_SOURCE_FIELDS, @@ -23,8 +23,8 @@ def to_public_source(row: dict[str, Any]) -> dict[str, Any]: return {field: row.get(field) for field in PUBLIC_SOURCE_FIELDS} -async def enrich_referenced_chunks_with_asset_urls(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: - return await enrich_rows_with_retrieval_asset_urls( +async def enrich_referenced_chunks_with_asset_url(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + return await enrich_rows_with_retrieval_asset_url( refs, log_context='agentic referenced chunk', ) @@ -48,7 +48,7 @@ async def project_public_retrieval_response(response: dict[str, Any]) -> dict[st if response.get('decision_trace') is not None: public_response['decision_trace'] = response['decision_trace'] - projected_rows = await enrich_rows_with_retrieval_asset_urls( + projected_rows = await enrich_rows_with_retrieval_asset_url( response.get('results', []), log_context='retrieval result', ) diff --git a/packages/shared-python/shared/services/retrieval/hydration/assets.py b/packages/shared-python/shared/services/retrieval/hydration/assets.py index 352aa2dc..8d305512 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/assets.py +++ b/packages/shared-python/shared/services/retrieval/hydration/assets.py @@ -1,13 +1,16 @@ from __future__ import annotations +import asyncio from typing import Any from loguru import logger from shared.services.retrieval.hydration.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type +from shared.services.storage.page_pdf_crop import crop_source_pdf_pages from shared.services.storage.result_storage import get_result_storage -AssetUrlValue = str | list[str] +AssetUrlValue = str +PagePdfRequestKey = tuple[str, tuple[int, ...]] def _normalize_artifact_ref(asset_ref: object) -> str | None: @@ -38,31 +41,39 @@ def _resolve_asset_request(row: dict[str, Any]) -> tuple[str, str] | None: return job_id, artifact_ref -def _resolve_page_asset_requests(row: dict[str, Any]) -> list[tuple[str, str]]: +def _coerce_page_nums(value: object) -> list[int]: + if isinstance(value, list): + raw_values = value + elif value is None: + raw_values = [] + else: + raw_values = str(value).split(",") + + pages: list[int] = [] + for item in raw_values: + try: + pages.append(int(str(item).strip())) + except (TypeError, ValueError): + continue + return pages + + +def _resolve_page_pdf_request(row: dict[str, Any]) -> PagePdfRequestKey | None: job_id = str(row.get("job_id") or "").strip() if not job_id or not _is_page_row(row): - return [] + return None metadata = row.get("chunk_metadata") or row.get("metadata") or {} if not isinstance(metadata, dict): - return [] - - requests: list[tuple[str, str]] = [] - seen_refs: set[str] = set() - raw_refs = metadata.get("page_image_uris") or [] - if not isinstance(raw_refs, list): - return [] - for raw_ref in raw_refs: - artifact_ref = _normalize_artifact_ref(raw_ref) - if ( - artifact_ref is None - or not artifact_ref.startswith("pages/") - or artifact_ref in seen_refs - ): - continue - seen_refs.add(artifact_ref) - requests.append((job_id, artifact_ref)) - return requests + return None + + pages = _coerce_page_nums(metadata.get("page_nums") or row.get("page_nums")) + if not pages: + return None + normalized_pages = tuple(sorted({page for page in pages if page > 0})) + if not normalized_pages: + return None + return job_id, normalized_pages async def _generate_retrieval_asset_url( @@ -85,35 +96,58 @@ async def _generate_retrieval_asset_url( return None -async def _generate_retrieval_asset_urls( +async def _generate_page_pdf_asset_url_for_request( + request: PagePdfRequestKey, *, - row: dict[str, Any], log_context: str, -) -> list[str]: - requests = _resolve_page_asset_requests(row) - if not requests: - return [] +) -> str | None: + job_id, pages = request + try: + return await asyncio.to_thread( + crop_source_pdf_pages, + job_id=job_id, + pages=list(pages), + ) + except Exception as exc: + logger.warning(f"Failed to generate {log_context} page PDF URL (ignored): {exc}") + return None - urls: list[str] = [] - for job_id, artifact_ref in requests: - try: - url = get_result_storage().generate_artifact_url( - job_id=job_id, - artifact_ref=artifact_ref, + +async def _build_page_pdf_url_lookup( + rows: list[dict[str, Any]], + *, + log_context: str, +) -> dict[PagePdfRequestKey, str]: + requests = { + request + for row in rows + if (request := _resolve_page_pdf_request(row)) is not None + } + if not requests: + return {} + sorted_requests = sorted(requests) + urls = await asyncio.gather( + *( + _generate_page_pdf_asset_url_for_request( + request, + log_context=log_context, ) - except Exception as exc: - logger.warning(f"Failed to generate {log_context} page asset URL (ignored): {exc}") - continue - if url: - urls.append(url) - return urls + for request in sorted_requests + ) + ) + return { + request: url + for request, url in zip(sorted_requests, urls, strict=True) + if url + } -async def enrich_rows_with_retrieval_asset_urls( +async def enrich_rows_with_retrieval_asset_url( rows: list[dict[str, Any]], *, log_context: str, ) -> list[dict[str, Any]]: + page_pdf_urls = await _build_page_pdf_url_lookup(rows, log_context=log_context) enriched_rows: list[dict[str, Any]] = [] for row in rows: enriched = dict(row) @@ -123,12 +157,11 @@ async def enrich_rows_with_retrieval_asset_urls( ) if asset_url: enriched["asset_url"] = asset_url - asset_urls = await _generate_retrieval_asset_urls( - row=row, - log_context=log_context, - ) - if asset_urls: - enriched["asset_urls"] = asset_urls + page_pdf_url = None + if page_request := _resolve_page_pdf_request(row): + page_pdf_url = page_pdf_urls.get(page_request) + if page_pdf_url: + enriched["asset_url"] = page_pdf_url enriched_rows.append(enriched) return enriched_rows @@ -138,18 +171,18 @@ async def build_retrieval_asset_url_map( *, log_context: str, ) -> dict[str, AssetUrlValue]: + page_pdf_urls = await _build_page_pdf_url_lookup(rows, log_context=log_context) url_map: dict[str, AssetUrlValue] = {} for row in rows: chunk_id = str(row.get("chunk_id") or "").strip() if not chunk_id: continue - asset_urls = await _generate_retrieval_asset_urls( - row=row, - log_context=log_context, - ) - if asset_urls: - url_map[chunk_id] = asset_urls + page_pdf_url = None + if page_request := _resolve_page_pdf_request(row): + page_pdf_url = page_pdf_urls.get(page_request) + if page_pdf_url: + url_map[chunk_id] = page_pdf_url continue asset_url = await _generate_retrieval_asset_url( diff --git a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py index e204ce26..d163836c 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py +++ b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py @@ -7,7 +7,7 @@ MEDIA_CHUNK_TYPES = {'image', 'table'} PUBLIC_RESULT_FIELDS = { - 'chunk_type', 'content', 'score', 'asset_url', 'asset_urls', + 'chunk_type', 'content', 'score', 'asset_url', } PUBLIC_SOURCE_FIELDS = { 'document_id', 'source_file_name', 'section_path', @@ -15,7 +15,7 @@ ReferenceLookupKey = tuple[str, str, str, str] -_PATH_REF_RE = re.compile(r'\[(?:images|tables|pages)/[^\]\n]+\]') +_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]') def clean_content(content: str) -> str: diff --git a/packages/shared-python/shared/services/storage/page_pdf_crop.py b/packages/shared-python/shared/services/storage/page_pdf_crop.py new file mode 100644 index 00000000..c3de8f51 --- /dev/null +++ b/packages/shared-python/shared/services/storage/page_pdf_crop.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import hashlib +import os +import tempfile +from pathlib import Path + +from loguru import logger + +from shared.services.storage.result_storage import JobResultStorage + +_SOURCE_PDF_REF = "source.pdf" +_PAGE_PDF_DIR = "page_pdfs" + + +def crop_source_pdf_pages( + *, + job_id: str, + pages: list[int], + storage: JobResultStorage | None = None, + expires_in: int = 3600, + temp_dir: str | None = None, +) -> str | None: + normalized_pages = _normalize_pages(pages) + if not job_id or not normalized_pages: + return None + + result_storage = storage or JobResultStorage() + cache_ref = _cache_ref(normalized_pages) + if result_storage.verify_raw_exists(job_id=job_id, relative_path=cache_ref): + return result_storage.generate_artifact_url( + job_id=job_id, + artifact_ref=cache_ref, + expires_in=expires_in, + ) + + if not result_storage.verify_raw_exists(job_id=job_id, relative_path=_SOURCE_PDF_REF): + logger.warning("[page_pdf_crop] source.pdf missing for job_id={}", job_id) + return None + + work_dir = temp_dir or tempfile.gettempdir() + try: + with tempfile.TemporaryDirectory(dir=work_dir) as local_dir: + source_path = result_storage.download_raw_to_temp( + job_id=job_id, + relative_path=_SOURCE_PDF_REF, + suffix=".pdf", + temp_dir=local_dir, + ) + cropped_path = os.path.join(local_dir, Path(cache_ref).name) + _write_cropped_pdf( + source_path=source_path, + output_path=cropped_path, + pages=normalized_pages, + ) + result_storage.upload_raw_file( + job_id=job_id, + relative_path=cache_ref, + local_file_path=cropped_path, + ) + except Exception as exc: + logger.warning( + "[page_pdf_crop] failed to crop source PDF for job_id={}, pages={}: {}", + job_id, + normalized_pages, + exc, + ) + return None + + return result_storage.generate_artifact_url( + job_id=job_id, + artifact_ref=cache_ref, + expires_in=expires_in, + ) + + +def _normalize_pages(pages: list[int]) -> list[int]: + normalized: set[int] = set() + for page in pages: + try: + value = int(page) + except (TypeError, ValueError): + continue + if value > 0: + normalized.add(value) + return sorted(normalized) + + +def _cache_ref(pages: list[int]) -> str: + digest = hashlib.sha1(",".join(str(page) for page in pages).encode("utf-8")).hexdigest() + return f"{_PAGE_PDF_DIR}/{digest}.pdf" + + +def _write_cropped_pdf(*, source_path: str, output_path: str, pages: list[int]) -> None: + from pypdf import PdfReader, PdfWriter + + reader = PdfReader(source_path) + writer = PdfWriter() + page_count = len(reader.pages) + selected_pages = [page for page in pages if 1 <= page <= page_count] + if not selected_pages: + raise ValueError("No requested pages exist in source PDF") + for page in selected_pages: + writer.add_page(reader.pages[page - 1]) + with open(output_path, "wb") as file_obj: + writer.write(file_obj) diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 614e1dd8..529dde30 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -13,7 +13,8 @@ _EXCLUDED_FILE_NAMES = {".DS_Store", "Thumbs.db"} _EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"} -_CLIENT_ARTIFACT_DIRS = {"images", "tables", "pages"} +_CLIENT_ARTIFACT_DIRS = {"images", "tables", "page_pdfs"} +_INTERNAL_RAW_FILES = {"source.pdf"} @dataclass(frozen=True) @@ -127,6 +128,39 @@ def generate_url(self, *, storage_key: str, expires_in: int = 3600) -> str | Non expires_in=expires_in, )["download_url"] + def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: + key = self.build_raw_key(job_id=job_id, relative_path=relative_path) + result = self._job_file_storage.verify_exists(key, bucket=self.results_bucket) + return bool(result.get("exists")) + + def upload_raw_file( + self, + *, + job_id: str, + relative_path: str, + local_file_path: str, + ) -> None: + self._job_file_storage.upload_local_file( + local_file_path, + self.build_raw_key(job_id=job_id, relative_path=relative_path), + bucket=self.results_bucket, + ) + + def download_raw_to_temp( + self, + *, + job_id: str, + relative_path: str, + suffix: str, + temp_dir: str, + ) -> str: + return self._job_file_storage.download_to_temp( + self.build_raw_key(job_id=job_id, relative_path=relative_path), + suffix=suffix, + temp_dir=temp_dir, + bucket=self.results_bucket, + ) + def generate_artifact_url( self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 ) -> str | None: @@ -171,10 +205,18 @@ def _normalize_artifact_refs(self, artifact_refs: set[str] | None) -> set[str] | normalized_refs = { normalized for ref in artifact_refs - if (normalized := self.normalize_artifact_ref(ref)) is not None + if (normalized := self._normalize_upload_artifact_ref(ref)) is not None } return normalized_refs + def _normalize_upload_artifact_ref(self, artifact_ref: str | None) -> str | None: + normalized = self._normalize_raw_relative_path(artifact_ref) + if not normalized: + return None + if normalized in _INTERNAL_RAW_FILES: + return normalized + return self.normalize_artifact_ref(normalized) + def _is_excluded_file(self, file_name: str) -> bool: return file_name in _EXCLUDED_FILE_NAMES or file_name.startswith(".") diff --git a/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py index 37343884..09a090d4 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -106,9 +106,6 @@ def format_chunks( ) metadata["tokens"] = [] elif chunk_type == "page": - # Page-memory page chunks preserve node/whole-doc navigation - # metadata from extra_metadata. - _merge_page_chunk_metadata(metadata, existing_metadata) metadata["keywords"] = existing_metadata.get("keywords") or [] metadata["tokens"] = [] @@ -130,18 +127,6 @@ def _normalize_chunk_type(value: Any) -> str: return raw_type.split("\n", 1)[0].lower() -_PAGE_CHUNK_METADATA_KEYS = ("page_image_uris",) - - -def _merge_page_chunk_metadata( - metadata: dict[str, Any], - existing_metadata: dict[str, Any], -) -> None: - for key in _PAGE_CHUNK_METADATA_KEYS: - if key in existing_metadata: - metadata[key] = existing_metadata[key] - - def _base_chunk_metadata( existing_metadata: dict[str, Any], chunk: dict[str, Any], diff --git a/packages/shared-python/shared/services/storage/zip_package_writer.py b/packages/shared-python/shared/services/storage/zip_package_writer.py index 42c78532..f3ec2337 100644 --- a/packages/shared-python/shared/services/storage/zip_package_writer.py +++ b/packages/shared-python/shared/services/storage/zip_package_writer.py @@ -57,6 +57,21 @@ def write(self, request: ZipPackageWriteRequest) -> ZipPackageArtifact: "toc_hierarchies.json", ): logger.info("Added toc_hierarchies.json to ZIP") + if self._write_optional_json_file( + zip_file, + request.add_dir, + "_doc_agent/trace.json", + "debug/trace.json", + compact_trace=True, + ): + logger.info("Added debug/trace.json to ZIP") + if self._write_optional_json_file( + zip_file, + request.add_dir, + "_doc_agent/anatomy_map.json", + "debug/anatomy_map.json", + ): + logger.info("Added debug/anatomy_map.json to ZIP") self._write_resource_files(zip_file, request.image_files, label="Image") self._write_resource_files(zip_file, request.table_files, label="Table") @@ -89,6 +104,32 @@ def _write_optional_file( zip_file.write(file_path, filename) return True + def _write_optional_json_file( + self, + zip_file: zipfile.ZipFile, + add_dir: str, + source_name: str, + zip_name: str, + *, + compact_trace: bool = False, + ) -> bool: + file_path = os.path.join(add_dir, source_name) + if not os.path.exists(file_path): + return False + try: + with open(file_path, encoding="utf-8") as file_obj: + payload = json.load(file_obj) + if compact_trace: + payload = _compact_trace_payload(payload) + zip_file.writestr( + zip_name, + json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8"), + ) + return True + except Exception as exc: + logger.warning(f"Failed to add {source_name} to ZIP: {exc}") + return False + def _write_resource_files( self, zip_file: zipfile.ZipFile, @@ -110,3 +151,23 @@ def _calculate_zip_checksum(zip_file_path: str) -> str: for byte_block in iter(lambda: file_obj.read(4096), b""): sha256_hash.update(byte_block) return sha256_hash.hexdigest().lower() + + +def _compact_trace_payload(payload: Any) -> Any: + if not isinstance(payload, dict): + return payload + compacted = dict(payload) + steps: list[Any] = [] + for step in compacted.get("steps") or []: + if not isinstance(step, dict): + steps.append(step) + continue + item = dict(step) + observation = item.get("observation") + if isinstance(observation, dict): + compact_observation = dict(observation) + compact_observation.pop("payload", None) + item["observation"] = compact_observation + steps.append(item) + compacted["steps"] = steps + return compacted diff --git a/page_memory_work_summary_20260624.md b/page_memory_work_summary_20260624.md new file mode 100644 index 00000000..7f7ed99d --- /dev/null +++ b/page_memory_work_summary_20260624.md @@ -0,0 +1,437 @@ +# Page-Memory 流程改造与测试进度汇总 + +日期：2026-06-24 +文档：`SJSYJ-SC-2024 企业制度汇编（上册）.pdf` +当前 debug 目录：`/Users/wuchengke/.knowhere/_debug_parse/SJSYJ-SC-2024 企业制度汇编（上册）.pdf/page_memory` + +## 1. 背景问题 + +这轮工作从“检查解析”对话中的异常开始： + +```text +The encrypted content Hand...run. could not be verified. +Reason: Encrypted content could not be decrypted or parsed. +traceid: f63bf5fe74f7bb1859be513dd7e54b22 +``` + +当时 debug 目录里只有： + +```text +page_memory_fine_hierarchy.json +``` + +但没有看到 PAGE-TAG 结果，也无法判断流程到底跑到哪一步。随后确认核心问题不是单个文件缺失，而是 page-memory 的中间产物、trace、scope 组织方式和生产链路不一致，导致 debug 不可控、不透明。 + +## 2. 已确认的目标流程 + +我们对齐后的 page-memory 生产/测试流程是： + +1. 先做 DOC_PROFILE，抽取文档页特征和 TOC hierarchy。 +2. 基于粗 TOC hierarchy 建立 coarse scopes。 +3. 对每个 coarse scope 构建 fine hierarchy。 +4. 只在 fine hierarchy 覆盖的页面范围内做 PAGE-TAG。 +5. 如开启图表资产提取，则同样在 hierarchy 覆盖范围内抽取资产。 +6. 最后汇总生成顶层 `hierarchy.json`、`page_tags.json`、`assets.json`，并继续支持后续 `chunks.json`、`doc_nav.json`、`manifest.json` 生成。 + +测试链路保持生产形状，但可以用 `--fat-only` 只选择最大的粗 scope 进行快速验证。 + +## 3. 已落地的主要改造 + +### 3.1 Trace 统一到顶层 + +已取消 `_doc_agent/trace.json` 的重复落盘，统一写到： + +```text +page_memory/trace.json +``` + +`_doc_agent/` 目前只保留 DOC_PROFILE/TOC 相关原始调试材料： + +```text +_doc_agent/anatomy_map.json +_doc_agent/toc_hierarchies.json +``` + +### 3.2 中间产物精简 + +当前 stop-at fine 的落盘结构已经收敛为： + +```text +page_memory/ + trace.json + hierarchy.json + page_tags.json + _doc_agent/ + anatomy_map.json + toc_hierarchies.json + scopes/ + p225-301/ + scope.json + fine_hierarchy.json + page_tags.json +``` + +不再输出大量重复的临时 JSON，例如旧的 `coarse_tag_scope.json`、`page_memory_fine_hierarchy.json`、`page_tags_pre_hierarchy.json` 等。 + +### 3.3 Scope 目录改为页码范围 + +scope 目录已从 hash/UUID 风格改为页码范围： + +```text +scopes/p225-301/ +``` + +这样 debug 时可以直接看出该 scope 覆盖的页域。 + +### 3.4 `hierarchy.json` 改为可读树形优先 + +顶层 `hierarchy.json` 和 scope 内 `fine_hierarchy.json` 都改成： + +```json +{ + "HIERARCHY": {}, + "nodes": [], + "stats": {} +} +``` + +其中 `HIERARCHY` 是第一字段，形态对齐最终 `manifest.json` 里的 `HIERARCHY` 字段，方便直接肉眼 debug。 + +scope 内的 `fine_hierarchy.json` 额外包含： + +```json +{ + "scope": {} +} +``` + +机器流程仍可继续使用 `nodes`。 + +### 3.5 页域字段精简 + +scope 和 trace 中不再同时记录 `page_ranges` 和完整 `pages` 列表。 + +当前约定： + +```json +{ + "document_page_count": 423, + "page_count": 77, + "page_ranges": [[225, 301]] +} +``` + +保留逐页 `page_index` 的地方仅限实体数据，例如： + +```text +page_tags.json +assets.json +``` + +因为这些文件本身就是逐页/逐资产记录。 + +## 4. 当前测试进度 + +本次从原始 PDF 重新开始测试： + +```text +/Users/wuchengke/Desktop/temp/test_docs/SJSYJ-SC-2024 企业制度汇编（上册）.pdf +``` + +执行目标： + +```text +从头跑到最大粗 scope 的 fine hierarchy +``` + +实际命令： + +```bash +uv run python apps/worker/scripts/debug_page_memory.py \ + --file '/Users/wuchengke/Desktop/temp/test_docs/SJSYJ-SC-2024 企业制度汇编（上册）.pdf' \ + --fat-only \ + --stop-at fine +``` + +结果：成功，最终状态为： + +```text +stopped_at_fine +``` + +### 4.1 DOC_PROFILE 结果 + +```text +page_count: 423 +toc_pages: [5, 6, 228] +native TOC TitleNode: 58 +native TOC leaf nodes: 44 +``` + +注意：C4 skeleton 阶段日志显示嵌入式目录页 `[228]` 被当前 global TOC 选择逻辑跳过： + +```text +embedded_toc_region_outside_front_cluster +``` + +这属于后续可优化点。 + +### 4.2 C4 Skeleton 定位 + +```text +skeleton_count: 25 +elapsed: 279.36s +``` + +残余定位阶段使用了多轮小窗口渲染 + VLM 确认，耗时较长。 + +### 4.3 最大粗 scope 选择 + +`--fat-only` 本次选中： + +```text +scope_id: p225-301 +page_ranges: [[225, 301]] +page_count: 77 +coarse skeletons before fine: 1 +``` + +### 4.4 Fine hierarchy 结果 + +title detection： + +```text +77 VLM calls +54 titles found +36 pages with observed_titles +``` + +fine hierarchy： + +```text +1 -> 52 skeletons +elapsed: 110.5s +``` + +最终 `hierarchy.json`： + +```json +{ + "stats": { + "node_count": 52, + "page_count": 77, + "page_ranges": [[225, 301]], + "max_depth": 6 + } +} +``` + +顶层 `HIERARCHY` 当前顶级节点： + +```text +安全类 +``` + +## 5. 当前已验证的文件 + +### 5.1 顶层 `hierarchy.json` + +路径： + +```text +page_memory/hierarchy.json +``` + +检查结果： + +```text +top_keys: HIERARCHY, nodes, stats +node_count: 52 +page_count: 77 +page_ranges: [[225, 301]] +max_depth: 6 +``` + +### 5.2 Scope `scope.json` + +路径： + +```text +page_memory/scopes/p225-301/scope.json +``` + +检查结果： + +```json +{ + "scope_id": "p225-301", + "strategy": "fat_only_coarse_scope:refined", + "document_page_count": 423, + "page_count": 77, + "page_ranges": [[225, 301]], + "skeleton_count": 52 +} +``` + +确认：没有 `pages` 长列表。 + +### 5.3 Scope `fine_hierarchy.json` + +路径： + +```text +page_memory/scopes/p225-301/fine_hierarchy.json +``` + +检查结果： + +```text +top_keys: HIERARCHY, nodes, stats, scope +node_count: 52 +page_count: 77 +page_ranges: [[225, 301]] +max_depth: 6 +``` + +### 5.4 `trace.json` + +路径： + +```text +page_memory/trace.json +``` + +检查结果： + +```text +final_status: stopped_at_fine +stage_count: 10 +summary.page_count: 423 +summary.scope_id: p225-301 +``` + +最后几个 stage 的 page_info 均为 compact range： + +```text +C4.coarse_scope page_count=77 page_ranges=[[225, 301]] +C1.render_pages.coarse page_count=77 page_ranges=[[225, 301]] +C2.page_plan.coarse page_count=77 page_ranges=[[225, 301]] +C3b.title_detection page_count=77 page_ranges=[[225, 301]] +C4b.fine_hierarchy fat_leaf.page_count=77 fat_leaf.page_ranges=[[225, 301]] +``` + +## 6. 已跑过的代码检查 + +最近一次相关检查通过： + +```bash +uv run ruff check \ + apps/worker/app/services/page_memory/memory_service.py \ + apps/worker/app/services/page_memory/fine_hierarchy.py \ + apps/worker/scripts/debug_page_memory.py +``` + +```bash +python -m py_compile \ + apps/worker/app/services/page_memory/memory_service.py \ + apps/worker/app/services/page_memory/fine_hierarchy.py \ + apps/worker/scripts/debug_page_memory.py +``` + +```bash +uv run pytest \ + apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py \ + apps/worker/tests/contract/test_page_memory_node_assembler_contract.py \ + apps/worker/tests/contract/test_document_agent_budget_contract.py \ + -q +``` + +结果： + +```text +17 passed +``` + +## 7. 当前待讨论/后续优化点 + +### 7.1 `parent_paths` 仍偏长 + +`scope.json` 里目前仍保留 `parent_paths`，虽然不是 page 长列表，但对 debug 阅读来说有些臃肿。 + +可选优化： + +```json +{ + "root_path": "...", + "parent_path_count": 14 +} +``` + +完整 `parent_paths` 可以放入 `trace.json`。 + +### 7.2 C4 skeleton 残余定位耗时较长 + +本次 C4 skeleton 定位耗时约 279 秒，明显比 fine hierarchy 更重。 + +后续可考虑： + +1. 对已定位的 TOC 节点减少 residual VLM。 +2. 对 debug 模式增加更明确的 residual cap。 +3. 对多个 coarse scope 并发定位/处理。 +4. 复用 anatomy + skeleton cache 做快速迭代。 + +### 7.3 嵌入式 TOC 页 228 被跳过 + +本次 DOC_PROFILE 找到目录页 `[5, 6, 228]`，但 skeleton 阶段跳过了嵌入式目录区域 `[228]`。 + +这可能影响后续更细粒度 scope 的粗 hierarchy 完整性，需要单独评估： + +```text +embedded_toc_region_outside_front_cluster +``` + +### 7.4 下一步测试建议 + +建议下一步直接测试： + +```text +fine hierarchy -> PAGE-TAG +``` + +即跑到： + +```text +--stop-at tag +``` + +重点检查： + +1. `scopes/p225-301/page_tags.json` +2. 顶层 `page_tags.json` +3. `trace.json` 中 PAGE-TAG 是否只覆盖 `[[225, 301]]` +4. PAGE-TAG 是否能支撑后续 `chunks.json` 和 `doc_nav.json` + +之后再开启图表资产提取，验证： + +```text +assets.json +scopes/p225-301/assets.json +``` + +## 8. 当前涉及的主要代码文件 + +本轮 page-memory 相关核心改动集中在： + +```text +apps/worker/app/services/page_memory/memory_service.py +apps/worker/app/services/page_memory/fine_hierarchy.py +apps/worker/app/services/page_memory/page_renderer.py +apps/worker/app/services/page_memory/page_assets.py +apps/worker/app/services/page_memory/node_assembler.py +apps/worker/app/services/document_agent/trace.py +apps/worker/app/services/document_agent/visual.py +apps/worker/scripts/debug_page_memory.py +``` + +其中 `apps/worker/scripts/debug_page_memory.py` 是当前测试入口。 + diff --git a/uv.lock b/uv.lock index 884ec07d..56392c29 100644 --- a/uv.lock +++ b/uv.lock @@ -1501,6 +1501,7 @@ dependencies = [ { name = "psycopg2-binary" }, { name = "pydantic", extra = ["email"] }, { name = "pydantic-settings" }, + { name = "pypdf" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-mock" }, @@ -1547,6 +1548,7 @@ requires-dist = [ { name = "pydantic", specifier = "==2.13.4" }, { name = "pydantic", extras = ["email"], specifier = "==2.13.4" }, { name = "pydantic-settings", specifier = "==2.14.1" }, + { name = "pypdf", specifier = "==6.10.2" }, { name = "pytest", specifier = "==9.0.3" }, { name = "pytest-asyncio", specifier = "==1.3.0" }, { name = "pytest-mock", specifier = "==3.15.1" }, From 3bf5d691dafd86c1f677ed7e591fb4d1559d2773 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 25 Jun 2026 15:55:14 +0800 Subject: [PATCH 06/27] feat: integrate summary engine for enhanced document parsing and asset handling - Updated document parser to utilize the new summary engine for images, tables, and text, improving the extraction of titles, summaries, and entities. - Refactored image and table handling to streamline the summarization process, replacing legacy methods with unified calls to the summary engine. - Enhanced markdown parsing to support new entity and asset title fields, ensuring comprehensive data capture during document processing. - Introduced serialization for typed entities to maintain structured data integrity across parsed documents. --- .gitignore | 1 + .../document_parser/formats/docx/parser.py | 57 +- .../formats/excel/table_parser.py | 16 +- .../document_parser/formats/image/parser.py | 116 ++-- .../formats/markdown/deferred_summary.py | 149 +++--- .../formats/markdown/parse_state.py | 2 + .../formats/markdown/parser.py | 2 +- .../document_parser/formats/text/parser.py | 111 +--- .../document_parser/support/parser_rows.py | 69 +++ .../services/page_memory/node_assembler.py | 196 +++---- .../app/services/page_memory/page_assets.py | 111 +++- .../app/services/page_memory/page_tagger.py | 198 ++----- .../shared-python/shared/core/config/ai.py | 18 +- .../shared/services/ai/llm_mock.py | 49 +- .../shared/services/ai/prompt_service.py | 344 ++++++------ .../shared/services/ai/summary/engine.py | 501 ++++++++++++++++++ .../shared/services/ai/summary/model.py | 109 ++++ .../chunks/dataframe_chunk_converter.py | 46 +- .../services/retrieval/graph/keywords.py | 66 +++ .../services/retrieval/graph/service.py | 137 ++++- .../services/retrieval/search/lexical_text.py | 37 +- .../services/storage/zip_chunk_schema.py | 2 +- 22 files changed, 1531 insertions(+), 806 deletions(-) create mode 100644 packages/shared-python/shared/services/ai/summary/engine.py create mode 100644 packages/shared-python/shared/services/ai/summary/model.py diff --git a/.gitignore b/.gitignore index 66e4ded6..428af9f0 100644 --- a/.gitignore +++ b/.gitignore @@ -113,6 +113,7 @@ LICENSE.txt .agent/ .agent-hooks/ .codex/ +CLAUDE.md # Codex-generated local files .venv_unstructured/ diff --git a/apps/worker/app/services/document_parser/formats/docx/parser.py b/apps/worker/app/services/document_parser/formats/docx/parser.py index fd05cb3f..32585db5 100755 --- a/apps/worker/app/services/document_parser/formats/docx/parser.py +++ b/apps/worker/app/services/document_parser/formats/docx/parser.py @@ -24,8 +24,6 @@ predict_heading_hierarchy, ) from app.services.document_parser.formats.image.parser import ( - _get_vision_client, - ask_image, perceptual_hash, ) from app.services.document_parser.tables.table_text_parser import sanitize_table_name_from_header @@ -125,7 +123,6 @@ def handle_image( logger.debug(f"Skipped duplicate image (hash={img_hash[:12]}...)") return headings_stack, df_list, False # False = cache hit, don't increment - client = _get_vision_client() last_context = _find_img_context(headings_stack) # Image index (always present) @@ -141,18 +138,19 @@ def handle_image( llm_title = None llm_summary = None if smart_summary: - from app.services.document_parser.formats.text.parser import split_title_summary - - # TODO: Risk of missing text content if the image is a screenshot of pure text. - # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. - llm_resp = ask_image( - client, - asset_store.image_dir, - [f"{raw_img_name}{img_ext}"], - title_text=last_context, + from shared.services.ai.summary.engine import summarize + + # Asset contract (§4.1): the engine's per-type image prompt covers + # charts/tables/diagrams and pure-text screenshots alike (§4.2). + image_path = os.path.join(asset_store.image_dir, f"{raw_img_name}{img_ext}") + result = summarize( + mode="asset", + image_paths=[image_path], + text=last_context, + usage_task="parser.docx.image", ) - if llm_resp: - llm_title, llm_summary = split_title_summary(llm_resp) + llm_title = result.title or None + llm_summary = result.summary or None # Fallback: LLM summary -> last_context -> None img_summary = llm_summary or last_context or None @@ -320,12 +318,19 @@ def handle_table( img_summary = None if summary_image: try: - client = _get_vision_client() - img_summary = ask_image( - client, - asset_store.image_dir, - [f"{img_name}{img_ext}"], - title_text=current_heading, + from shared.services.ai.summary.engine import summarize + + cell_image_path = os.path.join( + asset_store.image_dir, f"{img_name}{img_ext}" + ) + img_summary = ( + summarize( + mode="asset", + image_paths=[cell_image_path], + text=current_heading, + usage_task="parser.docx.table_image", + ).summary + or None ) except Exception as e: logger.warning(f"Failed to summarize table image: {e}") @@ -392,13 +397,13 @@ def handle_table( llm_summary = None tb_keywords = "" if summary_table: - from app.services.document_parser.formats.text.parser import ( - extract_title_keywords_summary, - ) + from shared.services.ai.summary.engine import summarize - llm_title, tb_keywords, llm_summary = extract_title_keywords_summary( - tb_html_str, max_keywords=3 - ) + # Tables are Contract B assets: title + summary + keywords from HTML. + result = summarize(mode="asset", text=tb_html_str, max_keywords=3) + llm_title = result.title or None + tb_keywords = result.keywords_str() + llm_summary = result.summary or None # Build tb_summary for df_list: table-n + optional LLM summary if llm_summary: diff --git a/apps/worker/app/services/document_parser/formats/excel/table_parser.py b/apps/worker/app/services/document_parser/formats/excel/table_parser.py index fd57d2ee..1a0fde31 100644 --- a/apps/worker/app/services/document_parser/formats/excel/table_parser.py +++ b/apps/worker/app/services/document_parser/formats/excel/table_parser.py @@ -311,15 +311,19 @@ def _summarize_excel_table( mechanical_keywords = parse_tb_keywords(table_frame) if llm_parameters["summary_table"]: - from app.services.document_parser.formats.text.parser import ( - extract_title_keywords_summary, - ) + from shared.services.ai.summary.engine import summarize - title, keywords, summary = extract_title_keywords_summary( - table_html, + # Tables are Contract B assets: title + summary + keywords from HTML. + result = summarize( + mode="asset", + text=table_html, max_keywords=3, ) - return title, keywords or mechanical_keywords, summary + return ( + result.title or None, + result.keywords_str() or mechanical_keywords, + result.summary or None, + ) return None, mechanical_keywords, None diff --git a/apps/worker/app/services/document_parser/formats/image/parser.py b/apps/worker/app/services/document_parser/formats/image/parser.py index fe4f29b0..ab078744 100755 --- a/apps/worker/app/services/document_parser/formats/image/parser.py +++ b/apps/worker/app/services/document_parser/formats/image/parser.py @@ -25,6 +25,7 @@ from shared.utils.chunk_refs import build_chunk_ref from app.services.common.file_loading import is_remote, load_file_bytes from app.services.common.file_utils import path_handle +from shared.services.ai.summary.engine import summarize, transcribe from shared.services.ai.openai_compatible_client_sync import ( OpenAICompatibleClientSync, get_openai_client, @@ -188,31 +189,18 @@ def ask_image( return None -def detect_summary_img_md(line, last_context, image_root_dir, mode=False): - client = _get_vision_client() +def detect_summary_img_md(line, last_context, image_root_dir): + """Collect markdown image refs with a placeholder summary. + + The actual per-image summary is produced later by the deferred task path + (``ImageDeferredSummaryTask`` → unified summary engine), so here we only emit + a positional placeholder derived from the surrounding context. + """ imgs = [] img_paths = re.findall(MD_IMAGE_PATTERN, line, flags=re.IGNORECASE) for i, ip in enumerate(img_paths): - if mode: - try: - llm_resp = ask_image(client, image_root_dir, paths_=[ip]) - if llm_resp: - from app.services.document_parser.formats.text.parser import ( - split_title_summary, - ) - - img_title, image_summary = split_title_summary(llm_resp) - else: - img_title = None - image_summary = last_context + str(i) - except Exception: - img_title = None - image_summary = last_context + str(i) - else: - img_title = None - image_summary = last_context + str(i) - if image_summary is not None: - imgs.append((ip, img_title, image_summary)) + image_summary = last_context + str(i) + imgs.append((ip, None, image_summary)) return imgs @@ -253,10 +241,10 @@ def parse_image( # Extract image content client = _get_vision_client() + abs_image_path = os.path.join(output_dir, relative_source_path) - ## Determine image category and task - img_task = "summary-images" - img_max_tokens = ProcessingConstants.IMG_MAX_TOKENS + ## Classify the image so text-heavy scans go through OCR (§4.2) and + ## everything else through the asset summary contract (§4.1). img_context = f"{filename}\n{base_llm_paras['frag_desc']}" type_resp = ask_image( client, @@ -266,52 +254,44 @@ def parse_image( task="judge-image-type", size_cut=False, ) - if type_resp is not None: - if type_resp["answer"] == "text": - img_task = "ocr-image" - img_max_tokens = ProcessingConstants.IMG_OCR_MAX_TOKENS - - if base_llm_paras[ - "summary_image" - ]: # Leave room for context to help understand the image - image_content = ask_image( - client, - output_dir, - paths_=[relative_source_path], - title_text=img_context, - task=img_task, - max_tokens=img_max_tokens, - size_cut=False, - ) - if image_content is None: - image_content = filename - else: - image_content = filename + is_text_image = bool( + isinstance(type_resp, dict) and type_resp.get("answer") == "text" + ) - if type_resp["answer"] == "text" and base_llm_paras["summary_image"]: - llm_resp = ask_image( - client, - output_dir, - paths_=[relative_source_path], - title_text=filename, - size_cut=False, + if not base_llm_paras["summary_image"]: + img_title = None + image_summary = filename + image_content = filename + elif is_text_image: + # Text scan: transcribe the body (→ content) and summarize for title. + transcribed = transcribe( + image_paths=[abs_image_path], + max_tokens=ProcessingConstants.IMG_OCR_MAX_TOKENS, + usage_task="parser.image.transcribe", ) - if llm_resp: - from app.services.document_parser.formats.text.parser import split_title_summary - - img_title, image_summary = split_title_summary(llm_resp) - else: - img_title = None - image_summary = image_content + image_content = transcribed or filename + asset = summarize( + mode="asset", + image_paths=[abs_image_path], + text=filename, + summary_len=ProcessingConstants.IMG_MAX_TOKENS, + usage_task="parser.image.summary", + ) + img_title = asset.title or None + image_summary = asset.summary or image_content else: - # For non-text images, split title from summary-images response - if base_llm_paras["summary_image"] and image_content != filename: - from app.services.document_parser.formats.text.parser import split_title_summary - - img_title, image_summary = split_title_summary(image_content) - else: - img_title = None - image_summary = image_content + # Figure/chart/diagram: asset summary provides title + summary, and + # the summary doubles as the chunk content. + asset = summarize( + mode="asset", + image_paths=[abs_image_path], + text=img_context, + summary_len=ProcessingConstants.IMG_MAX_TOKENS, + usage_task="parser.image.summary", + ) + img_title = asset.title or None + image_summary = asset.summary or filename + image_content = asset.summary or filename # 2. Decide whether to rename based on image title and filename img_name = path_handle((img_title or image_summary)[:20], mode="clean_single") diff --git a/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py b/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py index 394ed06c..d946db4f 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py +++ b/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py @@ -3,7 +3,7 @@ import os import re from dataclasses import dataclass -from typing import Literal, TypeGuard +from typing import Literal import gevent from app.services.document_parser.formats.markdown.deferred_task import ( @@ -12,30 +12,30 @@ TableDeferredSummaryTask, TextDeferredSummaryTask, ) -from app.services.document_parser.formats.image.parser import _get_vision_client, ask_image from app.services.document_parser.support.stage_profiler import stage_timer from app.services.document_parser.tables.table_text_parser import sanitize_table_name_from_header -from app.services.document_parser.formats.text.parser import ( - extract_title_keywords_summary, - split_title_summary, +from app.services.document_parser.support.parser_rows import ( + COL_ASSET_TITLE, + COL_ENTITIES, + COL_KEYWORDS, + COL_SUMMARY, + apply_body_summary, + serialize_entities, ) from gevent.pool import Pool as GeventPool from loguru import logger from shared.core.config import settings +from shared.services.ai.summary.engine import summarize +from shared.services.ai.summary.model import AssetSummary, BodySummary from shared.utils.chunk_refs import build_chunk_ref from app.services.common.file_utils import path_handle -DeferredResult = ( - tuple[ - int, - Literal["image", "table", "text"], - tuple[str | None, str | None] | tuple[str, str, str] | tuple[str, str], - ] -) -ImageSummaryResult = tuple[str | None, str | None] -TableSummaryResult = tuple[str, str, str] -TextSummaryResult = tuple[str, str] +# Each deferred task now carries the engine's typed contract straight through to +# the apply step (audit §4.5): assets → AssetSummary, text → BodySummary. The row +# is then written by name via apply_asset_summary / apply_body_summary, so we no +# longer poke magic positional offsets. +DeferredResult = tuple[int, Literal["image", "table", "text"], object] @dataclass(frozen=True) @@ -118,31 +118,36 @@ def _run_deferred_summary_task( ) -> DeferredResult | None: try: if isinstance(task, ImageDeferredSummaryTask): - client = _get_vision_client() - # TODO: Risk of missing text content if MinerU outputted a pure text image. - # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image. - llm_resp = ask_image( - client, deferred_input.output_dir, paths_=[task.relative_path] + # Asset contract (§4.1): title + summary + entities from the image. + # The engine's per-type image prompt covers charts/tables/diagrams and + # pure-text images alike, closing the old MinerU "pure text image" gap. + image_path = os.path.join(deferred_input.output_dir, task.relative_path) + result = summarize( + mode="asset", + image_paths=[image_path], + usage_task="parser.image.deferred", ) - if llm_resp: - img_title, img_summary = split_title_summary(llm_resp) - else: - img_title, img_summary = None, None - return task.row_index, "image", (img_title, img_summary) + return task.row_index, "image", result if isinstance(task, TableDeferredSummaryTask): - title, keywords, summary = extract_title_keywords_summary( - task.table_html, max_keywords=3 + # Tables are assets: summarize from HTML → title + summary + entities. + result = summarize( + mode="asset", + text=task.table_html, + usage_task="parser.table.deferred", ) - return task.row_index, "table", (title, keywords, summary) + return task.row_index, "table", result if isinstance(task, TextDeferredSummaryTask): - _, keywords, summary = extract_title_keywords_summary( - task.content, + # Body content (Contract A): summary + entities, no title. + result = summarize( + mode="text", + text=task.content, max_keywords=3, summary_len=deferred_input.summary_len, + usage_task="parser.text.deferred", ) - return task.row_index, "text", (keywords, summary) + return task.row_index, "text", result except Exception as exc: logger.warning( f"Deferred summary LLM call failed for idx={task.row_index}: {exc}" @@ -165,7 +170,7 @@ def _apply_deferred_summary_results( row_index, task_type, task_result = result if task_type == "image": - if not _is_image_summary_result(task_result): + if not isinstance(task_result, AssetSummary): logger.warning(f"Invalid image deferred result for idx={row_index}") continue _apply_image_summary_result( @@ -175,7 +180,7 @@ def _apply_deferred_summary_results( task_result, ) elif task_type == "table": - if not _is_table_summary_result(task_result): + if not isinstance(task_result, AssetSummary): logger.warning(f"Invalid table deferred result for idx={row_index}") continue _apply_table_summary_result( @@ -185,34 +190,10 @@ def _apply_deferred_summary_results( task_result, ) elif task_type == "text": - if not _is_text_summary_result(task_result): + if not isinstance(task_result, BodySummary): logger.warning(f"Invalid text deferred result for idx={row_index}") continue - _apply_text_summary_result(deferred_input.rows, row_index, task_result) - - -def _is_image_summary_result(result: object) -> TypeGuard[ImageSummaryResult]: - return ( - isinstance(result, tuple) - and len(result) == 2 - and all(isinstance(value, (str, type(None))) for value in result) - ) - - -def _is_table_summary_result(result: object) -> TypeGuard[TableSummaryResult]: - return ( - isinstance(result, tuple) - and len(result) == 3 - and all(isinstance(value, str) for value in result) - ) - - -def _is_text_summary_result(result: object) -> TypeGuard[TextSummaryResult]: - return ( - isinstance(result, tuple) - and len(result) == 2 - and all(isinstance(value, str) for value in result) - ) + apply_body_summary(deferred_input.rows[row_index], task_result) def _get_image_task(task: MarkdownDeferredSummaryTask) -> ImageDeferredSummaryTask: @@ -227,18 +208,41 @@ def _get_table_task(task: MarkdownDeferredSummaryTask) -> TableDeferredSummaryTa raise TypeError(f"Expected table deferred task, got {type(task).__name__}") +def _apply_asset_result_preserving_index( + row: list[str | int], result: AssetSummary +) -> None: + """Write an asset result onto a markdown row, keeping the summary's index tag. + + Markdown image/table rows seed the summary column with an ``image-N`` / + ``table-N`` index token on the first line (see ``image_asset`` / + ``table_asset``). We preserve that token and append the generated summary + below it, then write entities, asset_title, and the transitional keywords by + name (audit §4.5). + """ + while len(row) <= COL_ASSET_TITLE: + row.append("") + if result.summary: + existing = str(row[COL_SUMMARY]) + index_token = existing.split("\n", 1)[0] if existing else "" + row[COL_SUMMARY] = ( + f"{index_token}\n{result.summary}" if index_token else result.summary + ) + row[COL_ENTITIES] = serialize_entities(result.entities) + row[COL_KEYWORDS] = result.keywords_str() + if result.title: + row[COL_ASSET_TITLE] = result.title + + def _apply_image_summary_result( rows: list[list[str | int]], original_task: ImageDeferredSummaryTask, row_index: int, - result: ImageSummaryResult, + result: AssetSummary, ) -> None: - img_title, img_summary = result row = rows[row_index] - if img_summary: - image_index = str(row[5]).split("\n")[0] if row[5] else "image" - row[5] = f"{image_index}\n{img_summary}" + _apply_asset_result_preserving_index(row, result) + img_title = result.title if not img_title: return @@ -270,15 +274,12 @@ def _apply_table_summary_result( rows: list[list[str | int]], original_task: TableDeferredSummaryTask, row_index: int, - result: TableSummaryResult, + result: AssetSummary, ) -> None: - title, keywords, summary = result row = rows[row_index] - row[4] = keywords if isinstance(keywords, str) else "" - if summary: - table_index = str(row[5]) if "\n" not in str(row[5]) else str(row[5]).split("\n")[0] - row[5] = f"{table_index}\n{summary}" + _apply_asset_result_preserving_index(row, result) + title = result.title if not title: return @@ -298,11 +299,3 @@ def _apply_table_summary_result( new_relative_path = f"tables/{new_table_name}.html" replace_chunk_ref_in_rows(rows, str(row[1]), new_relative_path) row[1] = new_relative_path - - -def _apply_text_summary_result( - rows: list[list[str | int]], row_index: int, result: TextSummaryResult -) -> None: - keywords, summary = result - rows[row_index][4] = keywords if isinstance(keywords, str) else "" - rows[row_index][5] = summary if isinstance(summary, str) else "" diff --git a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py index 95e5088e..8afab023 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py @@ -167,6 +167,8 @@ def to_dataframe(self) -> pd.DataFrame: connectto=str(row_values[8]), addtime=str(row_values[9]), page_nums=str(row_values[10]), + entities=str(row_values[11]) if len(row_values) > 11 else "", + asset_title=str(row_values[12]) if len(row_values) > 12 else "", ) ) return process_dup_paths_df(rows_builder.to_dataframe()) diff --git a/apps/worker/app/services/document_parser/formats/markdown/parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py index 8c887899..1d197827 100755 --- a/apps/worker/app/services/document_parser/formats/markdown/parser.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py @@ -371,7 +371,7 @@ def parse_md( else: # no path change, remain in the same hierarchy # a. handle lines containing images (LLM deferred to post-loop parallel batch) # Always skip inline LLM — vision calls are deferred to parallel batch - imgs = detect_summary_img_md(line, last_context, output_dir, mode=False) + imgs = detect_summary_img_md(line, last_context, output_dir) image_name = build_markdown_image_name( image_count=parser_state.image_count, last_context=last_context, diff --git a/apps/worker/app/services/document_parser/formats/text/parser.py b/apps/worker/app/services/document_parser/formats/text/parser.py index b819bd2f..fa6d80e5 100755 --- a/apps/worker/app/services/document_parser/formats/text/parser.py +++ b/apps/worker/app/services/document_parser/formats/text/parser.py @@ -1,6 +1,5 @@ # pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportGeneralTypeIssues=false import re -import uuid import gevent import pandas as pd @@ -9,11 +8,8 @@ from loguru import logger from shared.core.config import settings -from shared.services.ai.prompt_service import build_prompt -from shared.services.ai.response_process_service import eval_response from shared.utils.chunk_refs import CHUNK_REF_PATTERN from app.services.common.file_loading import load_file_bytes -from shared.services.ai.openai_compatible_client_sync import get_openai_client def clean_texts_by_form(text, form="html"): @@ -66,112 +62,39 @@ def divide_long_contents(texts, max_threshold=None, min_threshold=None): return sublists, len(sublists) -def split_title_summary(text): - """Split a title+summary response into (title, summary). - - Expected format: first line is title, remaining lines are summary. - Fallback: if only one line, it serves as both title and summary. - - Returns: - tuple: (title, summary) — both may be None if input is empty - """ - if not text or not text.strip(): - return None, None - parts = text.strip().split("\n", 1) - title = parts[0].strip() - summary = parts[1].strip() if len(parts) > 1 else title - return title, summary - - def extract_title_keywords_summary(texts, max_keywords=3, summary_len=None): - """Extract title + keywords + summary in ONE LLM call. + """Extract keywords + summary for a text chunk via the unified engine (§4.1). - Uses the 'summary-full' prompt to get all three fields at once, - reducing LLM calls from 2-3 to 1. + Delegates to ``summarize(mode="text")``. The engine centralizes the + ``summary-full`` prompt, deterministic language locking, JSON parsing, and + retry. Title is no longer produced here — body content (Contract A) carries + only summary + entities; the section title already lives on the node path. Args: texts: Input text (may include HTML tables or structured data). - max_keywords: Maximum number of keywords to extract (default 3). + max_keywords: Maximum number of keywords to extract. summary_len: Maximum summary length in characters. Returns: tuple: (title, keywords_str, summary) - - title: short title (≤15 chars), or None - - keywords_str: semicolon-separated keywords, or "" + - title: always None (kept for signature compatibility; §4.5 removes it) + - keywords_str: semicolon-separated entity surface forms, or "" - summary: summary text, or "" """ from shared.core.constants import ProcessingConstants - from shared.services.ai.prompt_service import _detect_text_language + from shared.services.ai.summary.engine import summarize if summary_len is None: summary_len = ProcessingConstants.SUMMARY_LEN - try: - # Deterministic language lock: LLMs (especially deepseek-v4-flash) often - # default to Chinese on numeric / structured inputs even when asked to - # match input language. We detect the input's dominant language here - # and pass it as a HARD constraint to the prompt. - detected_lang = _detect_text_language(texts) - prompt, temperature, top_p, max_tokens = build_prompt( - task="summary-full", - texts=texts, - query="", - paras={ - "max_tokens": summary_len, - "kw_num": max_keywords, - "lang": detected_lang, - }, - ) - messages = [ - {"role": "system", "content": "you are a helpful assistant"}, - {"role": "user", "content": prompt}, - ] - - import os - if os.getenv("LOCAL_DEBUG", "0") != "1": - from shared.services.redis.redis_sync_service import SyncRedisServiceFactory - - redis_service = SyncRedisServiceFactory.get_service() - ctx_task_id = str(uuid.uuid4()) - redis_service.set(f"task:{ctx_task_id}:status", "processing", ttl=7200) - - resp = get_openai_client().chat_completion( - messages=messages, - timeout=90, - max_tokens=max_tokens, - usage_task="parser.text_summary", - ) - - # Handle null/none response - if resp is None: - return None, "", "" - if isinstance(resp, str): - resp_stripped = resp.strip().lower() - if resp_stripped in ("null", "none"): - return None, "", "" - - # Parse JSON response - parsed = eval_response(resp) - if isinstance(parsed, dict): - title = parsed.get("title") or None - keywords = parsed.get("keywords", "") - summary = parsed.get("summary", "") - # Normalize title - if title and isinstance(title, str): - title = title.strip() - if title.lower() in ("null", "none", ""): - title = None - return ( - title, - keywords if isinstance(keywords, str) else "", - summary if isinstance(summary, str) else "", - ) - - return None, "", "" - - except Exception as e: - print(f"❌ failed to extract title/keywords/summary: {e}") - return None, "", "" + result = summarize( + mode="text", + text=texts, + summary_len=summary_len, + max_keywords=max_keywords, + usage_task="parser.text_summary", + ) + return None, result.keywords_str(), result.summary def postprocess_leaf_dics( diff --git a/apps/worker/app/services/document_parser/support/parser_rows.py b/apps/worker/app/services/document_parser/support/parser_rows.py index 839ddd9f..9a12298b 100644 --- a/apps/worker/app/services/document_parser/support/parser_rows.py +++ b/apps/worker/app/services/document_parser/support/parser_rows.py @@ -1,6 +1,8 @@ from __future__ import annotations +import json from dataclasses import dataclass +from typing import Any import pandas as pd from pandas import Index @@ -10,6 +12,29 @@ PARSER_ROW_COLUMNS: tuple[str, ...] = tuple(settings.ALL_DF_COLS.split(",")) +def serialize_entities(entities: Any) -> str: + """Serialize typed entities (§4.4) into the JSON string stored in the row. + + Accepts a list of ``{"text","type"}`` dicts (or objects exposing + ``to_dict``). Returns ``""`` for an empty/invalid value so empty stays empty + rather than becoming ``"[]"`` noise in the column. + """ + if not entities: + return "" + normalized: list[dict[str, str]] = [] + for item in entities: + if hasattr(item, "to_dict"): + item = item.to_dict() + if isinstance(item, dict): + text = str(item.get("text", "")).strip() + if not text: + continue + normalized.append({"text": text, "type": str(item.get("type", "")).strip()}) + if not normalized: + return "" + return json.dumps(normalized, ensure_ascii=False) + + @dataclass(frozen=True) class ParsedRow: content: str @@ -23,6 +48,8 @@ class ParsedRow: connectto: str = "" page_nums: str = "" length: int | None = None + entities: str = "" + asset_title: str = "" def to_list(self) -> list[object]: content_length = self.length if self.length is not None else len(self.content) @@ -38,6 +65,8 @@ def to_list(self) -> list[object]: self.connectto, self.addtime, self.page_nums, + self.entities, + self.asset_title, ] def to_dict(self) -> dict[str, object]: @@ -59,3 +88,43 @@ def to_dataframe(self) -> pd.DataFrame: [row.to_list() for row in self._rows], columns=Index(PARSER_ROW_COLUMNS), ) + + +# Column indices into the positional raw-row lists (markdown track). Named here +# so callers stop poking magic offsets like row[4]/row[5] (audit P5). Kept in +# sync with ``ParsedRow.to_list`` / ``PARSER_ROW_COLUMNS``. +COL_KEYWORDS = PARSER_ROW_COLUMNS.index("keywords") +COL_SUMMARY = PARSER_ROW_COLUMNS.index("summary") +COL_ENTITIES = PARSER_ROW_COLUMNS.index("entities") +COL_ASSET_TITLE = PARSER_ROW_COLUMNS.index("asset_title") + + +def apply_body_summary(row: list[Any], result: Any) -> None: + """Write a ``BodySummary`` (Contract A) onto a positional raw row by name. + + Sets ``summary`` and ``entities`` (and the transitional flattened + ``keywords``). Replaces fragile ``row[5] = ...`` writes (audit §4.5). + """ + _ensure_row_width(row) + row[COL_SUMMARY] = result.summary or "" + row[COL_ENTITIES] = serialize_entities(getattr(result, "entities", None)) + row[COL_KEYWORDS] = result.keywords_str() + + +def apply_asset_summary(row: list[Any], result: Any) -> None: + """Write an ``AssetSummary`` (Contract B) onto a positional raw row by name. + + Sets ``asset_title``, ``summary``, ``entities`` and the transitional + ``keywords``. The caller owns any file-rename logic keyed off the title. + """ + _ensure_row_width(row) + row[COL_ASSET_TITLE] = result.title or "" + row[COL_SUMMARY] = result.summary or "" + row[COL_ENTITIES] = serialize_entities(getattr(result, "entities", None)) + row[COL_KEYWORDS] = result.keywords_str() + + +def _ensure_row_width(row: list[Any]) -> None: + """Pad a legacy 11-column row out to the current schema width in place.""" + while len(row) < len(PARSER_ROW_COLUMNS): + row.append("") diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index 881c6cd3..d27385d7 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -20,24 +20,23 @@ from __future__ import annotations -import base64 import json import os from dataclasses import dataclass, field -from typing import Any, cast +from typing import Any from loguru import logger from app.services.document_agent.budget import BudgetTracker from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import serialize_entities from app.services.page_memory.page_assets import ( PageAsset, build_asset_rows, ) from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton -from shared.services.ai.prompt_service import build_prompt -from shared.utils.token_estimate import estimate_tokens +from shared.services.ai.summary.engine import summarize, transcribe SAME_AS_PREFIX = "SAME-AS" @@ -250,15 +249,6 @@ def _node_summary_max_pages() -> int: ) -def _read_image_b64(image_path: str) -> str | None: - try: - with open(image_path, "rb") as handle: - return base64.b64encode(handle.read()).decode() - except Exception as exc: # pragma: no cover - filesystem edge - logger.warning("[node_assembler] failed to read image {}: {}", image_path, exc) - return None - - def resolve_page_text( *, page: int, @@ -270,65 +260,21 @@ def resolve_page_text( """Body text for an owned page: PyMuPDF text, or VLM OCR for scanned pages. Electronic PDFs already have PyMuPDF text; scanned pages have (near) empty - text and fall back to a one-shot VLM OCR transcription. + text and fall back to the shared ``transcribe()`` OCR primitive (§4.2). """ 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 "" - - img_b64 = _read_image_b64(image_path) - if img_b64 is None: - return "" - - prompt, temperature, _top_p, max_tokens = build_prompt( - "page-memory-vlm-ocr", "", "", paras={"max_tokens": 1500} + return transcribe( + image_paths=[image_path], + model=vlm_model, + max_tokens=1500, + usage_task="page_memory.node_ocr", + budget=budget, + budget_stage=_BUDGET_STAGE, ) - est = estimate_tokens(prompt) + 1000 - if budget is not None and not budget.try_reserve( - "visual", est, stage=_BUDGET_STAGE - ): - logger.debug("[node_assembler] OCR budget exhausted for page {}", page) - return "" - - try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - client = get_openai_client(model=vlm_model) - content_parts = [ - {"type": "text", "text": prompt}, - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{img_b64}"}, - }, - ] - raw_response, usage = client.chat_completion_with_usage( - messages=cast(Any, [{"role": "user", "content": content_parts}]), - model=vlm_model, - temperature=temperature, - max_tokens=max_tokens, - response_format={"type": "json_object"}, - usage_task="page_memory.node_ocr", - ) - if budget is not None: - budget.commit( - "visual", - actual=usage.get("total_tokens", est), - est=est, - stage=_BUDGET_STAGE, - ) - try: - data = json.loads(raw_response) - return str(data.get("text", "")).strip() - except (json.JSONDecodeError, TypeError): - # VLM returned non-JSON; use raw text directly - return raw_response.strip() if raw_response else "" - except Exception as exc: - logger.warning("[node_assembler] OCR failed for page {}: {}", page, exc) - if budget is not None: - budget.refund("visual", est=est, stage=_BUDGET_STAGE) - return "" def compute_node_summary( @@ -339,17 +285,21 @@ def compute_node_summary( image_path_by_page: dict[int, str], vlm_model: str | None, budget: BudgetTracker | None, -) -> tuple[str, list[str]]: - """Settle a node's summary/keywords. +) -> tuple[str, list[str], list[dict[str, str]]]: + """Settle a node's summary, keywords, and typed entities (§4.4). Reuses the per-page tag when the node is a single page that no sibling leaf shares. Otherwise asks the VLM to summarize the node as a whole, bounding the slice with the next sibling title when the page hosts multiple nodes. Falls back to combining per-page tags when the VLM is unavailable. + + Returns ``(summary, keywords, entities)`` where ``keywords`` is the flattened + surface-form list (transitional) and ``entities`` is the typed + ``{"text","type"}`` list. """ pages = view.pages if not pages: - return "", [] + return "", [], [] single_page = len(pages) == 1 shared = any(len(page_to_leaves.get(page, [])) > 1 for page in pages) @@ -357,8 +307,8 @@ def compute_node_summary( if single_page and not shared: tag = tag_by_page.get(pages[0]) if tag is not None: - return tag.summary, list(tag.keywords) - return "", [] + return tag.summary, list(tag.keywords), list(tag.entities) + return "", [], [] if vlm_model: result = _vlm_node_summary( @@ -378,10 +328,12 @@ def _combine_page_tags( *, pages: list[int], tag_by_page: dict[int, PageTagResult], -) -> tuple[str, list[str]]: +) -> tuple[str, list[str], list[dict[str, str]]]: summaries: list[str] = [] keywords: list[str] = [] + entities: list[dict[str, str]] = [] seen: set[str] = set() + seen_entities: set[str] = set() for page in pages: tag = tag_by_page.get(page) if tag is None: @@ -394,7 +346,17 @@ def _combine_page_tags( if key and key not in seen: seen.add(key) keywords.append(keyword.strip()) - return " ".join(summaries).strip(), keywords + for entity in tag.entities: + entity_text = str(entity.get("text", "")).strip() + if not entity_text: + continue + entity_key = entity_text.casefold() + if entity_key not in seen_entities: + seen_entities.add(entity_key) + entities.append( + {"text": entity_text, "type": str(entity.get("type", "")).strip()} + ) + return " ".join(summaries).strip(), keywords, entities def _vlm_node_summary( @@ -404,7 +366,7 @@ def _vlm_node_summary( image_path_by_page: dict[int, str], vlm_model: str, budget: BudgetTracker | None, -) -> tuple[str, list[str]] | None: +) -> tuple[str, list[str], list[dict[str, str]]] | None: leaf = view.leaf pages = view.pages[: _node_summary_max_pages()] @@ -416,78 +378,36 @@ def _vlm_node_summary( leaves_on_page=page_to_leaves.get(leaf.start_page, []), ) - image_parts: list[dict[str, Any]] = [] - for page in pages: - path = image_path_by_page.get(page) - if not path or not os.path.exists(path): - continue - img_b64 = _read_image_b64(path) - if img_b64 is None: - continue - image_parts.append( - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{img_b64}"}, - } - ) - if not image_parts: + image_paths = [ + path + for page in pages + if (path := image_path_by_page.get(page)) and os.path.exists(path) + ] + if not image_paths: return None - prompt, temperature, _top_p, max_tokens = build_prompt( - "page-memory-node-summary", - "", - "", - paras={ + result = summarize( + mode="page", + image_paths=image_paths, + model=vlm_model, + usage_task="page_memory.node_summary", + budget=budget, + budget_stage=_BUDGET_STAGE, + prompt_task="page-memory-node-summary", + prompt_paras={ "max_tokens": 400, "node_title": leaf.title, "next_title": next_title or "", "kw_num": 5, }, ) - est = estimate_tokens(prompt) + 800 * len(image_parts) - if budget is not None and not budget.try_reserve( - "visual", est, stage=_BUDGET_STAGE - ): - logger.debug( - "[node_assembler] node summary budget exhausted for {}", - leaf.section_path, - ) - return None - - try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - client = get_openai_client(model=vlm_model) - content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] - content_parts.extend(image_parts) - raw_response, usage = client.chat_completion_with_usage( - messages=cast(Any, [{"role": "user", "content": content_parts}]), - model=vlm_model, - temperature=temperature, - max_tokens=max_tokens, - response_format={"type": "json_object"}, - usage_task="page_memory.node_summary", - ) - if budget is not None: - budget.commit( - "visual", - actual=usage.get("total_tokens", est), - est=est, - stage=_BUDGET_STAGE, - ) - data = json.loads(raw_response) - kw_str = str(data.get("keywords", "")) - keywords = [k.strip() for k in kw_str.split(";") if k.strip()] - return str(data.get("summary", "")).strip(), keywords - except Exception as exc: - logger.warning( - "[node_assembler] node summary VLM failed for {}: {}", - leaf.section_path, - exc, - ) - if budget is not None: - budget.refund("visual", est=est, stage=_BUDGET_STAGE) + 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], + ) # ── Orchestration ──────────────────────────────────────────────────── @@ -533,7 +453,7 @@ def build_node_rows( page_owner=page_owner, page_text=resolved_text, ) - summary, keywords = compute_node_summary( + summary, keywords, entities = compute_node_summary( view=view, page_to_leaves=page_to_leaves, tag_by_page=tag_by_page, @@ -554,6 +474,8 @@ def build_node_rows( "connectto": "", "addtime": get_str_time(), "page_nums": ",".join(str(page) for page in view.pages), + "entities": serialize_entities(entities), + "asset_title": "", "extra_metadata": {}, } rows.append(row) diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index b04c71b5..8d164c26 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -19,6 +19,7 @@ from app.services.document_agent.budget import BudgetTracker from app.services.document_agent.visual import visual_debug_enabled from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import serialize_entities from app.services.page_memory.page_renderer import PageRenderResult from shared.services.ai.prompt_service import build_prompt from shared.utils.token_estimate import estimate_tokens @@ -46,6 +47,8 @@ class PageAsset: title: str = "" summary: str = "" keywords: list[str] = field(default_factory=list) + entities: list[dict[str, str]] = field(default_factory=list) + chart: dict[str, Any] | None = None image_uri: str = "" html_uri: str = "" image_path: str = "" @@ -59,6 +62,17 @@ def page_asset_extraction_enabled() -> bool: ).strip().lower() in {"1", "true", "yes", "on"} +def page_asset_summary_enabled() -> bool: + """Whether detected page assets are summarized via the engine (§4.3). + + Gated separately from detection so the richer chart/figure summarization can + be rolled out independently. Defaults off. + """ + return os.environ.get( + "PAGE_MEMORY_ASSET_SUMMARY_ENABLED", "false" + ).strip().lower() in {"1", "true", "yes", "on"} + + def get_asset_confidence_threshold() -> float: try: return float(os.environ.get("PAGE_MEMORY_ASSET_CONFIDENCE_THRESHOLD", "0.3")) @@ -338,6 +352,67 @@ def extract_table_html( return asset +def summarize_page_asset( + *, + asset: PageAsset, + model_name: str | None, + budget: BudgetTracker | None, +) -> PageAsset: + """Summarize a cropped asset via the unified engine (§4.3). + + Routes tables (with extracted HTML) through the text asset path and figures / + charts through the image asset path. Populates ``summary``, ``entities``, + ``chart`` (statistical content), and a ``title`` when the asset had none. + Numbers and entities come entirely from the model reading the asset; nothing + is hard-coded here. + """ + from shared.services.ai.summary.engine import summarize + + table_html = "" + if asset.kind == "table" and asset.html_path and os.path.exists(asset.html_path): + try: + table_html = Path(asset.html_path).read_text(encoding="utf-8") + except Exception as exc: + logger.warning("[page_assets] failed to read table html: {}", exc) + + if table_html: + result = summarize( + mode="asset", + text=table_html, + model=model_name, + usage_task="page_memory.asset_summary", + budget=budget, + budget_stage=_BUDGET_STAGE, + asset_title_hint=asset.title, + ) + elif asset.image_path and os.path.exists(asset.image_path): + result = summarize( + mode="asset", + image_paths=[asset.image_path], + text=asset.title, + model=model_name, + usage_task="page_memory.asset_summary", + budget=budget, + budget_stage=_BUDGET_STAGE, + asset_title_hint=asset.title, + ) + else: + return asset + + if result.title and not asset.title: + asset.title = result.title + if result.summary: + asset.summary = result.summary + if result.entities: + asset.entities = [e.to_dict() for e in result.entities] + asset.keywords = [e.text for e in result.entities if e.text] + if result.chart is not None: + asset.chart = result.chart.to_dict() + if result.summary or result.entities: + asset.extraction_status = "summarized" + return asset + + def extract_page_assets_from_renders( *, pdf_path: str, @@ -360,6 +435,7 @@ def extract_page_assets_from_renders( confidence_threshold=confidence_threshold, ) page_assets: list[PageAsset] = [] + summary_enabled = page_asset_summary_enabled() for asset in detected: crop_page_asset( asset=asset, @@ -372,6 +448,12 @@ def extract_page_assets_from_renders( pdf_path=pdf_path, output_dir=output_dir, ) + if summary_enabled and (asset.image_uri or asset.html_uri): + summarize_page_asset( + asset=asset, + model_name=model_name, + budget=budget, + ) if asset.image_uri or asset.html_uri: page_assets.append(asset) if page_assets: @@ -451,9 +533,6 @@ def build_asset_rows( continue row_type = "table" if asset.kind == "table" and asset.html_uri else "image" content = _asset_content(asset, row_type=row_type) - summary = "\n".join( - part for part in [asset.title.strip(), asset.summary.strip()] if part - ) rows.append( { "content": content, @@ -461,26 +540,30 @@ def build_asset_rows( "type": row_type, "length": len(content), "keywords": ";".join(asset.keywords), - "summary": summary, + "summary": asset.summary.strip(), "know_id": asset.asset_id, "tokens": "", "connectto": "", "addtime": get_str_time(), "page_nums": str(asset.page_index), - "extra_metadata": { - "asset_index": asset.asset_index, - "bbox_px": asset.bbox_px, - "confidence": asset.confidence, - "caption": asset.caption, - "image_uri": asset.image_uri, - "extraction_status": asset.extraction_status, - "table_engine": "tabula" if asset.kind == "table" else "", - }, + "entities": serialize_entities(asset.entities), + "asset_title": asset.title.strip(), + "extra_metadata": _asset_extra_metadata(asset), } ) return rows +def _asset_extra_metadata(asset: PageAsset) -> dict[str, Any]: + metadata: dict[str, Any] = { + "caption": asset.caption, + "image_uri": asset.image_uri, + } + if asset.chart: + metadata["chart"] = asset.chart + return metadata + + def asset_reference(asset: PageAsset) -> str: uri = ( asset.html_uri if asset.kind == "table" and asset.html_uri else asset.image_uri @@ -629,4 +712,6 @@ def _safe_float(value: object, *, default: float) -> float: "get_asset_max_pages", "get_asset_model", "page_asset_extraction_enabled", + "page_asset_summary_enabled", + "summarize_page_asset", ] diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index 9428de10..8173313a 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -28,6 +28,7 @@ from app.services.page_memory.page_plan import PagePlan, PageProcessingStrategy from app.services.page_memory.page_renderer import PageRenderResult from shared.services.ai.prompt_service import build_prompt +from shared.services.ai.summary.engine import summarize from shared.utils.token_estimate import estimate_tokens @@ -39,6 +40,9 @@ class PageTagResult: summary: str = "" keywords: list[str] = field(default_factory=list) strategy_used: str = "" + entities: list[dict[str, str]] = field(default_factory=list) + """Typed entities (§4.4): ``{"text","type"}`` dicts. ``keywords`` is kept as + the flattened surface-form view for transitional keyword-overlap consumers.""" observed_titles: list[dict[str, Any]] = field(default_factory=list) """Step 2: verbatim title candidates observed on this page. @@ -50,7 +54,6 @@ class PageTagResult: _BUDGET_STAGE = "page_tagging" _BUDGET_STAGE_TITLES = "page_title_detection" _MAX_JSON_RETRIES = 1 -_RAW_TEXT_SUMMARY_LIMIT = 500 _DEFAULT_FINE_MIN_PAGES = 4 @@ -135,9 +138,11 @@ def _tag_skip(page: PageRenderResult) -> PageTagResult: def _tag_text_only(page: PageRenderResult) -> PageTagResult: - """Use existing ``summary-full`` LLM prompt to extract summary + keywords. + """Extract summary + keywords from raw page text via the unified engine. - Falls back to raw text truncation if LLM is not available or fails. + Returns an EMPTY-marked result when the page has no extractable text. On LLM + failure the engine returns an empty summary; we surface that directly (no + raw-text truncation fallback — empty means the model could not summarize). """ raw = page.raw_text.strip() if not raw: @@ -148,51 +153,19 @@ def _tag_text_only(page: PageRenderResult) -> PageTagResult: strategy_used="text_only", ) - # Try the existing summary-full LLM call (same as text chunk pipeline) - try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - text_model = os.environ.get("NORMOL_MODEL", "deepseek-v4-flash") - prompt, temperature, top_p, max_tokens = build_prompt( - "summary-full", - raw[:3000], # limit input to avoid token overflow - "", - paras={"max_tokens": 200, "kw_num": 5}, - ) - client = get_openai_client(model=text_model) - raw_response, _ = client.chat_completion_with_usage( - messages=[{"role": "user", "content": prompt}], - model=text_model, - temperature=temperature, - max_tokens=max_tokens, - usage_task="page_memory.text_only_summary", - ) - - if raw_response and raw_response.strip().lower() != "null": - data = json.loads(raw_response) - summary = str(data.get("summary", "")) - kw_str = str(data.get("keywords", "")) - keywords = [k.strip() for k in kw_str.split(";") if k.strip()] - return PageTagResult( - page_index=page.page_index, - summary=summary, - keywords=keywords, - strategy_used="text_only", - ) - except Exception as exc: - logger.warning( - "[page_tagger] summary-full LLM failed for page {}: {}; " - "falling back to raw text truncation", - page.page_index, exc, - ) - - # Fallback: raw text truncation - summary = " ".join(raw.split())[:_RAW_TEXT_SUMMARY_LIMIT] + result = summarize( + mode="text", + text=raw[:3000], # limit input to avoid token overflow + summary_len=200, + max_keywords=5, + usage_task="page_memory.text_only_summary", + ) return PageTagResult( page_index=page.page_index, - summary=summary, - keywords=[], - strategy_used="text_only_fallback", + summary=result.summary, + keywords=[e.text for e in result.entities], + entities=[e.to_dict() for e in result.entities], + strategy_used="text_only", ) @@ -202,121 +175,40 @@ def _tag_vlm_lite( model: str, budget: BudgetTracker | None, ) -> PageTagResult: - """Send page PNG to VLM and parse JSON response.""" - prompt, temperature, _top_p, max_tokens = build_prompt( - "page-memory-vlm-tag", - "", - "", - paras={"max_tokens": 600}, - ) - est = estimate_tokens(prompt) + 800 # ~800 tokens for image - - if budget is not None: - if not budget.try_reserve("visual", est, stage=_BUDGET_STAGE): - logger.warning( - "[page_tagger] insufficient budget for page {}; text_only fallback", - page.page_index, - ) - result = _tag_text_only(page) - result = PageTagResult( - page_index=result.page_index, - summary=result.summary, - keywords=result.keywords, - strategy_used="text_only_budget_fallback", - ) - return result - + """Send page PNG to the VLM via the unified engine; text_only on miss.""" if not page.image_path or not os.path.exists(page.image_path): logger.warning( "[page_tagger] no PNG for page {}; text_only fallback", page.page_index, ) - if budget is not None: - budget.refund("visual", est=est, stage=_BUDGET_STAGE) return _tag_text_only(page) - try: - with open(page.image_path, "rb") as f: - img_b64 = base64.b64encode(f.read()).decode() - except Exception as exc: - logger.warning( - "[page_tagger] failed to read PNG for page {}: {}", - page.page_index, exc, + result = summarize( + mode="page", + image_paths=[page.image_path], + model=model, + usage_task="page_memory.tag", + budget=budget, + budget_stage=_BUDGET_STAGE, + ) + if not result.summary and not result.entities: + # Engine returned empty (budget exhausted, image unreadable, or JSON + # miss). Fall back to text_only so the page still gets a summary. + fallback = _tag_text_only(page) + return PageTagResult( + page_index=fallback.page_index, + summary=fallback.summary, + keywords=fallback.keywords, + entities=fallback.entities, + strategy_used="vlm_lite_fallback", ) - if budget is not None: - budget.refund("visual", est=est, stage=_BUDGET_STAGE) - return _tag_text_only(page) - - content_parts: list[dict[str, Any]] = [ - {"type": "text", "text": prompt}, - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{img_b64}"}, - }, - ] - - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - client = get_openai_client(model=model) - - for attempt in range(_MAX_JSON_RETRIES + 1): - try: - raw_response, usage = client.chat_completion_with_usage( - messages=cast(Any, [{"role": "user", "content": content_parts}]), - model=model, - temperature=temperature, - max_tokens=max_tokens, - response_format={"type": "json_object"}, - usage_task="page_memory.tag", - ) - if budget is not None: - budget.commit( - "visual", - actual=usage.get("total_tokens", est), - est=est, - stage=_BUDGET_STAGE, - ) - - data = json.loads(raw_response) - kw_str = str(data.get("keywords", "")) - keywords = [k.strip() for k in kw_str.split(";") if k.strip()] - return PageTagResult( - page_index=page.page_index, - summary=str(data.get("summary", "")), - keywords=keywords, - strategy_used="vlm_lite", - ) - except json.JSONDecodeError: - if attempt < _MAX_JSON_RETRIES: - logger.warning( - "[page_tagger] JSON parse failed for page {} (attempt {}/{}), retrying", - page.page_index, attempt + 1, _MAX_JSON_RETRIES + 1, - ) - continue - # Final attempt failed: fallback to text_only - logger.warning( - "[page_tagger] JSON retry exhausted for page {}; text_only fallback", - page.page_index, - ) - result = _tag_text_only(page) - result = PageTagResult( - page_index=result.page_index, - summary=result.summary, - keywords=result.keywords, - strategy_used="vlm_lite_json_fallback", - ) - return result - except Exception as exc: - logger.warning( - "[page_tagger] VLM call failed for page {}: {}", - page.page_index, exc, - ) - if budget is not None: - budget.refund("visual", est=est, stage=_BUDGET_STAGE) - return _tag_text_only(page) - - # Should not reach here, but safety net - return _tag_text_only(page) + return PageTagResult( + page_index=page.page_index, + summary=result.summary, + keywords=[e.text for e in result.entities], + entities=[e.to_dict() for e in result.entities], + strategy_used="vlm_lite", + ) # ── Step 2: Independent title candidate extraction ─────────────────── diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index d823b8a9..c7732a51 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -163,6 +163,20 @@ class AIConfig(BaseModel): ) SPLIT_CHAR: str = Field(default="/", description="Path separator") ALL_DF_COLS: str = Field( - default="content,path,type,length,keywords,summary,know_id,tokens,connectto,addtime,page_nums", - description="All dataframe columns (compatibility field)", + default="content,path,type,length,keywords,summary,know_id,tokens,connectto,addtime,page_nums,entities,asset_title", + description=( + "All dataframe columns. `entities` (JSON-encoded typed entities, §4.4) " + "and `asset_title` (asset caption/label, §4.5) are additive trailing " + "columns; legacy `keywords` is retained transitionally." + ), + ) + ENTITY_TYPES: str = Field( + default="person,location,organization", + description=( + "Comma-separated seed list of entity types the summarizer may emit " + "(§4.4). Extend this to broaden extraction (e.g. product, date, money) " + "without code changes. Order is not significant; matching is " + "case-insensitive. An empty value disables type guidance and lets the " + "model choose, but the seed list keeps cross-document links consistent." + ), ) diff --git a/packages/shared-python/shared/services/ai/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py index ccb84143..28748575 100644 --- a/packages/shared-python/shared/services/ai/llm_mock.py +++ b/packages/shared-python/shared/services/ai/llm_mock.py @@ -113,17 +113,17 @@ def _detect_mock_task(prompt_text: str) -> str: return "detect-table-headers" if '"answer"' in normalized_prompt and '"text" or "image"' in normalized_prompt: return "judge-image-type" + if ( + "you will receive a single image extracted from a document" in normalized_prompt + and '"chart"' in normalized_prompt + ): + return "summary-images" if ( '"title"' in normalized_prompt - and '"keywords"' in normalized_prompt and '"summary"' in normalized_prompt + and '"entities"' in normalized_prompt ): return "summary-full" - if ( - 'json dictionary format with key "answer"' in normalized_prompt - and "keywords" in normalized_prompt - ): - return "summary-keywords" if ( "json array only" in normalized_prompt and "toc, not body text" in normalized_prompt @@ -136,21 +136,15 @@ def _detect_mock_task(prompt_text: str) -> str: return "eval-headings" if "scanned page from an engineering atlas" in normalized_prompt: return "atlas-page-info" - if "perform ocr operation" in normalized_prompt: - return "ocr-image" - if ( - "you will receive an image from a document" in normalized_prompt - and "identify the image type" in normalized_prompt - ): - return "summary-images" + if "you are transcribing a document page or image" in normalized_prompt: + return "transcribe" + if "annotating a single rendered document page" in normalized_prompt: + return "page-memory-vlm-tag" + if "you are summarizing one section of a document" in normalized_prompt: + return "page-memory-node-summary" if "summaries of sub-sections from a document section" in normalized_prompt: return "file-summary" - if ( - "line 1: output a short title" in normalized_prompt - and "line 2 onward" in normalized_prompt - ): - return "summary-titled" if "extract the main content of the material" in normalized_prompt: return "summary" @@ -291,16 +285,23 @@ def _build_mock_response(task_name: str) -> str: "detect-toc-range": '{"toc_start": null, "toc_end": null, "confidence": "low"}', "detect-table-headers": '{"answer": [0]}', "judge-image-type": '{"answer": "image"}', - "summary-full": '{"title": "Mock Title", "keywords": "mock", "summary": "Mock summary"}', - "summary-keywords": '{"answer": "mock"}', + "summary-full": ( + '{"title": "Mock Title", "summary": "Mock summary", ' + '"entities": [], "chart": null}' + ), "eval-headings": "[]", "eval-toc-headings": "[]", "atlas-page-info": "Mock atlas page info", - "ocr-image": "Mock OCR text", - "summary-images": "Mock Image Title\nMock image summary", - + "transcribe": '{"text": "Mock transcribed text"}', + "summary-images": ( + '{"title": "Mock Image Title", "summary": "Mock image summary", ' + '"entities": [], "chart": null}' + ), + "page-memory-vlm-tag": '{"summary": "Mock page summary", "entities": []}', + "page-memory-node-summary": ( + '{"summary": "Mock section summary", "entities": []}' + ), "file-summary": "Mock section summary", - "summary-titled": "Mock Title\nMock summary", "summary": "Mock summary", "default": "Mock LLM response", } diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index debc79eb..0c5ba58f 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -63,6 +63,81 @@ def _language_directive(lang) -> str: return "" +# ────────────────────────────────────────────────────────────────────────────── +# Entity-extraction & chart-numeric directives (§4.3 / §4.4) +# ────────────────────────────────────────────────────────────────────────────── +# These build the shared, GENERAL-PURPOSE instructions injected into every +# summary prompt. The type vocabulary is read from the ``ENTITY_TYPES`` config so +# it can be extended without editing prompts, and the wording deliberately avoids +# baked-in examples, sample values, or magic counts — extraction must generalize +# across arbitrary documents, not fit any one corpus. + + +def _entity_types() -> list[str]: + """The configured entity type vocabulary (lower-cased, de-duplicated).""" + from shared.core.config import settings + + raw = getattr(settings, "ENTITY_TYPES", "") or "" + seen: dict[str, None] = {} + for part in raw.split(","): + label = part.strip().lower() + if label and label not in seen: + seen[label] = None + return list(seen.keys()) + + +def _entity_instruction() -> str: + """Build the ``entities`` field instruction from the configured vocabulary. + + Returns a JSON-field directive that asks for typed entities and treats an + empty result as valid. No entity names or counts are hard-coded; the only + corpus-specific input is the configurable type list. + """ + types = _entity_types() + if types: + type_clause = ( + "Set \"type\" to the single best-fitting label from this allowed list: " + + ", ".join(types) + + ". If an entity fits none of them, omit that entity." + ) + else: + type_clause = ( + 'Set "type" to a short lower-case category label you judge appropriate.' + ) + return ( + '- "entities": a JSON array of the salient named entities explicitly ' + "present in the content. Each element is an object with keys \"text\" and " + '"type". Use the exact surface form from the content for "text". ' + f"{type_clause} " + "Do not infer, translate, or invent entities. Return an empty array [] " + "when none are present — an empty result is valid and expected, so never " + "force extraction." + ) + + +def _chart_instruction() -> str: + """Build the optional ``chart`` numeric-extraction directive (§4.3). + + General-purpose: it tells the model to read whatever quantitative extremes and + feature values are actually shown, without prescribing any specific metric, + unit, count, or example value. + """ + return ( + '- "chart": include this object ONLY when the content is a statistical ' + "chart or a data table carrying quantitative values; otherwise omit the " + "key or set it to null. When present it has the shape " + '{"metric": , "extremes": [{"label": , "value": , "kind": "max"|"min"|' + '"peak"}], "features": [{"label": , "value": }], "period":