From 6542e9e27f655b0f8360547b5364b8346ae7b17a Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 16 Jun 2026 20:04:35 +0800 Subject: [PATCH 1/5] Enhance document agent with new page_locate stage and lazy loading of ProfileAgent. - Added "page_locate" to BudgetStage for improved budget management. - Introduced lazy loading for ProfileAgent to optimize imports in document_agent. - Updated tools initialization to include page_locate functionality. --- .../app/services/document_agent/__init__.py | 14 +- .../app/services/document_agent/budget.py | 1 + .../services/document_agent/coordinator.py | 6 + .../document_agent/structure/__init__.py | 2 + .../structure/hierarchy_locator.py | 924 ++++++++++++++++++ .../structure/page_locate_agent.py | 358 +++++++ .../structure/page_locate_subagent.py | 440 +++++++++ .../structure/page_locate_tools.py | 156 +++ .../services/document_agent/tools/__init__.py | 1 + .../document_agent/tools/page_locate.py | 16 + .../page_memory/skeleton_extractor.py | 357 +++++++ 11 files changed, 2274 insertions(+), 1 deletion(-) create mode 100644 apps/worker/app/services/document_agent/structure/__init__.py create mode 100644 apps/worker/app/services/document_agent/structure/hierarchy_locator.py create mode 100644 apps/worker/app/services/document_agent/structure/page_locate_agent.py create mode 100644 apps/worker/app/services/document_agent/structure/page_locate_subagent.py create mode 100644 apps/worker/app/services/document_agent/structure/page_locate_tools.py create mode 100644 apps/worker/app/services/document_agent/tools/page_locate.py create mode 100644 apps/worker/app/services/page_memory/skeleton_extractor.py diff --git a/apps/worker/app/services/document_agent/__init__.py b/apps/worker/app/services/document_agent/__init__.py index ee28ee9f6..59e0afc99 100644 --- a/apps/worker/app/services/document_agent/__init__.py +++ b/apps/worker/app/services/document_agent/__init__.py @@ -1,12 +1,24 @@ """Page anatomy agent for hierarchy-first PDF profiling.""" +from typing import TYPE_CHECKING + from app.services.document_agent.manifest import ( PageAnatomyMap, PageFeature, PageLabel, ShardPlan, ) -from app.services.document_agent.profile_agent import ProfileAgent + +if TYPE_CHECKING: + from app.services.document_agent.profile_agent import ProfileAgent + + +def __getattr__(name: str): + if name == "ProfileAgent": + from app.services.document_agent.profile_agent import ProfileAgent + + return ProfileAgent + raise AttributeError(name) __all__ = [ "PageAnatomyMap", diff --git a/apps/worker/app/services/document_agent/budget.py b/apps/worker/app/services/document_agent/budget.py index 0e7320272..718130e6c 100644 --- a/apps/worker/app/services/document_agent/budget.py +++ b/apps/worker/app/services/document_agent/budget.py @@ -10,6 +10,7 @@ "toc_confirm", "coarse_planner", "structural_react", + "page_locate", "page_tagging", ] diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 9d871866c..38f5093b8 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -64,6 +64,12 @@ def __init__( ), cap=int(os.environ.get("PARSE_AGENT_STRUCTURAL_REACT_CAP", "64000")), ), + "page_locate": StageEnvelope( + min_guarantee=int( + os.environ.get("PARSE_AGENT_PAGE_LOCATE_MIN_BUDGET", "0") + ), + cap=int(os.environ.get("PARSE_AGENT_PAGE_LOCATE_CAP", "0")) or None, + ), "page_tagging": StageEnvelope( min_guarantee=int( os.environ.get("PARSE_AGENT_PAGE_TAGGING_MIN_BUDGET", "0") diff --git a/apps/worker/app/services/document_agent/structure/__init__.py b/apps/worker/app/services/document_agent/structure/__init__.py new file mode 100644 index 000000000..df50ed481 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/__init__.py @@ -0,0 +1,2 @@ +"""Shared document-structure helpers for profile-time anatomy.""" + diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py new file mode 100644 index 000000000..db2fd1c0c --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -0,0 +1,924 @@ +"""Locate hierarchy titles on PDF pages and resolve page ranges. + +This module is intentionally deterministic: it performs strict title anchoring, +candidate collection, and range assembly. The page-memory residual agent calls +into these primitives for grep-like tools and adds VLM verification outside this +module. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Any, Literal + +from app.services.document_parser.structure.body_boundary import ( + clean_toc_title, + normalize_heading_text, +) + +TitleMatchSource = Literal[ + "exact", + "anchored", + "page_compact", + "normalized", + "token", + "printed_prior", + "h1_result", + "agent_vlm", + "agent_heuristic", +] + + +@dataclass(frozen=True) +class PageRange: + start: int + end: int + + def pages(self) -> list[int]: + if self.end < self.start: + return [] + return list(range(self.start, self.end + 1)) + + +@dataclass(frozen=True) +class TitleNode: + title: str + level: int + printed_page: int | None = None + physical_page_hint: int | None = None + children: list["TitleNode"] = field(default_factory=list) + + +@dataclass(frozen=True) +class TitleMatch: + page: int + confidence: float + source: TitleMatchSource + matched_line: str + score: float + candidates: list[int] + evidence: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ResolvedHierarchyRange: + title: str + level: int + start_page: int + end_page: int + path_titles: tuple[str, ...] + match: TitleMatch | None + evidence: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class _LineHit: + page: int + line_index: int + line: str + source: TitleMatchSource + score: float + + +_STOPWORDS = { + "a", + "an", + "and", + "are", + "as", + "for", + "in", + "of", + "on", + "or", + "the", + "to", + "with", +} + + +def locate_title_start_page( + title: str, + *, + scope_pages: list[int], + page_texts: dict[int, str], + printed_page: int | None = None, + page_offset_hint: int | None = None, +) -> TitleMatch | None: + """Locate *title* using deterministic weak evidence. + + This is a candidate-gathering primitive. The C4 page-memory path should only + directly accept :func:`locate_title_strict_exact`; weak results from this + function are meant for residual agent/VLM arbitration. + """ + matches = collect_title_candidate_matches( + title, + scope_pages=scope_pages, + page_texts=page_texts, + printed_page=printed_page, + page_offset_hint=page_offset_hint, + ) + return matches[0] if matches else None + + +def locate_title_strict_exact( + title: str, + *, + scope_pages: list[int], + page_texts: dict[int, str], +) -> TitleMatch | None: + """Return a direct anchor only when a cleaned heading line has one page hit.""" + hits = _find_anchored_hits(title, scope_pages, page_texts) + pages = sorted({hit.page for hit in hits}) + if len(pages) != 1: + return None + return _choose_best_hit( + hits, + source="anchored", + printed_page=None, + page_offset_hint=None, + extra_evidence={"accept": "strict_exact_unique"}, + ) + + +def collect_title_candidate_matches( + title: str, + *, + scope_pages: list[int], + page_texts: dict[int, str], + printed_page: int | None = None, + page_offset_hint: int | None = None, + limit: int | None = None, +) -> list[TitleMatch]: + """Collect grep-style candidate pages for a title without final arbitration.""" + normalized_title = normalize_heading_text(title) + if not normalized_title or not scope_pages: + return [] + + hits: list[_LineHit] = [] + for _source, finder in ( + ("anchored", _find_anchored_hits), + ("page_compact", _find_page_compact_hits), + ("normalized", _find_normalized_hits), + ("token", _find_token_hits), + ): + hits.extend(finder(normalized_title, scope_pages, page_texts)) + + by_page: dict[int, list[_LineHit]] = {} + for hit in hits: + by_page.setdefault(hit.page, []).append(hit) + + matches = [ + _choose_best_hit( + page_hits, + source=_preferred_source(page_hits), + printed_page=printed_page, + page_offset_hint=page_offset_hint, + ) + for page_hits in by_page.values() + ] + + prior_page = _resolve_printed_prior( + printed_page=printed_page, + page_offset_hint=page_offset_hint, + scope_pages=scope_pages, + ) + if prior_page is not None and prior_page not in by_page: + matches.append( + TitleMatch( + page=prior_page, + confidence=0.35, + source="printed_prior", + matched_line="", + score=0.35, + candidates=[prior_page], + evidence={ + "printed_page": printed_page, + "page_offset_hint": page_offset_hint, + }, + ) + ) + + matches.sort( + key=lambda match: ( + match.score, + match.confidence, + -abs(match.page - (printed_page + page_offset_hint)) + if printed_page is not None and page_offset_hint is not None + else 0, + -match.page, + ), + reverse=True, + ) + if limit is not None: + return matches[: max(int(limit), 0)] + return matches + + +def resolve_hierarchy_page_ranges( + nodes: list[TitleNode], + *, + page_count: int, + page_texts: dict[int, str], + body_pages: list[int] | None = None, + page_offset_hint: int | None = None, + match_overrides: dict[tuple[str, ...], TitleMatch] | None = None, + use_weak_fallback: bool = False, +) -> list[ResolvedHierarchyRange]: + """Resolve leaf hierarchy nodes into closed page ranges. + + The emitted ranges are leaf-first and intentionally closed-closed: if the + next leaf starts on page N, the previous leaf may also include page N. This + preserves page-to-section many-to-many mapping for dense documents. + """ + if page_count <= 0 or not nodes: + return [] + + pages = sorted(set(body_pages or list(range(1, page_count + 1)))) + pages = [page for page in pages if 1 <= page <= page_count] + if not pages: + return [] + + allowed_pages = set(pages) + scope = PageRange(start=pages[0], end=pages[-1]) + resolved: list[ResolvedHierarchyRange] = [] + _resolve_siblings( + nodes, + parent_scope=scope, + allowed_pages=allowed_pages, + parent_titles=(), + page_texts=page_texts, + page_offset_hint=page_offset_hint, + match_overrides=match_overrides or {}, + use_weak_fallback=use_weak_fallback, + resolved=resolved, + ) + return resolved + + +def extract_toc_nodes(toc_hierarchies: list[dict[str, Any]] | None) -> list[TitleNode]: + """Build a title tree from supported TOC hierarchy payloads.""" + flat_entries: list[dict[str, Any]] = [] + for hierarchy in toc_hierarchies or []: + entries = _extract_flat_entries(hierarchy.get("toc_with_level")) + if not entries and hierarchy.get("toc_tree"): + entries = _flatten_tree_entries(hierarchy["toc_tree"]) + flat_entries.extend(entries) + return _entries_to_tree(flat_entries) + + +def _resolve_siblings( + nodes: list[TitleNode], + *, + parent_scope: PageRange, + allowed_pages: set[int], + parent_titles: tuple[str, ...], + page_texts: dict[int, str], + page_offset_hint: int | None, + match_overrides: dict[tuple[str, ...], TitleMatch], + use_weak_fallback: bool, + resolved: list[ResolvedHierarchyRange], +) -> None: + located: list[tuple[TitleNode, int, TitleMatch | None]] = [] + lower_bound = parent_scope.start + + for index, node in enumerate(nodes): + path_titles = (*parent_titles, node.title) + pages = _allowed_pages_between(lower_bound, parent_scope.end, allowed_pages) + match = _match_override(path_titles, match_overrides, pages) + if match is None: + match = _match_physical_hint(node=node, scope_pages=pages) + if match is None: + match = locate_title_strict_exact( + node.title, + scope_pages=pages, + page_texts=page_texts, + ) + if match is None and use_weak_fallback: + match = locate_title_start_page( + node.title, + scope_pages=pages, + page_texts=page_texts, + printed_page=node.printed_page, + page_offset_hint=page_offset_hint, + ) + if match is None: + start_page = lower_bound + else: + start_page = max(parent_scope.start, min(match.page, parent_scope.end)) + located.append((node, start_page, match)) + if match is not None: + lower_bound = start_page + elif index + 1 < len(nodes): + next_match = _find_next_located_sibling( + nodes=nodes, + start_index=index + 1, + lower_bound=lower_bound, + parent_end=parent_scope.end, + allowed_pages=allowed_pages, + page_texts=page_texts, + page_offset_hint=page_offset_hint, + match_overrides=match_overrides, + use_weak_fallback=use_weak_fallback, + parent_titles=parent_titles, + ) + if next_match is not None: + lower_bound = next_match.page + + for index, (node, start_page, match) in enumerate(located): + next_start = _next_located_start(located, index + 1) + end_page = next_start if next_start is not None else parent_scope.end + if end_page < start_page: + end_page = start_page + + path_titles = (*parent_titles, node.title) + evidence = _range_evidence(match) + if match is None: + evidence.update( + _unlocated_warning_evidence( + title=node.title, + path_titles=path_titles, + start_page=start_page, + end_page=end_page, + parent_scope=parent_scope, + ) + ) + + if node.children: + _resolve_siblings( + node.children, + parent_scope=PageRange(start_page, end_page), + allowed_pages=allowed_pages, + parent_titles=path_titles, + page_texts=page_texts, + page_offset_hint=page_offset_hint, + match_overrides=match_overrides, + use_weak_fallback=use_weak_fallback, + resolved=resolved, + ) + continue + + resolved.append( + ResolvedHierarchyRange( + title=node.title, + level=node.level, + start_page=start_page, + end_page=end_page, + path_titles=path_titles, + match=match, + evidence=evidence, + ) + ) + + +def _find_next_located_sibling( + *, + nodes: list[TitleNode], + start_index: int, + lower_bound: int, + parent_end: int, + allowed_pages: set[int], + page_texts: dict[int, str], + page_offset_hint: int | None, + match_overrides: dict[tuple[str, ...], TitleMatch], + use_weak_fallback: bool, + parent_titles: tuple[str, ...], +) -> TitleMatch | None: + pages = _allowed_pages_between(lower_bound, parent_end, allowed_pages) + for sibling in nodes[start_index:]: + path_titles = (*parent_titles, sibling.title) + match = _match_override(path_titles, match_overrides, pages) + if match is None: + match = _match_physical_hint(node=sibling, scope_pages=pages) + if match is not None: + return match + match = locate_title_strict_exact( + sibling.title, + scope_pages=pages, + page_texts=page_texts, + ) + if match is None and use_weak_fallback: + match = locate_title_start_page( + sibling.title, + scope_pages=pages, + page_texts=page_texts, + printed_page=sibling.printed_page, + page_offset_hint=page_offset_hint, + ) + if match is not None: + return match + return None + + +def _match_override( + path_titles: tuple[str, ...], + match_overrides: dict[tuple[str, ...], TitleMatch], + scope_pages: list[int], +) -> TitleMatch | None: + match = match_overrides.get(path_titles) + if match is None or match.page not in scope_pages: + return None + return match + + +def iter_leaf_title_nodes( + nodes: list[TitleNode], + *, + parent_titles: tuple[str, ...] = (), +) -> list[tuple[tuple[str, ...], TitleNode]]: + leaves: list[tuple[tuple[str, ...], TitleNode]] = [] + for node in nodes: + path_titles = (*parent_titles, node.title) + if node.children: + leaves.extend(iter_leaf_title_nodes(node.children, parent_titles=path_titles)) + else: + leaves.append((path_titles, node)) + return leaves + + +def max_title_depth(nodes: list[TitleNode]) -> int: + if not nodes: + return 0 + return max( + max(node.level, max_title_depth(node.children)) + if node.children + else node.level + for node in nodes + ) + + +def prune_title_nodes_for_emit_depth( + nodes: list[TitleNode], + *, + emit_depth: int, +) -> list[TitleNode]: + pruned: list[TitleNode] = [] + for node in nodes: + if node.level >= emit_depth or not node.children: + pruned.append( + TitleNode( + title=node.title, + level=node.level, + printed_page=node.printed_page, + physical_page_hint=node.physical_page_hint, + children=[], + ) + ) + continue + pruned.append( + TitleNode( + title=node.title, + level=node.level, + printed_page=node.printed_page, + physical_page_hint=node.physical_page_hint, + children=prune_title_nodes_for_emit_depth( + node.children, + emit_depth=emit_depth, + ), + ) + ) + return pruned + + +def _next_located_start( + located: list[tuple[TitleNode, int, TitleMatch | None]], + start_index: int, +) -> int | None: + for _node, start_page, match in located[start_index:]: + if match is not None: + return start_page + return None + + +def _range_evidence(match: TitleMatch | None) -> dict[str, Any]: + if match is None: + return {"source": "unlocated", "confidence": 0.0, "candidates": []} + return { + "source": match.source, + "confidence": match.confidence, + "matched_line": match.matched_line, + "candidates": match.candidates, + "score": match.score, + **match.evidence, + } + + +def _unlocated_warning_evidence( + *, + title: str, + path_titles: tuple[str, ...], + start_page: int, + end_page: int, + parent_scope: PageRange, +) -> dict[str, Any]: + warning = { + "code": "section_title_unlocated", + "title": title, + "path_titles": list(path_titles), + "assigned_range": [start_page, end_page], + "parent_scope": [parent_scope.start, parent_scope.end], + "message": ( + "Section title was not found on any allowed body page; assigned range " + "from neighboring hierarchy boundaries." + ), + } + return { + "status": "inherited_unlocated", + "warning": warning, + "warnings": [warning], + } + + +def _match_physical_hint( + *, + node: TitleNode, + scope_pages: list[int], +) -> TitleMatch | None: + if node.physical_page_hint is None or node.physical_page_hint not in scope_pages: + return None + return TitleMatch( + page=node.physical_page_hint, + confidence=0.88, + source="h1_result", + matched_line="", + score=0.88, + candidates=[node.physical_page_hint], + evidence={"physical_page_hint": node.physical_page_hint}, + ) + + +def _allowed_pages_between(start: int, end: int, allowed_pages: set[int]) -> list[int]: + if end < start: + return [] + return [page for page in range(start, end + 1) if page in allowed_pages] + + +def _find_exact_hits( + title: str, + scope_pages: list[int], + page_texts: dict[int, str], +) -> list[_LineHit]: + hits: list[_LineHit] = [] + needle = normalize_heading_text(title).casefold() + for page, line_index, line in _iter_lines(scope_pages, page_texts): + normalized_line = normalize_heading_text(line).casefold() + if needle and needle in normalized_line: + base = 1.0 + cleaned_line = normalize_heading_text(clean_toc_title(line)).casefold() + if normalized_line == needle or cleaned_line == needle: + base = 1.18 + hits.append( + _LineHit( + page=page, + line_index=line_index, + line=line.strip(), + source="exact", + score=_line_score(line=line, line_index=line_index, base=base), + ) + ) + return hits + + +def _find_anchored_hits( + title: str, + scope_pages: list[int], + page_texts: dict[int, str], +) -> list[_LineHit]: + hits: list[_LineHit] = [] + needle = normalize_heading_text(clean_toc_title(title) or title).casefold() + if not needle: + return hits + for page, line_index, line in _iter_lines(scope_pages, page_texts): + cleaned_line = normalize_heading_text(clean_toc_title(line)).casefold() + if cleaned_line == needle: + hits.append( + _LineHit( + page=page, + line_index=line_index, + line=line.strip(), + source="anchored", + score=_line_score(line=line, line_index=line_index, base=0.96), + ) + ) + return hits + + +def _find_page_compact_hits( + title: str, + scope_pages: list[int], + page_texts: dict[int, str], +) -> list[_LineHit]: + hits: list[_LineHit] = [] + needles = _compact_title_variants(title) + if not needles: + return hits + + for page in scope_pages: + raw_text = page_texts.get(page, "") + compact_text = _compact_match_text(raw_text) + if not compact_text: + continue + matched_needle = next((needle for needle in needles if needle in compact_text), None) + if matched_needle is None: + continue + line_index, evidence = _compact_match_evidence(raw_text, matched_needle) + hits.append( + _LineHit( + page=page, + line_index=line_index, + line=evidence, + source="page_compact", + score=_line_score(line=evidence, line_index=line_index, base=0.94), + ) + ) + return hits + + +def _find_normalized_hits( + title: str, + scope_pages: list[int], + page_texts: dict[int, str], +) -> list[_LineHit]: + hits: list[_LineHit] = [] + needle = normalize_heading_text(clean_toc_title(title)).casefold() + if len(needle) < 2: + return hits + for page, line_index, line in _iter_lines(scope_pages, page_texts): + cleaned_line = normalize_heading_text(clean_toc_title(line)).casefold() + if not cleaned_line: + continue + if needle in cleaned_line or _is_strong_reverse_match(cleaned_line, needle): + hits.append( + _LineHit( + page=page, + line_index=line_index, + line=line.strip(), + source="normalized", + score=_line_score(line=line, line_index=line_index, base=0.9), + ) + ) + return hits + + +def _find_token_hits( + title: str, + scope_pages: list[int], + page_texts: dict[int, str], +) -> list[_LineHit]: + title_tokens = _significant_tokens(clean_toc_title(title) or title) + if not title_tokens: + return [] + + hits: list[_LineHit] = [] + for page, line_index, line in _iter_lines(scope_pages, page_texts): + line_tokens = _significant_tokens(line) + if not line_tokens: + continue + coverage = len(title_tokens & line_tokens) / len(title_tokens) + if coverage < 0.8: + continue + hits.append( + _LineHit( + page=page, + line_index=line_index, + line=line.strip(), + source="token", + score=_line_score(line=line, line_index=line_index, base=0.78) + + coverage, + ) + ) + return hits + + +def _choose_best_hit( + hits: list[_LineHit], + *, + source: TitleMatchSource, + printed_page: int | None, + page_offset_hint: int | None, + extra_evidence: dict[str, Any] | None = None, +) -> TitleMatch: + expected_page = ( + printed_page + page_offset_hint + if printed_page is not None and page_offset_hint is not None + else None + ) + + def sort_key(hit: _LineHit) -> tuple[float, int, int, int]: + printed_bonus = 0 + if expected_page is not None: + printed_bonus = -abs(hit.page - expected_page) + return (hit.score, printed_bonus, -hit.line_index, -hit.page) + + ordered = sorted(hits, key=sort_key, reverse=True) + best = ordered[0] + pages = sorted({hit.page for hit in ordered}) + confidence_by_source = { + "exact": 0.95, + "anchored": 0.92, + "page_compact": 0.9, + "normalized": 0.84, + "token": 0.72, + "printed_prior": 0.35, + "h1_result": 0.88, + } + return TitleMatch( + page=best.page, + confidence=confidence_by_source[source], + source=source, + matched_line=best.line[:160], + score=best.score, + candidates=pages, + evidence={ + "line_index": best.line_index, + "candidate_count": len(pages), + "printed_page": printed_page, + "page_offset_hint": page_offset_hint, + **(extra_evidence or {}), + }, + ) + + +def _preferred_source(hits: list[_LineHit]) -> TitleMatchSource: + priority = { + "anchored": 50, + "page_compact": 40, + "normalized": 30, + "token": 20, + "printed_prior": 10, + "exact": 5, + "h1_result": 60, + "agent_vlm": 70, + "agent_heuristic": 15, + } + return max(hits, key=lambda hit: (priority.get(hit.source, 0), hit.score)).source + + +def _resolve_printed_prior( + *, + printed_page: int | None, + page_offset_hint: int | None, + scope_pages: list[int], +) -> int | None: + if printed_page is None or page_offset_hint is None or not scope_pages: + return None + page = printed_page + page_offset_hint + if page in scope_pages: + return page + return None + + +def _line_score(*, line: str, line_index: int, base: float) -> float: + stripped = normalize_heading_text(line) + short_line_bonus = max(0.0, 1.0 - (len(stripped) / 140.0)) + top_bonus = max(0.0, 1.0 - (line_index / 18.0)) + return base + short_line_bonus * 0.12 + top_bonus * 0.1 + + +def _iter_lines( + scope_pages: list[int], + page_texts: dict[int, str], +) -> list[tuple[int, int, str]]: + rows: list[tuple[int, int, str]] = [] + for page in scope_pages: + for line_index, line in enumerate(page_texts.get(page, "").splitlines()): + if line.strip(): + rows.append((page, line_index, line)) + return rows + + +def _compact_title_variants(title: str) -> list[str]: + normalized = normalize_heading_text(clean_toc_title(title) or title).casefold() + compacted: list[str] = [] + compact = _compact_match_text(normalized) + if compact: + compacted.append(compact) + return compacted + + +def _compact_match_text(text: str) -> str: + return re.sub(r"\s+", "", normalize_heading_text(text)).casefold() + + +def _compact_match_evidence(raw_text: str, compact_needle: str) -> tuple[int, str]: + lines = [line.strip() for line in raw_text.splitlines() if line.strip()] + if not lines: + return 0, "" + + compact_so_far = "" + start_index = 0 + for index, line in enumerate(lines): + line_compact = _compact_match_text(line) + if not compact_so_far: + start_index = index + compact_so_far += line_compact + if compact_needle in compact_so_far: + return start_index, " ".join(lines[start_index : index + 1])[:160] + if len(compact_so_far) > len(compact_needle) * 3: + compact_so_far = line_compact + start_index = index + + return 0, " ".join(lines[:3])[:160] + + +def _is_strong_reverse_match(fragment: str, title: str) -> bool: + return len(fragment) >= 6 and fragment in title + + +def _significant_tokens(text: str) -> set[str]: + normalized = normalize_heading_text(clean_toc_title(text) or text).casefold() + latin = { + token + for token in re.findall(r"[a-z0-9][a-z0-9_-]+", normalized) + if token not in _STOPWORDS + } + cjk = set(re.findall(r"[\u4e00-\u9fff]", normalized)) + return latin | cjk + + +def _extract_flat_entries(payload: Any) -> list[dict[str, Any]]: + if isinstance(payload, list): + return [ + entry + for entry in payload + if isinstance(entry, dict) and entry.get("heading") + ] + if not isinstance(payload, str): + return [] + return _parse_markdown_toc_entries(payload) + + +def _parse_markdown_toc_entries(markdown: str) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] + headers: list[str] | None = None + for raw_line in markdown.splitlines(): + line = raw_line.strip() + if not line.startswith("|") or not line.endswith("|"): + continue + cells = [cell.strip() for cell in line.strip("|").split("|")] + if not cells or all(set(cell) <= {"-", ":"} for cell in cells): + continue + if headers is None: + headers = [cell.lower() for cell in cells] + continue + row = dict(zip(headers, cells)) + level = _safe_int(row.get("level")) + heading = row.get("heading") + if heading and level: + entries.append( + { + "heading": heading, + "level": level, + "page_number": _safe_int(row.get("page_number")), + } + ) + return entries + + +def _flatten_tree_entries( + tree: dict[str, Any], + *, + level: int = 1, +) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] + for title, children in tree.items(): + entries.append({"heading": title, "level": level}) + if isinstance(children, dict): + entries.extend(_flatten_tree_entries(children, level=level + 1)) + return entries + + +def _entries_to_tree(entries: list[dict[str, Any]]) -> list[TitleNode]: + roots: list[TitleNode] = [] + stack: list[tuple[int, TitleNode]] = [] + + for entry in entries: + raw_title = str(entry.get("heading") or "").strip() + title = clean_toc_title(raw_title) or normalize_heading_text(raw_title) + level = _safe_int(entry.get("level")) or 1 + if not title or len(title) < 2: + continue + node = TitleNode( + title=title, + level=level, + printed_page=_safe_int(entry.get("page_number")), + ) + while stack and stack[-1][0] >= level: + stack.pop() + if stack: + stack[-1][1].children.append(node) + else: + roots.append(node) + stack.append((level, node)) + + return roots + + +def _safe_int(value: Any) -> int | None: + if value is None or value == "": + return None + try: + return int(value) + except (TypeError, ValueError): + return None 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 new file mode 100644 index 000000000..5aecab463 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/page_locate_agent.py @@ -0,0 +1,358 @@ +"""Residual page-location agent for hierarchy titles.""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from dataclasses import dataclass, field +from typing import Any, cast + +from loguru import logger + +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.structure.hierarchy_locator import ( + TitleMatch, + TitleNode, + collect_title_candidate_matches, + iter_leaf_title_nodes, + locate_title_strict_exact, + max_title_depth, + prune_title_nodes_for_emit_depth, +) +from app.services.document_agent.structure.page_locate_subagent import ( + PageLocateSubAgent, + SubAgentConfig, +) + + +@dataclass(frozen=True) +class PageLocateConfig: + residual_agent_limit: int = 50 + max_emit_depth: int = 5 + min_emit_depth: int = 2 + vlm_candidate_page_cap: int = 4 + full_leaf_sections: bool = False + + @classmethod + def from_env(cls) -> "PageLocateConfig": + return cls( + residual_agent_limit=int(os.environ.get("PAGE_MEMORY_RESIDUAL_AGENT_LIMIT", "50")), + max_emit_depth=int(os.environ.get("PAGE_MEMORY_MAX_EMIT_DEPTH", "5")), + min_emit_depth=int(os.environ.get("PAGE_MEMORY_MIN_EMIT_DEPTH", "2")), + vlm_candidate_page_cap=int( + os.environ.get("PAGE_MEMORY_VLM_CANDIDATE_PAGE_CAP", "4") + ), + full_leaf_sections=os.environ.get("PAGE_MEMORY_FULL_LEAF_SECTIONS", "").lower() + in {"1", "true", "yes", "on"}, + ) + + +@dataclass(frozen=True) +class ResidualRequest: + path_titles: tuple[str, ...] + node: TitleNode + + +@dataclass(frozen=True) +class PageLocatePrepareResult: + nodes: list[TitleNode] + match_overrides: dict[tuple[str, ...], TitleMatch] + summary: dict[str, Any] = field(default_factory=dict) + + +class PageLocateResidualAgent: + """Batch residual resolver for C4 title-to-page anchoring.""" + + def __init__( + self, + *, + ctx: ToolContext | None, + page_texts: dict[int, str], + body_pages: list[int], + page_count: int, + page_offset_hint: int | None, + config: PageLocateConfig | None = None, + ) -> None: + self.ctx = ctx + self.page_texts = page_texts + self.body_pages = sorted({page for page in body_pages if 1 <= page <= page_count}) + self.page_count = page_count + self.page_offset_hint = page_offset_hint + self.config = config or PageLocateConfig.from_env() + + def prepare(self, nodes: list[TitleNode]) -> PageLocatePrepareResult: + if not nodes: + return PageLocatePrepareResult(nodes=[], match_overrides={}, summary={}) + + actual_max_depth = max_title_depth(nodes) + emit_depth = min(actual_max_depth, self.config.max_emit_depth) + emit_depth = max(emit_depth, self.config.min_emit_depth) + selected_nodes = nodes + direct_matches: dict[tuple[str, ...], TitleMatch] = {} + residuals: list[ResidualRequest] = [] + + while True: + selected_nodes = prune_title_nodes_for_emit_depth(nodes, emit_depth=emit_depth) + direct_matches, residuals = self._classify_residuals(selected_nodes) + if ( + self.config.full_leaf_sections + or len(residuals) <= self.config.residual_agent_limit + or emit_depth <= self.config.min_emit_depth + ): + break + emit_depth -= 1 + + residual_matches = self._resolve_residuals(residuals) + match_overrides = {**direct_matches, **residual_matches} + summary = { + "agent": "page_locate_residual", + "emit_depth": emit_depth, + "actual_max_depth": actual_max_depth, + "direct_exact_count": len(direct_matches), + "residual_count": len(residuals), + "resolved_residual_count": len(residual_matches), + "residual_limit": self.config.residual_agent_limit, + "vlm_candidate_page_cap": self.config.vlm_candidate_page_cap, + "full_leaf_sections": self.config.full_leaf_sections, + } + logger.info("[page_locate.agent] summary={}", summary) + return PageLocatePrepareResult( + nodes=selected_nodes, + match_overrides=match_overrides, + summary=summary, + ) + + def _classify_residuals( + self, + nodes: list[TitleNode], + ) -> tuple[dict[tuple[str, ...], TitleMatch], list[ResidualRequest]]: + direct: dict[tuple[str, ...], TitleMatch] = {} + residuals: list[ResidualRequest] = [] + for path_titles, node in iter_leaf_title_nodes(nodes): + match = locate_title_strict_exact( + node.title, + scope_pages=self.body_pages, + page_texts=self.page_texts, + ) + if match is not None: + direct[path_titles] = TitleMatch( + page=match.page, + confidence=match.confidence, + source=match.source, + matched_line=match.matched_line, + score=match.score, + candidates=match.candidates, + evidence={ + **match.evidence, + "page_locate_agent": { + "decision": "direct_strict_exact", + "path_titles": list(path_titles), + }, + }, + ) + else: + residuals.append(ResidualRequest(path_titles=path_titles, node=node)) + return direct, residuals + + def _resolve_residuals( + self, + residuals: list[ResidualRequest], + ) -> dict[tuple[str, ...], TitleMatch]: + matches: dict[tuple[str, ...], TitleMatch] = {} + if not residuals or self.ctx is None: + # No agent runtime (ctx/budget) → leave residuals for physical-hint / + # neighbor-boundary fallback in the resolver. The residual sub-agent + # only runs when a real ToolContext is available. + return matches + + sub_config = SubAgentConfig( + candidate_cap=max(self.config.vlm_candidate_page_cap * 2, 4), + verify_page_cap=max(self.config.vlm_candidate_page_cap, 1), + ) + for residual in residuals: + agent = PageLocateSubAgent( + ctx=self.ctx, + scope_pages=self.body_pages, + page_count=self.page_count, + config=sub_config, + page_offset_hint=self.page_offset_hint, + ) + result = agent.locate( + title=residual.node.title, + printed_page=residual.node.printed_page, + ) + if result.match is None: + logger.warning( + "[page_locate.agent] unresolved title={!r} path_titles={} stop={}", + residual.node.title, + residual.path_titles, + result.stop_reason, + ) + continue + base = result.match + agent_evidence = dict(base.evidence.get("page_locate_agent", {})) + agent_evidence["path_titles"] = list(residual.path_titles) + agent_evidence["stop_reason"] = result.stop_reason + matches[residual.path_titles] = TitleMatch( + page=base.page, + confidence=base.confidence, + source=base.source, + matched_line=base.matched_line, + score=base.score, + candidates=base.candidates, + evidence={**base.evidence, "page_locate_agent": agent_evidence}, + ) + return matches + + +def grep_title_page_candidates( + *, + title: str, + scope_pages: list[int], + page_texts: dict[int, str], + printed_page: int | None = None, + page_offset_hint: int | None = None, + limit: int | None = None, +) -> list[TitleMatch]: + return collect_title_candidate_matches( + title, + scope_pages=scope_pages, + page_texts=page_texts, + printed_page=printed_page, + page_offset_hint=page_offset_hint, + limit=limit, + ) + + +def verify_section_page_choice( + *, + ctx: ToolContext | None, + title: str, + candidate_matches: list[TitleMatch], + candidate_page_cap: int, +) -> dict[str, Any]: + candidates = candidate_matches[: max(candidate_page_cap, 1)] + if not candidates: + return { + "selected_page": None, + "confidence": 0.0, + "source": "agent_heuristic", + "reason": "no grep candidates", + } + + model = None + if ctx is not None: + model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") + if ctx is None or not model or ctx.budget is None: + best = candidates[0] + return { + "selected_page": best.page, + "confidence": min(best.confidence, 0.62), + "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( + ctx, + pages, + folder_name="page_locate_pages", + prefix="locate", + timeout=120, + ) + if not rendered: + best = candidates[0] + return { + "selected_page": best.page, + "confidence": min(best.confidence, 0.58), + "source": "agent_heuristic", + "reason": "render failed; selected top grep candidate", + } + + prompt = _build_verify_prompt(title=title, candidates=candidates) + est = 800 * len(rendered) + 800 + stage = "page_locate" + if not ctx.budget.try_reserve("visual", est, stage=stage): + best = candidates[0] + return { + "selected_page": best.page, + "confidence": min(best.confidence, 0.56), + "source": "agent_heuristic", + "reason": "page_locate visual budget exhausted; selected top grep candidate", + } + + content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + for item in rendered: + with open(str(item["png_path"]), "rb") as image_file: + img_b64 = base64.b64encode(image_file.read()).decode() + content_parts.append({"type": "text", "text": f"\n--- Page {item['page']} ---"}) + content_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + + start = time.monotonic() + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=0.0, + max_tokens=400, + response_format={"type": "json_object"}, + usage_task="page_memory.page_locate", + ) + ctx.budget.commit( + "visual", + actual=usage.get("total_tokens", est), + est=est, + stage=stage, + ) + payload = json.loads(raw) + selected_page = payload.get("selected_page") + if selected_page is not None: + selected_page = int(selected_page) + if selected_page not in pages: + selected_page = None + return { + "selected_page": selected_page, + "confidence": float(payload.get("confidence") or 0.75), + "source": "agent_vlm", + "reason": str(payload.get("reason") or ""), + "latency_ms": int((time.monotonic() - start) * 1000), + "tokens_used": usage.get("total_tokens", 0), + } + except Exception as exc: + ctx.budget.refund("visual", est=est, stage=stage) + best = candidates[0] + logger.warning("[page_locate.agent] VLM failed for title={!r}: {}", title, exc) + return { + "selected_page": best.page, + "confidence": min(best.confidence, 0.54), + "source": "agent_heuristic", + "reason": f"VLM failed ({type(exc).__name__}); selected top grep candidate", + } + + +def _build_verify_prompt(*, title: str, candidates: list[TitleMatch]) -> str: + candidate_lines = "\n".join( + f"- page {match.page}: source={match.source}, line={match.matched_line!r}" + for match in candidates + ) + return ( + "You are a page-location sub-agent for a PDF hierarchy parser.\n" + "Choose which candidate page is the true START page of the section title, " + "not a table-of-contents entry, page header, footer, or body-text mention.\n" + f"Section title: {title!r}\n" + f"Candidates:\n{candidate_lines}\n" + "Return strict JSON: {\"selected_page\": number|null, " + "\"confidence\": number, \"reason\": \"brief explanation\"}." + ) diff --git a/apps/worker/app/services/document_agent/structure/page_locate_subagent.py b/apps/worker/app/services/document_agent/structure/page_locate_subagent.py new file mode 100644 index 000000000..b6c3cbfc9 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/page_locate_subagent.py @@ -0,0 +1,440 @@ +"""Per-title ReAct sub-agent for residual hierarchy page location. + +This is intentionally NOT a fixed pipeline. For every title that strict anchoring +could not resolve, an LLM runs a bounded reasoning loop and decides — round by +round — which tool to call: + +* ``grep.title_pages`` — find candidate body pages for a (possibly rewritten) + query. The agent may shorten the title to a distinctive core when the full + title carries trailing document-reference codes / ``《》`` wrappers / version + notes that never appear contiguously in the body, or when the heading is split + across lines. +* ``verify.section_page`` — render candidate pages and ask a VLM which one truly + *starts* the section (vs a TOC row, header/footer, or a body citation). + +The agent then ``submit``s a confirmed start page (or gives up). Tools are +dispatched through the shared registry, so the agent reuses the exact same +grep/VLM primitives the rest of the profile agent uses. +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass, field +from typing import Any, Callable, cast + +from loguru import logger + +from app.services.document_agent.manifest import ToolContext +from app.services.document_agent.registry import REGISTRY +from app.services.document_agent.structure.hierarchy_locator import TitleMatch + +DecideFn = Callable[..., dict[str, Any] | None] + +_ALLOWED_ACTIONS = {"grep", "verify", "submit", "give_up"} + + +@dataclass(frozen=True) +class SubAgentConfig: + max_rounds: int = 6 + candidate_cap: int = 8 + verify_page_cap: int = 4 + decide_max_tokens: int = 500 + + +@dataclass +class SubAgentResult: + match: TitleMatch | None + transcript: list[dict[str, Any]] = field(default_factory=list) + rounds: int = 0 + stop_reason: str = "exhausted" + + +class PageLocateSubAgent: + """Bounded ReAct loop that locates ONE residual title via grep + VLM tools.""" + + def __init__( + self, + *, + ctx: ToolContext | None, + scope_pages: list[int], + page_count: int, + config: SubAgentConfig | None = None, + page_offset_hint: int | None = None, + decide: DecideFn | None = None, + ) -> None: + self.ctx = ctx + self.page_count = page_count + self.scope_pages = sorted({p for p in scope_pages if 1 <= p <= page_count}) + self.config = config or SubAgentConfig() + self.page_offset_hint = page_offset_hint + self._decide = decide or decide_next_action + # Ensure the page-locate tools are registered. Importing the light + # ``structure`` module (not the heavy ``tools`` package) avoids pulling + # in S3/DB-dependent modules and avoids import cycles. + import app.services.document_agent.structure.page_locate_tools # noqa: F401 + + def locate(self, *, title: str, printed_page: int | None = None) -> SubAgentResult: + if not self.scope_pages: + return SubAgentResult(match=None, stop_reason="empty_scope") + + transcript: list[dict[str, Any]] = [] + seen_grep: dict[int, dict[str, Any]] = {} + last_verify: dict[str, Any] | None = None + + for round_index in range(self.config.max_rounds): + decision = self._decide( + ctx=self.ctx, + title=title, + scope_pages=self.scope_pages, + printed_page=printed_page, + page_offset_hint=self.page_offset_hint, + observations=transcript, + seen_pages=sorted(seen_grep.keys()), + last_verify=last_verify, + round_index=round_index, + max_rounds=self.config.max_rounds, + ) + if decision is None: + break + + action = str(decision.get("action") or "") + entry: dict[str, Any] = {"round": round_index, "decision": decision} + transcript.append(entry) + + if action == "grep": + query = str(decision.get("query_title") or title).strip() or title + pages = self._coerce_pages(decision.get("pages")) or self.scope_pages + result = self._dispatch( + "grep.title_pages", + { + "title": query, + "pages": pages, + "candidate_cap": self.config.candidate_cap, + }, + ) + candidates = [] + if result.status == "ok" and result.payload: + candidates = list(result.payload.get("candidates") or []) + for cand in candidates: + seen_grep[int(cand["page"])] = cand + entry["observation"] = { + "tool": "grep.title_pages", + "query": query, + "candidates": candidates[: self.config.candidate_cap], + } + + elif action == "verify": + pages = self._coerce_pages(decision.get("pages")) or sorted(seen_grep.keys()) + pages = pages[: max(self.config.verify_page_cap, 1)] + if not pages: + entry["observation"] = { + "tool": "verify.section_page", + "error": "no candidate pages to verify", + } + continue + result = self._dispatch("verify.section_page", {"title": title, "pages": pages}) + choice = result.payload if (result.status == "ok" and result.payload) else {} + last_verify = choice + entry["observation"] = { + "tool": "verify.section_page", + "pages": pages, + "choice": choice, + } + + elif action == "submit": + selected = decision.get("selected_page") + if selected is None: + return SubAgentResult(None, transcript, round_index + 1, "submit_null") + page = int(selected) + if page not in self.scope_pages: + entry["observation"] = {"error": f"submit page {page} outside scope"} + continue + match = self._build_match( + title=title, + page=page, + seen_grep=seen_grep, + last_verify=last_verify, + decision=decision, + transcript=transcript, + ) + return SubAgentResult(match, transcript, round_index + 1, "submit") + + elif action == "give_up": + return SubAgentResult(None, transcript, round_index + 1, "give_up") + + else: + entry["observation"] = {"error": f"unknown action {action!r}"} + + # Out of rounds / budget: accept a VLM-confirmed page if one exists. + if last_verify and last_verify.get("selected_page") is not None: + page = int(last_verify["selected_page"]) + if page in self.scope_pages: + match = self._build_match( + title=title, + page=page, + seen_grep=seen_grep, + last_verify=last_verify, + decision={"confidence": last_verify.get("confidence")}, + transcript=transcript, + ) + return SubAgentResult(match, transcript, self.config.max_rounds, "round_cap_with_verify") + return SubAgentResult(None, transcript, len(transcript), "exhausted") + + def _dispatch(self, name: str, args: dict[str, Any]): + return REGISTRY.dispatch(name, self.ctx, args) + + def _coerce_pages(self, raw: Any) -> list[int]: + if not isinstance(raw, (list, tuple)): + return [] + pages: list[int] = [] + for value in raw: + try: + page = int(value) + except (TypeError, ValueError): + continue + if page in self.scope_pages and page not in pages: + pages.append(page) + return pages + + def _build_match( + self, + *, + title: str, + page: int, + seen_grep: dict[int, dict[str, Any]], + last_verify: dict[str, Any] | None, + decision: dict[str, Any], + transcript: list[dict[str, Any]], + ) -> TitleMatch: + grep_hit = seen_grep.get(page) or {} + vlm_confirmed = bool(last_verify and last_verify.get("selected_page") == page) + if vlm_confirmed: + source = cast(Any, last_verify.get("source") or "agent_vlm") + confidence = float(last_verify.get("confidence") or 0.7) + reason = last_verify.get("reason") + else: + source = cast(Any, "agent_heuristic") + confidence = float(decision.get("confidence") or grep_hit.get("confidence") or 0.55) + reason = decision.get("reason") + return TitleMatch( + page=page, + confidence=confidence, + source=source, + matched_line=str(grep_hit.get("matched_line") or ""), + score=float(grep_hit.get("score") or confidence), + candidates=sorted(seen_grep.keys()), + evidence={ + "page_locate_agent": { + "decision": source, + "selected_page": page, + "reason": reason, + "vlm_confirmed": vlm_confirmed, + "rounds": len(transcript), + "seen_pages": sorted(seen_grep.keys()), + "transcript": transcript[-8:], + } + }, + ) + + +# ── Decision policy ──────────────────────────────────────────────────────── + + +def decide_next_action( + *, + ctx: ToolContext | None, + title: str, + scope_pages: list[int], + printed_page: int | None, + page_offset_hint: int | None, + observations: list[dict[str, Any]], + seen_pages: list[int], + last_verify: dict[str, Any] | None, + round_index: int, + max_rounds: int, +) -> dict[str, Any] | None: + """Pick the next action. + + Primary path is LLM-driven (the agent may rewrite the query, expand scope, + verify with a VLM, then submit). When no reasoning model is configured we + fall back to a minimal deterministic policy (strict grep → verify → submit) + that performs no query rewriting — only the LLM is allowed to relax the + query, so offline mode never silently mis-anchors. + """ + model = None + if ctx is not None: + model = ctx.settings.get("executor_model") or ctx.settings.get("model") + if ctx is None or not model or getattr(ctx, "budget", None) is None: + return _deterministic_decide( + title=title, + observations=observations, + seen_pages=seen_pages, + last_verify=last_verify, + ) + + from shared.utils.token_estimate import estimate_tokens + + prompt = _build_prompt( + title=title, + scope_pages=scope_pages, + printed_page=printed_page, + page_offset_hint=page_offset_hint, + observations=observations, + seen_pages=seen_pages, + last_verify=last_verify, + round_index=round_index, + max_rounds=max_rounds, + ) + est = estimate_tokens(prompt) + if not ctx.budget.try_reserve("plan", est): + logger.warning("[page_locate.subagent] planner budget exhausted for title={!r}", title) + return None + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=[{"role": "user", "content": prompt}], + model=model, + temperature=0.0, + max_tokens=500, + response_format={"type": "json_object"}, + usage_task="page_memory.page_locate_decide", + ) + ctx.budget.commit("plan", actual=usage.get("total_tokens", est), est=est) + return _normalize_decision(json.loads(raw)) + except Exception as exc: + ctx.budget.refund("plan", est=est) + logger.warning("[page_locate.subagent] decide failed for title={!r}: {}", title, exc) + return None + + +def _deterministic_decide( + *, + title: str, + observations: list[dict[str, Any]], + seen_pages: list[int], + last_verify: dict[str, Any] | None, +) -> dict[str, Any]: + has_grep = any( + entry.get("observation", {}).get("tool") == "grep.title_pages" + for entry in observations + ) + if not has_grep: + return {"action": "grep", "query_title": title, "reason": "initial strict grep"} + if last_verify is None: + if seen_pages: + return {"action": "verify", "pages": seen_pages, "reason": "verify grep candidates"} + return {"action": "give_up", "reason": "strict grep found no candidates"} + selected = last_verify.get("selected_page") + if selected is not None: + return { + "action": "submit", + "selected_page": selected, + "confidence": last_verify.get("confidence", 0.6), + "reason": "submit VLM-confirmed page", + } + return {"action": "give_up", "reason": "verify returned no page"} + + +def _normalize_decision(data: dict[str, Any]) -> dict[str, Any]: + action = str(data.get("action") or "").strip().lower() + if action not in _ALLOWED_ACTIONS: + action = "give_up" + decision: dict[str, Any] = {"action": action, "reason": data.get("reason")} + if action == "grep": + decision["query_title"] = data.get("query_title") + decision["pages"] = data.get("pages") + elif action == "verify": + decision["pages"] = data.get("pages") + elif action == "submit": + selected = data.get("selected_page") + decision["selected_page"] = None if selected in (None, "", "null") else selected + decision["confidence"] = data.get("confidence") + return decision + + +def _build_prompt( + *, + title: str, + scope_pages: list[int], + printed_page: int | None, + page_offset_hint: int | None, + observations: list[dict[str, Any]], + seen_pages: list[int], + last_verify: dict[str, Any] | None, + round_index: int, + max_rounds: int, +) -> str: + scope_desc = ( + f"{scope_pages[0]}-{scope_pages[-1]} ({len(scope_pages)} pages)" + if scope_pages + else "none" + ) + hint = "" + if printed_page is not None: + guessed = printed_page + page_offset_hint if page_offset_hint is not None else None + hint = f"\nPrinted page number in the title's TOC entry: {printed_page}" + ( + f" (rough physical-page guess: {guessed})" if guessed is not None else "" + ) + compact_obs = json.dumps( + [_compact_obs(entry) for entry in observations[-4:]], + ensure_ascii=False, + ) + return ( + "You are a sub-agent that locates the START page of ONE section title " + "inside a PDF body. Decide the single next action and return strict JSON.\n\n" + f"Title to locate: {title!r}\n" + f"Allowed body pages: {scope_desc}{hint}\n" + f"Round {round_index + 1} of {max_rounds}. " + f"Candidate pages seen so far: {seen_pages}. " + f"Last VLM verification: {json.dumps(last_verify, ensure_ascii=False)}\n" + f"Observations (most recent last): {compact_obs}\n\n" + "Available actions (pick exactly ONE, as a JSON object):\n" + '1. {"action":"grep","query_title":"","reason":"..."}\n' + " Searches body pages for the query (exact line, whitespace-insensitive " + "compact text, and token overlap) and returns candidate pages with matched " + "lines. IMPORTANT: the title may carry a trailing document-reference code " + "like （陕十一建[2022]30 号）, a 《》 wrapper, or a （试行）/（2020 版）note " + "that never appears contiguously in the body, and the heading may be split " + "across two lines. If a strict search of the full title returns nothing, " + "retry grep with a shortened, distinctive CORE of the title (drop the " + "trailing brackets/codes/notes).\n" + '2. {"action":"verify","pages":[,...],"reason":"..."}\n' + " Renders those pages as images and asks a vision model which one truly " + "STARTS the section (not a table-of-contents row, a running header/footer, " + "or a body mention/citation). Use this to disambiguate.\n" + '3. {"action":"submit","selected_page":,"confidence":<0..1>,"reason":"..."}\n' + " Finalize. selected_page must be a page you have already seen as a " + "candidate; use null only if the section is genuinely absent.\n" + '4. {"action":"give_up","reason":"..."}\n\n' + "Rules: when there is more than one candidate or any ambiguity, verify with " + "the vision model before submitting. Never invent page numbers. " + "Return ONLY the JSON object." + ) + + +def _compact_obs(entry: dict[str, Any]) -> dict[str, Any]: + obs = entry.get("observation") or {} + tool = obs.get("tool") + if tool == "grep.title_pages": + return { + "action": "grep", + "query": obs.get("query"), + "candidates": [ + {"page": c.get("page"), "source": c.get("source"), "line": (c.get("matched_line") or "")[:60]} + for c in (obs.get("candidates") or []) + ], + } + if tool == "verify.section_page": + choice = obs.get("choice") or {} + return { + "action": "verify", + "pages": obs.get("pages"), + "selected_page": choice.get("selected_page"), + "source": choice.get("source"), + "reason": (choice.get("reason") or "")[:80], + } + return {"action": entry.get("decision", {}).get("action"), "note": obs.get("error")} diff --git a/apps/worker/app/services/document_agent/structure/page_locate_tools.py b/apps/worker/app/services/document_agent/structure/page_locate_tools.py new file mode 100644 index 000000000..3f91868a4 --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/page_locate_tools.py @@ -0,0 +1,156 @@ +"""Registry tools for page-memory residual title location. + +These live under ``structure/`` (not ``tools/``) so the page-memory ReAct +sub-agent can register and dispatch them without importing the full profiling +``tools`` package (which pulls in S3/DB-heavy modules). ``tools/page_locate.py`` +re-exports these so the profile executor and ``tools/__init__`` keep working. +""" + +from __future__ import annotations + +import time +from typing import Any + +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.registry import register_tool +from app.services.document_agent.structure import page_locate_agent as _pla +from app.services.document_agent.structure.hierarchy_locator import TitleMatch + + +@register_tool( + name="grep.title_pages", + description=( + "Find candidate body pages for a section title using strict heading, " + "normalized, compact, and token grep variants." + ), + parameters={ + "type": "object", + "properties": { + "title": {"type": "string"}, + "pages": {"type": "array", "items": {"type": "integer"}}, + "candidate_cap": {"type": "integer"}, + }, + "required": ["title", "pages"], + }, +) +def grep_title_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + title = str(args.get("title") or "").strip() + pages = _valid_pages(ctx, args.get("pages") or []) + if not title or not pages: + return ToolResult( + status="error", + error="grep.title_pages requires title and pages", + latency_ms=int((time.monotonic() - start) * 1000), + ) + page_texts = _get_page_texts(ctx, pages) + matches = _pla.grep_title_page_candidates( + title=title, + scope_pages=pages, + page_texts=page_texts, + limit=int(args.get("candidate_cap") or 8), + ) + return ToolResult( + status="ok", + payload={ + "title": title, + "candidates": [_match_to_payload(match) for match in matches], + }, + latency_ms=int((time.monotonic() - start) * 1000), + ) + + +@register_tool( + name="verify.section_page", + description=( + "Use VLM page screenshots to choose which candidate page starts a section." + ), + parameters={ + "type": "object", + "properties": { + "title": {"type": "string"}, + "pages": {"type": "array", "items": {"type": "integer"}}, + }, + "required": ["title", "pages"], + }, +) +def verify_section_page(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + title = str(args.get("title") or "").strip() + pages = _valid_pages(ctx, args.get("pages") or []) + if not title or not pages: + return ToolResult( + status="error", + error="verify.section_page requires title and pages", + latency_ms=int((time.monotonic() - start) * 1000), + ) + page_texts = _get_page_texts(ctx, pages) + grep_by_page = { + match.page: match + for match in _pla.grep_title_page_candidates( + title=title, + scope_pages=pages, + page_texts=page_texts, + limit=len(pages), + ) + } + # Always verify the *requested* pages, even when grep found nothing on them + # (the heading may be split across lines / wrapped, so render + VLM decides). + candidates = [ + grep_by_page.get(page) + or TitleMatch( + page=page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[page], + evidence={"synthesized": True}, + ) + for page in pages + ] + choice = _pla.verify_section_page_choice( + ctx=ctx, + title=title, + candidate_matches=candidates, + candidate_page_cap=len(pages), + ) + return ToolResult( + status="ok", + payload=choice, + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=int(choice.get("tokens_used") or 0), + ) + + +def _valid_pages(ctx: ToolContext, raw_pages: Any) -> list[int]: + page_count = max(int(ctx.blackboard.page_count or 0), 0) + return sorted( + { + int(page) + for page in raw_pages + if 1 <= int(page) <= page_count + } + ) + + +def _get_page_texts(ctx: ToolContext, pages: list[int]) -> dict[int, str]: + cache = getattr(ctx.blackboard, "page_full_text_cache", {}) + missing = [page for page in pages if page not in cache] + if missing: + from app.services.document_agent.pdf_text import read_page_texts + + cache.update(read_page_texts(ctx.pdf_path, missing, timeout=300)) + return {page: cache.get(page, "") for page in pages} + + +def _match_to_payload(match: Any) -> dict[str, Any]: + return { + "page": match.page, + "confidence": match.confidence, + "source": match.source, + "matched_line": match.matched_line, + "score": match.score, + "candidates": match.candidates, + "evidence": match.evidence, + } diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 00260327d..f322f6d6a 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -7,6 +7,7 @@ from . import grep_text as grep_text # noqa: F401 from . import inspect_pages as inspect_pages # noqa: F401 from . import match_h1_pages as match_h1_pages # noqa: F401 +from . import page_locate as page_locate # noqa: F401 from . import propose_shard_plan as propose_shard_plan # noqa: F401 from . import validate_anatomy_map as validate_anatomy_map # noqa: F401 from . import verdict as verdict # noqa: F401 diff --git a/apps/worker/app/services/document_agent/tools/page_locate.py b/apps/worker/app/services/document_agent/tools/page_locate.py new file mode 100644 index 000000000..0e67f1435 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/page_locate.py @@ -0,0 +1,16 @@ +"""Page-memory residual title-location tools. + +The implementations live in ``structure/page_locate_tools.py`` so the +page-memory sub-agent can register/dispatch them without importing this heavy +``tools`` package. Importing this module (e.g. via ``tools/__init__``) ensures +the tools are registered for the profile executor as well. +""" + +from __future__ import annotations + +from app.services.document_agent.structure.page_locate_tools import ( # noqa: F401 + grep_title_pages, + verify_section_page, +) + +__all__ = ["grep_title_pages", "verify_section_page"] diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py new file mode 100644 index 000000000..0b76d946f --- /dev/null +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -0,0 +1,357 @@ +"""Build page-memory section skeletons from profile-time anatomy.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any + +from app.services.document_agent.manifest import ( + H1Candidate, + PageAnatomyMap, + Shard, + ToolContext, +) +from app.services.document_agent.structure.page_locate_agent import ( + PageLocateResidualAgent, +) +from app.services.document_agent.structure.hierarchy_locator import ( + ResolvedHierarchyRange, + TitleNode, + extract_toc_nodes, + resolve_hierarchy_page_ranges, +) +from app.services.document_parser.structure.body_boundary import ( + clean_toc_title, + normalize_heading_text, +) +from loguru import logger + +_FRONT_TOC_REGION_GAP_PAGES = 5 + + +@dataclass(frozen=True) +class SectionSkeleton: + section_path: str + level: int + start_page: int + end_page: int + title: str + parent_path: str | None = None + evidence: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def extract_section_skeletons( + *, + anatomy: PageAnatomyMap | Any | None, + filename: str, + page_texts: dict[int, str], + ctx: ToolContext | None = None, +) -> list[SectionSkeleton]: + """Convert PageAnatomyMap hierarchy evidence into leaf section skeletons.""" + page_count = _page_count(anatomy) + root_path = f"{filename}/Root" + if page_count <= 0: + return [_root_skeleton(root_path=root_path, filename=filename, page_count=0)] + + toc_hierarchies, toc_selection = _select_global_toc_hierarchies( + anatomy=anatomy, + filename=filename, + ) + toc_nodes = extract_toc_nodes(toc_hierarchies) + nodes = toc_nodes or _h1_nodes(anatomy) + if not nodes: + return [ + _root_skeleton( + root_path=root_path, + filename=filename, + page_count=page_count, + reason="no_hierarchy", + ) + ] + + body_pages = _body_pages(anatomy=anatomy, page_count=page_count) + offset_hint = _estimate_page_offset(nodes=nodes, anatomy=anatomy) + locate_result = PageLocateResidualAgent( + ctx=ctx, + page_texts=page_texts, + body_pages=body_pages, + page_count=page_count, + page_offset_hint=offset_hint, + ).prepare(nodes) + ranges = resolve_hierarchy_page_ranges( + locate_result.nodes, + page_count=page_count, + page_texts=page_texts, + body_pages=body_pages, + page_offset_hint=offset_hint, + match_overrides=locate_result.match_overrides, + ) + if not ranges: + return [ + _root_skeleton( + root_path=root_path, + filename=filename, + page_count=page_count, + reason="unresolved_hierarchy", + ) + ] + + skeletons = [ + _range_to_skeleton( + item, + filename=filename, + page_count=page_count, + shards=_shards(anatomy), + locate_summary=locate_result.summary, + toc_selection=toc_selection, + ) + for item in ranges + ] + _log_unlocated_title_warnings(filename=filename, skeletons=skeletons) + return skeletons + + +def _range_to_skeleton( + item: ResolvedHierarchyRange, + *, + 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, + ) + 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 + evidence = { + **item.evidence, + "resolver": "hierarchy_locator", + "path_titles": path_titles, + "page_locate_summary": locate_summary, + } + 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, + start_page=start_page, + end_page=end_page, + title=item.title, + parent_path=parent_path, + evidence=evidence, + ) + + +def _root_skeleton( + *, + root_path: str, + filename: str, + page_count: int, + reason: str = "no_pages", +) -> SectionSkeleton: + end_page = max(page_count, 1) + return SectionSkeleton( + section_path=root_path, + level=1, + start_page=1, + end_page=end_page, + title="Root", + parent_path=filename, + evidence={"source": "fallback_root", "reason": reason}, + ) + + +def _page_count(anatomy: Any | None) -> int: + return max(int(getattr(anatomy, "page_count", 0) or 0), 0) + + +def _toc_hierarchies(anatomy: Any | None) -> list[dict[str, Any]] | None: + return getattr(anatomy, "toc_hierarchies", None) if anatomy is not None else None + + +def _select_global_toc_hierarchies( + *, + anatomy: Any | None, + filename: str, +) -> tuple[list[dict[str, Any]] | None, dict[str, Any]]: + """Keep the front/global TOC cluster and skip later embedded TOCs. + + Profile-time TOC extraction can find local TOCs inside a long document + (for example, an embedded standard with its own English outline). Page + memory C4 currently emits a document-level skeleton, so later page-based + TOC regions must not be concatenated as root siblings. + """ + hierarchies = list(_toc_hierarchies(anatomy) or []) + if len(hierarchies) <= 1: + return (hierarchies or None), {} + + page_based = [ + hierarchy + for hierarchy in hierarchies + if hierarchy.get("toc_range_unit") == "page" and _toc_range_start(hierarchy) is not None + ] + if not page_based or len(page_based) != len(hierarchies): + return hierarchies, {} + + sorted_items = sorted(enumerate(hierarchies), key=lambda item: _toc_range_start(item[1]) or 0) + selected_indices: set[int] = set() + skipped: list[dict[str, Any]] = [] + cluster_end: int | None = None + + for original_index, hierarchy in sorted_items: + start = _toc_range_start(hierarchy) + end = _toc_range_end(hierarchy) + if start is None or end is None: + selected_indices.add(original_index) + continue + if cluster_end is None: + selected_indices.add(original_index) + cluster_end = end + continue + if start <= cluster_end + _FRONT_TOC_REGION_GAP_PAGES: + selected_indices.add(original_index) + cluster_end = max(cluster_end, end) + continue + skipped.append( + { + "index": original_index, + "toc_range": [start, end], + "scan_range": hierarchy.get("scan_range"), + "reason": "embedded_toc_region_outside_front_cluster", + } + ) + + selected = [ + hierarchy + for index, hierarchy in enumerate(hierarchies) + if index in selected_indices + ] + if skipped: + logger.warning( + "[page_memory.skeleton] skipped embedded toc regions filename={} skipped={}", + filename, + skipped, + ) + summary = { + "strategy": "front_page_toc_cluster", + "input_count": len(hierarchies), + "selected_count": len(selected), + "skipped_count": len(skipped), + "skipped": skipped, + } + return (selected or None), summary + + +def _toc_range_start(hierarchy: dict[str, Any]) -> int | None: + toc_range = hierarchy.get("toc_range") + if not isinstance(toc_range, (list, tuple)) or not toc_range: + return None + try: + return int(toc_range[0]) + except (TypeError, ValueError): + return None + + +def _toc_range_end(hierarchy: dict[str, Any]) -> int | None: + toc_range = hierarchy.get("toc_range") + if not isinstance(toc_range, (list, tuple)) or not toc_range: + return None + try: + return int(toc_range[-1]) + except (TypeError, ValueError): + return None + + +def _h1_nodes(anatomy: Any | None) -> list[TitleNode]: + h1_result = getattr(anatomy, "h1_result", None) + candidates: list[H1Candidate] = list(getattr(h1_result, "h1_candidates", []) or []) + nodes: list[TitleNode] = [] + for candidate in sorted(candidates, key=lambda item: item.page): + title = clean_toc_title(candidate.title) or normalize_heading_text(candidate.title) + if not title: + continue + nodes.append(TitleNode(title=title, level=1, physical_page_hint=candidate.page)) + return nodes + + +def _body_pages(*, anatomy: Any | None, page_count: int) -> list[int]: + excluded: set[int] = set() + toc_result = getattr(anatomy, "toc_result", None) + excluded.update(int(page) for page in getattr(toc_result, "toc_pages", []) or []) + return [page for page in range(1, page_count + 1) if page not in excluded] + + +def _estimate_page_offset(*, nodes: list[TitleNode], anatomy: Any | None) -> int | None: + printed_by_title: dict[str, int] = {} + for node in _walk_nodes(nodes): + if node.printed_page is None: + continue + printed_by_title[_title_key(node.title)] = node.printed_page + + offsets: list[int] = [] + h1_result = getattr(anatomy, "h1_result", None) + for candidate in getattr(h1_result, "h1_candidates", []) or []: + printed_page = printed_by_title.get(_title_key(candidate.title)) + if printed_page is not None: + offsets.append(int(candidate.page) - printed_page) + if not offsets: + return None + offsets.sort() + return offsets[len(offsets) // 2] + + +def _walk_nodes(nodes: list[TitleNode]) -> list[TitleNode]: + walked: list[TitleNode] = [] + for node in nodes: + walked.append(node) + walked.extend(_walk_nodes(node.children)) + return walked + + +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)) + + +def _log_unlocated_title_warnings( + *, + filename: str, + skeletons: list[SectionSkeleton], +) -> None: + for skeleton in skeletons: + for warning in skeleton.evidence.get("warnings", []) or []: + logger.warning( + "[page_memory.skeleton] title unlocated filename={} title={!r} " + "assigned_range={} parent_scope={} path_titles={}", + filename, + warning.get("title"), + warning.get("assigned_range"), + warning.get("parent_scope"), + warning.get("path_titles"), + ) From 5d586d6e49e9ee057e69fe5e205c2025c3602a02 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 17 Jun 2026 09:24:23 +0800 Subject: [PATCH 2/5] =?UTF-8?q?feat:=20PR4=20page-memory=20page=20mode=20?= =?UTF-8?q?=E2=80=94=20C1-C7=20modules=20+=20GAP-1/2/3=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New page_memory modules: - C1 page_renderer: PNG + thumbnail + raw_text per page - C2 page_plan: rule-based vlm_lite/text_only/skip_tagging strategy - C3 page_tagger: VLM per-page annotation with JSON retry + blurry degradation - C6 page_section_mapper: skeleton × tagger → section_path (primary/spans/inherited) - C7 memory_service: unified page/shard_page builder via full C1-C7 pipeline shared-python GAP fixes: - GAP-1: zip_chunk_schema recognizes 'page' chunk type (no collapse to text) - GAP-2: zip_result_resources collects pages/ directory - GAP-3: zip_doc_navigation counts page_chunks in stats --- .gitignore | 1 + .../services/page_memory/memory_service.py | 215 +++++++++++-- .../app/services/page_memory/page_plan.py | 83 ++++++ .../app/services/page_memory/page_renderer.py | 210 +++++++++++++ .../page_memory/page_section_mapper.py | 134 +++++++++ .../app/services/page_memory/page_tagger.py | 282 ++++++++++++++++++ .../services/storage/zip_chunk_schema.py | 10 + .../services/storage/zip_doc_navigation.py | 10 + .../services/storage/zip_result_resources.py | 23 ++ 9 files changed, 938 insertions(+), 30 deletions(-) create mode 100644 apps/worker/app/services/page_memory/page_plan.py create mode 100644 apps/worker/app/services/page_memory/page_renderer.py create mode 100644 apps/worker/app/services/page_memory/page_section_mapper.py create mode 100644 apps/worker/app/services/page_memory/page_tagger.py diff --git a/.gitignore b/.gitignore index 8a56b4f57..66e4ded6b 100644 --- a/.gitignore +++ b/.gitignore @@ -123,3 +123,4 @@ TODO.md doc-deploy/ todos/ *.egg-info/ +*.rdb diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 816ea70b2..1808b619e 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -16,13 +16,10 @@ from app.services.document_parser.support.parser_rows import PARSER_ROW_COLUMNS from app.services.page_memory.normalizer import normalize_to_pdf +from loguru import logger from shared.core.exceptions.domain_exceptions import ValidationException -_SUPPORTED_GRANULARITY = "whole_doc" -_UNSUPPORTED_GRANULARITY_REASON = "PAGE_MEMORY_GRANULARITY_NOT_IMPLEMENTED" - - @dataclass(frozen=True) class PageMemoryInput: file_path: str @@ -36,8 +33,9 @@ class PageMemoryInput: def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: """Run the page-memory track. - PR3 intentionally implements only the whole_doc skeleton. Full page mode, - tagging, and section mapping land in PR4. + Supports three granularity verdicts: + - ``whole_doc`` (≤6 pages, no TOC) → single whole-document chunk + - ``page`` / ``shard_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) @@ -54,13 +52,22 @@ def run(request: PageMemoryInput) -> tuple[str, pd.DataFrame]: output_dir=full_output_dir, ) verdict = _decide_granularity(profile) - if verdict != _SUPPORTED_GRANULARITY: - _raise_unsupported_granularity(verdict) - return full_output_dir, _build_whole_doc_dataframe( + + if verdict == "whole_doc": + return full_output_dir, _build_whole_doc_dataframe( + pdf_path=pdf_path, + filename=request.filename, + output_dir=full_output_dir, + page_count=max(int(profile.page_count or 0), 0), + verdict=verdict, + ) + + # page / shard_page → unified per-page pipeline + return full_output_dir, _build_page_dataframe( pdf_path=pdf_path, filename=request.filename, output_dir=full_output_dir, - page_count=max(int(profile.page_count or 0), 0), + profile=profile, verdict=verdict, ) @@ -81,28 +88,175 @@ def _decide_granularity(profile: Any) -> str: return "page" -def _raise_unsupported_granularity(verdict: str) -> None: - raise ValidationException( - user_message=( - "page_memory is enabled, but this PR only supports whole-document " - "page memory. Per-page and shard-page modes are intentionally gated " - "until the page renderer, tagger, and section mapper land." - ), - violations=[ - { - "field": "parse_track", - "description": ( - f"{_UNSUPPORTED_GRANULARITY_REASON}: " - f"granularity={verdict}; supported={_SUPPORTED_GRANULARITY}" - ), - } - ], - internal_message=( - f"{_UNSUPPORTED_GRANULARITY_REASON}: granularity={verdict}; " - f"supported={_SUPPORTED_GRANULARITY}" - ), +# ── page builder (C1→C2→C3→C4→C6→C7) ──────────────────────────────── + + +def _build_page_dataframe( + *, + pdf_path: str, + filename: str, + output_dir: str, + profile: Any, + verdict: str, +) -> pd.DataFrame: + """Build per-page DataFrame via the full C1-C7 pipeline. + + Steps: + C4 skeleton_extractor → SectionSkeleton[] + C1 page_renderer → PageRenderResult[] + C2 page_plan → PagePlan[] + C3 page_tagger → PageTagResult[] + C6 page_section_mapper → PageSectionMapping[] + C7 assemble 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 + + 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"])) + + # ── budget (page_tagging envelope for VLM calls) ────────────────── + page_tagging_budget = int( + os.environ.get("PAGE_MEMORY_TAG_BUDGET", str(page_count * 1200)) + ) + budget = BudgetTracker( + plan_budget=0, + visual_budget=page_tagging_budget, + visual_stage_envelopes={ + "page_tagging": StageEnvelope( + min_guarantee=page_tagging_budget, + cap=None, + ), + }, ) + # ── C4: skeleton (from profile anatomy) ─────────────────────────── + if anatomy is not None: + page_texts = read_page_texts( + pdf_path, list(range(1, page_count + 1)), timeout=300, + ) + skeletons = extract_section_skeletons( + anatomy=anatomy, + filename=filename, + page_texts=page_texts, + ctx=None, # no VLM for skeleton in page mode + ) + else: + skeletons = [] + page_texts = read_page_texts( + pdf_path, list(range(1, page_count + 1)), timeout=300, + ) + logger.info( + "[page_memory] C4 skeleton: {} sections from anatomy", + 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, + ) + + # ── 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, + ) + + # ── C6: page → section mapping ─────────────────────────────────── + mappings = map_pages_to_sections( + page_count=page_count, + skeletons=skeletons, + tag_results=tags, + filename=filename, + ) + + # ── C7: assemble DataFrame rows ────────────────────────────────── + tag_map = {t.page_index: t for t in tags} + render_map = {r.page_index: r for r in rendered} + plan_map = {p.page_index: p for p in plans} + + rows: list[dict[str, Any]] = [] + for mapping in mappings: + page = mapping.page_index + tag = tag_map.get(page) + rend = render_map.get(page) + plan = plan_map.get(page) + + raw_text = rend.raw_text if rend else page_texts.get(page, "") + summary = tag.summary if tag else "" + content = f"[SUMMARY]\n{summary}\n\n[RAW]\n{raw_text}".strip() + + doc_hash = gen_str_codes(f"{filename}::{page}") + know_id = f"page_{doc_hash}" + + image_uri = "" + thumb_uri = "" + if rend: + if rend.image_path and os.path.exists(rend.image_path): + image_uri = str( + Path(rend.image_path).relative_to(output_dir) + ) + if rend.thumb_path and os.path.exists(rend.thumb_path): + thumb_uri = str( + Path(rend.thumb_path).relative_to(output_dir) + ) + + rows.append({ + "content": content, + "path": mapping.section_path, + "type": "page", + "length": len(content), + "keywords": "", + "summary": summary, + "know_id": know_id, + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": str(page), + "extra_metadata": { + "granularity": "page", + "strategy_used": tag.strategy_used if tag else "", + "source_verdict": verdict, + "page_index": page, + "page_image_uri": image_uri, + "thumb_uri": thumb_uri, + "status": tag.status if tag else "clear", + "kind": plan.reason if plan else "", + "observed_titles": tag.observed_titles if tag else [], + "section_roles": mapping.section_roles, + }, + }) + + logger.info( + "[page_memory] C7 assembled {} page rows (verdict={})", + len(rows), verdict, + ) + return pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) + + +# ── whole_doc builder (PR3, unchanged) ──────────────────────────────── + def _build_whole_doc_dataframe( *, @@ -180,3 +334,4 @@ def _render_page_images( for item in rendered if item.get("png_path") ] + diff --git a/apps/worker/app/services/page_memory/page_plan.py b/apps/worker/app/services/page_memory/page_plan.py new file mode 100644 index 000000000..30557941a --- /dev/null +++ b/apps/worker/app/services/page_memory/page_plan.py @@ -0,0 +1,83 @@ +"""Page processing plan: maps page labels to tagging strategy. + +Reads ``PageLabel.kind`` + ``PageFeature`` and assigns each page +one of three strategies: + +- ``vlm_lite`` — send the page image to VLM for summary/status/titles +- ``text_only`` — use raw text heuristics (no VLM call, saves cost) +- ``skip_tagging`` — blank-like page, preserve image only +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + +from app.services.document_agent.manifest import PageFeature, PageLabel + + +class PageProcessingStrategy(str, Enum): + """Tagging strategy for a single page.""" + + VLM_LITE = "vlm_lite" + TEXT_ONLY = "text_only" + SKIP_TAGGING = "skip_tagging" + + +@dataclass(frozen=True) +class PagePlan: + """Processing plan for a single page.""" + + page_index: int + strategy: PageProcessingStrategy + reason: str + + +def derive_page_processing_plan( + *, + page_count: int, + page_labels: list[PageLabel], + page_features: list[PageFeature], +) -> list[PagePlan]: + """Assign a processing strategy to every page. + + Rules (PR4 version): + - ``low_content`` + ``is_blank_like`` → ``skip_tagging`` + - ``table_heavy`` with sufficient raw text → ``text_only`` + - everything else (normal / image_heavy / landscape) → ``vlm_lite`` + + Returns one ``PagePlan`` per page (1-indexed), ordered by page_index. + """ + label_map: dict[int, PageLabel] = {label.page: label for label in page_labels} + feature_map: dict[int, PageFeature] = {feat.page: feat for feat in page_features} + + plans: list[PagePlan] = [] + for page in range(1, page_count + 1): + label = label_map.get(page) + feature = feature_map.get(page) + strategy, reason = _classify_page(label, feature) + plans.append(PagePlan(page_index=page, strategy=strategy, reason=reason)) + + return plans + + +# ── minimum raw text length for table_heavy → text_only ────────────── +_TABLE_TEXT_THRESHOLD = 200 + + +def _classify_page( + label: PageLabel | None, + feature: PageFeature | None, +) -> tuple[PageProcessingStrategy, str]: + """Determine the strategy for a single page.""" + kind = label.kind if label else "normal" + is_blank = feature.is_blank_like if feature else False + text_len = feature.raw_text_length if feature else 0 + + if kind == "low_content" and is_blank: + return PageProcessingStrategy.SKIP_TAGGING, "low_content + blank_like" + + if kind == "table_heavy" and text_len >= _TABLE_TEXT_THRESHOLD: + return PageProcessingStrategy.TEXT_ONLY, f"table_heavy with {text_len} chars" + + return PageProcessingStrategy.VLM_LITE, f"kind={kind}" diff --git a/apps/worker/app/services/page_memory/page_renderer.py b/apps/worker/app/services/page_memory/page_renderer.py new file mode 100644 index 000000000..fca800c85 --- /dev/null +++ b/apps/worker/app/services/page_memory/page_renderer.py @@ -0,0 +1,210 @@ +"""Page renderer: produce PNG, thumbnail, and raw text for each page. + +Wraps the existing ``document_agent/visual.render_pages`` and +``pdf_text.read_page_texts`` utilities, adding thumbnail generation +(72 dpi), dimensions, and landscape detection. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from loguru import logger + +from app.services.document_agent.manifest import PageFeature, ToolContext +from app.services.document_agent.pdf_text import read_page_texts +from app.services.document_agent.visual import render_pages + + +@dataclass(frozen=True) +class PageRenderResult: + """Output of rendering a single page.""" + + page_index: int + """1-based page number.""" + + image_path: str + """Absolute path to full-resolution PNG (``pages/page-N.png``).""" + + thumb_path: str + """Absolute path to thumbnail JPEG (``pages/thumb-N.jpg``).""" + + raw_text: str + """PyMuPDF extracted text for this page.""" + + width: float + """Page width in points.""" + + height: float + """Page height in points.""" + + is_landscape: bool + """``True`` when *width > height*.""" + + +def render_document_pages( + *, + pdf_path: str, + page_count: int, + output_dir: str, + page_features: list[PageFeature] | None = None, + ctx: ToolContext | None = None, + dpi: int = 144, + thumb_dpi: int = 72, + timeout: int = 300, +) -> list[PageRenderResult]: + """Render every page in *pdf_path* as PNG + thumb + raw text. + + Parameters + ---------- + pdf_path: + Local filesystem path to the PDF. + page_count: + Total page count (avoids re-opening the PDF). + output_dir: + Root output directory; pages are written to ``output_dir/pages/``. + page_features: + If available, dimensions are read from here (avoiding a second + PyMuPDF open). Otherwise falls back to 0/0/False. + ctx: + An optional ``ToolContext`` forwarded to ``render_pages``. If + *None*, a lightweight context is constructed internally. + dpi: + Full-resolution DPI (default 144). + thumb_dpi: + Thumbnail DPI (default 72). + timeout: + PyMuPDF subprocess timeout in seconds. + + Returns + ------- + list[PageRenderResult] + One entry per page, ordered by page_index. + """ + pages = list(range(1, page_count + 1)) + if not pages: + return [] + + # ── raw text ────────────────────────────────────────────────────── + page_texts = read_page_texts(pdf_path, pages, timeout=timeout) + + # ── full-resolution PNGs ────────────────────────────────────────── + if ctx is not None: + pngs = render_pages( + ctx, + pages, + folder_name="pages", + prefix="page", + dpi=dpi, + timeout=timeout, + ) + else: + from app.services.document_agent.state import AgentBlackboard + + blackboard = AgentBlackboard() + blackboard.page_count = page_count + tmp_ctx = ToolContext( + pdf_path=pdf_path, + job_id="page_renderer", + blackboard=blackboard, + budget=None, + trace=None, + output_dir=output_dir, + settings={"agent_png_dpi": str(dpi)}, + ) + pngs = render_pages( + tmp_ctx, + pages, + folder_name="pages", + prefix="page", + dpi=dpi, + timeout=timeout, + ) + + png_map: dict[int, str] = { + item["page"]: item["png_path"] for item in pngs if item.get("png_path") + } + + # ── thumbnails (downsample PNG → JPEG at thumb_dpi) ─────────────── + thumb_map = _generate_thumbnails( + png_map, output_dir=output_dir, dpi=dpi, thumb_dpi=thumb_dpi, + ) + + # ── page dimensions from page_features ──────────────────────────── + feature_map: dict[int, PageFeature] = {} + if page_features: + for feat in page_features: + feature_map[feat.page] = feat + + # ── assemble results ────────────────────────────────────────────── + results: list[PageRenderResult] = [] + for page in pages: + feat = feature_map.get(page) + results.append( + PageRenderResult( + page_index=page, + image_path=png_map.get(page, ""), + thumb_path=thumb_map.get(page, ""), + raw_text=page_texts.get(page, ""), + width=feat.width if feat else 0.0, + height=feat.height if feat else 0.0, + is_landscape=( + feat.orientation == "landscape" if feat else False + ), + ) + ) + + logger.info( + "[page_renderer] rendered {} pages ({} PNGs, {} thumbs)", + len(results), + len(png_map), + len(thumb_map), + ) + return results + + +def _generate_thumbnails( + png_map: dict[int, str], + *, + output_dir: str, + dpi: int, + thumb_dpi: int, +) -> dict[int, str]: + """Downsample full-resolution PNGs to JPEG thumbnails.""" + thumb_dir = os.path.join(output_dir, "pages") + os.makedirs(thumb_dir, exist_ok=True) + + try: + from PIL import Image # type: ignore[import-untyped] + except ImportError: + logger.warning( + "[page_renderer] Pillow not available; skipping thumbnail generation" + ) + return {} + + scale = thumb_dpi / max(dpi, 1) + thumb_map: dict[int, str] = {} + + for page, png_path in png_map.items(): + if not os.path.exists(png_path): + continue + try: + with Image.open(png_path) as img: + new_size = ( + max(int(img.width * scale), 1), + max(int(img.height * scale), 1), + ) + thumb = img.resize(new_size, Image.LANCZOS) + thumb_name = f"thumb_{page}.jpg" + thumb_path = os.path.join(thumb_dir, thumb_name) + thumb.convert("RGB").save(thumb_path, "JPEG", quality=75) + thumb_map[page] = thumb_path + except Exception as exc: + logger.warning( + "[page_renderer] thumbnail failed for page {}: {}", page, exc, + ) + + return thumb_map diff --git a/apps/worker/app/services/page_memory/page_section_mapper.py b/apps/worker/app/services/page_memory/page_section_mapper.py new file mode 100644 index 000000000..06fc76b2f --- /dev/null +++ b/apps/worker/app/services/page_memory/page_section_mapper.py @@ -0,0 +1,134 @@ +"""Page-to-section mapper: assigns each page a section_path and role. + +Given ``SectionSkeleton`` (from C4) and ``PageTagResult`` (from C3), +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 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 + + +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], + tag_results: list[PageTagResult] | None = None, + 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``. + tag_results: + Optional tagger output; ``observed_titles`` can refine primary + detection (future use — currently unused). + 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/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py new file mode 100644 index 000000000..cf9c1dc2d --- /dev/null +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -0,0 +1,282 @@ +"""Page tagger: VLM per-page annotation for summary/status/observed_titles. + +For ``vlm_lite`` pages, sends the page PNG to the VLM and expects a JSON +response. ``text_only`` pages get a rules-based summary from raw text. +``skip_tagging`` pages are left empty. + +Budget is drawn from the ``page_tagging`` stage envelope. +""" + +from __future__ import annotations + +import base64 +import json +import os +from dataclasses import dataclass, field +from typing import Any, Literal, cast + +from loguru import logger + +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.utils.token_estimate import estimate_tokens + + +PageStatus = Literal["clear", "blurry", "skipped"] + + +@dataclass +class PageTagResult: + """Tagging output for a single page.""" + + page_index: int + summary: str = "" + status: PageStatus = "clear" + observed_titles: list[str] = field(default_factory=list) + strategy_used: str = "" + + +# ── prompt placeholder — NEEDS USER CONFIRMATION ───────────────────── + +_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>", + "status": "clear" | "blurry", + "observed_titles": ["", ...] +} + +Rules: +- "summary" should describe the main content visible on the page. +- "status" should be "clear" if the page is legible, "blurry" if the text is + unreadable or the image is too low quality to summarize. +- "observed_titles" should list any headings, section titles, or chapter titles + visible on the page. Return an empty array if none are visible. +- Return ONLY the JSON object, no markdown fences or extra text. +""" + +_BUDGET_STAGE = "page_tagging" +_MAX_JSON_RETRIES = 1 +_RAW_TEXT_SUMMARY_LIMIT = 500 + + +def tag_pages( + *, + pages: list[PageRenderResult], + plans: list[PagePlan], + budget: BudgetTracker | None = None, + vlm_model: str | None = None, +) -> list[PageTagResult]: + """Tag all pages according to their processing plan. + + Parameters + ---------- + pages: + Rendered page results (from ``page_renderer``). + plans: + Processing plans (from ``page_plan``). + budget: + Optional budget tracker with a ``page_tagging`` stage envelope. + vlm_model: + VLM model name; falls back to ``$IMAGE_MODEL``. + + Returns + ------- + list[PageTagResult] + One result per page, ordered by page_index. + """ + plan_map = {plan.page_index: plan for plan in plans} + model = vlm_model or os.environ.get("IMAGE_MODEL") + + results: list[PageTagResult] = [] + vlm_calls = 0 + + for page in pages: + plan = plan_map.get(page.page_index) + strategy = plan.strategy if plan else PageProcessingStrategy.VLM_LITE + + if strategy == PageProcessingStrategy.SKIP_TAGGING: + results.append( + PageTagResult( + page_index=page.page_index, + status="skipped", + strategy_used="skip_tagging", + ) + ) + continue + + if strategy == PageProcessingStrategy.TEXT_ONLY: + results.append(_tag_text_only(page)) + continue + + # vlm_lite + if not model: + logger.warning( + "[page_tagger] no VLM model for page {}; falling back to text_only", + page.page_index, + ) + results.append(_tag_text_only(page)) + continue + + tag = _tag_vlm_lite(page, model=model, budget=budget) + results.append(tag) + vlm_calls += 1 + + logger.info( + "[page_tagger] tagged {} pages ({} VLM calls, {} text_only, {} skipped)", + len(results), + vlm_calls, + sum(1 for r in results if r.strategy_used == "text_only"), + sum(1 for r in results if r.status == "skipped"), + ) + return results + + +def _tag_text_only(page: PageRenderResult) -> PageTagResult: + """Rules-based tag from raw text (no VLM).""" + raw = page.raw_text.strip() + summary = " ".join(raw.split())[:_RAW_TEXT_SUMMARY_LIMIT] + if not summary: + summary = f"Page {page.page_index} (no extractable text)" + + # Simple heuristic: lines that look like headings (short, no trailing punct) + observed: list[str] = [] + for line in raw.splitlines()[:30]: + stripped = line.strip() + if ( + stripped + and 3 < len(stripped) < 100 + and not stripped.endswith((".", "。", ",", "，", ";", "；")) + and not stripped[0].isdigit() + and stripped[0].isupper() or not stripped[0].isascii() + ): + # Very rough heuristic; the real heading detection is in skeleton + pass + # For text_only mode, we don't attempt heading detection + return PageTagResult( + page_index=page.page_index, + summary=summary, + status="clear", + observed_titles=observed, + strategy_used="text_only", + ) + + +def _tag_vlm_lite( + page: PageRenderResult, + *, + model: str, + budget: BudgetTracker | None, +) -> PageTagResult: + """Send page PNG to VLM and parse JSON response.""" + est = estimate_tokens(_VLM_TAG_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, + status=result.status, + observed_titles=result.observed_titles, + strategy_used="text_only_budget_fallback", + ) + return result + + 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, + ) + 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": _VLM_TAG_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=0.0, + max_tokens=600, + 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) + return PageTagResult( + page_index=page.page_index, + summary=str(data.get("summary", "")), + status="clear" if data.get("status") != "blurry" else "blurry", + observed_titles=list(data.get("observed_titles") or []), + 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: blurry degradation + logger.warning( + "[page_tagger] JSON retry exhausted for page {}; blurry fallback", + page.page_index, + ) + raw_summary = " ".join(page.raw_text.split())[:_RAW_TEXT_SUMMARY_LIMIT] + return PageTagResult( + page_index=page.page_index, + summary=raw_summary or f"Page {page.page_index} (VLM parse failed)", + status="blurry", + observed_titles=[], + strategy_used="vlm_lite_blurry_fallback", + ) + 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) 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 2dbf973d1..5b4ea35c5 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -35,6 +35,9 @@ def calculate_statistics(self, chunks: list[dict[str, Any]]) -> dict[str, Any]: "text_chunks": text_chunks, "image_chunks": image_chunks, "table_chunks": table_chunks, + "page_chunks": sum( + 1 for c in chunks if _normalize_chunk_type(c.get("type", "")) == "page" + ), "total_pages": None, } @@ -61,6 +64,8 @@ def format_chunks( chunk_type = "image" elif normalized_type == "table": chunk_type = "table" + elif normalized_type == "page": + chunk_type = "page" else: chunk_type = "text" @@ -98,6 +103,11 @@ def format_chunks( "keywords", [] ) metadata["tokens"] = [] + elif chunk_type == "page": + # Page chunks carry page_image_uri and section_roles in + # extra_metadata; minimal metadata here. + metadata["keywords"] = existing_metadata.get("keywords") or [] + metadata["tokens"] = [] formatted.append( { 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 e4e2f8f13..2cd721a33 100644 --- a/packages/shared-python/shared/services/storage/zip_doc_navigation.py +++ b/packages/shared-python/shared/services/storage/zip_doc_navigation.py @@ -47,6 +47,7 @@ def build_doc_nav( "text_chunks": 0, "image_chunks": 0, "table_chunks": 0, + "page_chunks": 0, "max_depth": 0, } @@ -76,6 +77,15 @@ def build_doc_nav( "summary": summary or content_preview, } ) + elif chunk_type == "page": + stats["page_chunks"] += 1 + # Page chunks participate in section tree like text chunks + text_chunks.append( + { + "path": path, + "summary": summary or content_preview, + } + ) else: stats["text_chunks"] += 1 text_chunks.append( 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 e84e0cb5a..dd96c4e67 100644 --- a/packages/shared-python/shared/services/storage/zip_result_resources.py +++ b/packages/shared-python/shared/services/storage/zip_result_resources.py @@ -18,6 +18,7 @@ class ZipPackageResources: image_files: tuple[ZipResourceFileInfo, ...] table_files: tuple[ZipResourceFileInfo, ...] + page_files: tuple[ZipResourceFileInfo, ...] = () @property def image_files_map(self) -> dict[str, ZipResourceFileInfo]: @@ -39,9 +40,11 @@ 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( @@ -213,6 +216,26 @@ 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] = {} From 5caf41a606c4fe572d26fc333ac6fee30114af01 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 17 Jun 2026 09:49:51 +0800 Subject: [PATCH 3/5] refactor: align page chunk fields with V2 spec Field changes: - content = raw PyMuPDF text only (no [SUMMARY]/[RAW] markers) - summary = VLM or LLM-generated summary (metadata only) - keywords = VLM or summary-full LLM extracted (semicolon-separated) - kind = PageLabel.kind from Profile Agent (not plan.reason) - observed_titles = from C4 skeleton primary titles (not VLM) - Remove thumb_uri (only page_image_uri kept) - Remove status field (strategy_used covers quality info) Strategy changes: - text_only: calls existing summary-full LLM for summary+keywords - skip_tagging: preserves raw text content, marks EMPTY if blank - vlm_lite: outputs summary+keywords (no observed_titles) - Mapper no longer depends on PageTagResult --- .../services/page_memory/memory_service.py | 45 ++--- .../app/services/page_memory/page_renderer.py | 68 +------- .../page_memory/page_section_mapper.py | 9 +- .../app/services/page_memory/page_tagger.py | 155 +++++++++++------- 4 files changed, 129 insertions(+), 148 deletions(-) diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 1808b619e..12852e882 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -187,7 +187,6 @@ def _build_page_dataframe( mappings = map_pages_to_sections( page_count=page_count, skeletons=skeletons, - tag_results=tags, filename=filename, ) @@ -196,38 +195,46 @@ def _build_page_dataframe( render_map = {r.page_index: r for r in rendered} plan_map = {p.page_index: p for p in plans} + # Build page → PageLabel.kind lookup + label_map: dict[int, str] = {} + if page_labels: + 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) + rows: list[dict[str, Any]] = [] for mapping in mappings: page = mapping.page_index tag = tag_map.get(page) rend = render_map.get(page) - plan = plan_map.get(page) raw_text = rend.raw_text if rend else page_texts.get(page, "") + content = raw_text.strip() summary = tag.summary if tag else "" - content = f"[SUMMARY]\n{summary}\n\n[RAW]\n{raw_text}".strip() + keywords_list = tag.keywords if tag else [] + keywords_str = ";".join(keywords_list) doc_hash = gen_str_codes(f"{filename}::{page}") know_id = f"page_{doc_hash}" image_uri = "" - thumb_uri = "" - if rend: - if rend.image_path and os.path.exists(rend.image_path): - image_uri = str( - Path(rend.image_path).relative_to(output_dir) - ) - if rend.thumb_path and os.path.exists(rend.thumb_path): - thumb_uri = str( - Path(rend.thumb_path).relative_to(output_dir) - ) + if rend and rend.image_path and os.path.exists(rend.image_path): + image_uri = str( + Path(rend.image_path).relative_to(output_dir) + ) rows.append({ "content": content, "path": mapping.section_path, "type": "page", "length": len(content), - "keywords": "", + "keywords": keywords_str, "summary": summary, "know_id": know_id, "tokens": "", @@ -236,15 +243,13 @@ def _build_page_dataframe( "page_nums": str(page), "extra_metadata": { "granularity": "page", - "strategy_used": tag.strategy_used if tag else "", - "source_verdict": verdict, "page_index": page, "page_image_uri": image_uri, - "thumb_uri": thumb_uri, - "status": tag.status if tag else "clear", - "kind": plan.reason if plan else "", - "observed_titles": tag.observed_titles if tag else [], + "strategy_used": tag.strategy_used if tag else "", + "kind": label_map.get(page, "normal"), + "observed_titles": skeleton_titles.get(page, []), "section_roles": mapping.section_roles, + "source_verdict": verdict, }, }) diff --git a/apps/worker/app/services/page_memory/page_renderer.py b/apps/worker/app/services/page_memory/page_renderer.py index fca800c85..ca6a02371 100644 --- a/apps/worker/app/services/page_memory/page_renderer.py +++ b/apps/worker/app/services/page_memory/page_renderer.py @@ -1,15 +1,14 @@ -"""Page renderer: produce PNG, thumbnail, and raw text for each page. +"""Page renderer: produce PNG and raw text for each page. Wraps the existing ``document_agent/visual.render_pages`` and -``pdf_text.read_page_texts`` utilities, adding thumbnail generation -(72 dpi), dimensions, and landscape detection. +``pdf_text.read_page_texts`` utilities, adding dimensions and +landscape detection from page_features. """ from __future__ import annotations import os from dataclasses import dataclass -from pathlib import Path from typing import Any from loguru import logger @@ -29,9 +28,6 @@ class PageRenderResult: image_path: str """Absolute path to full-resolution PNG (``pages/page-N.png``).""" - thumb_path: str - """Absolute path to thumbnail JPEG (``pages/thumb-N.jpg``).""" - raw_text: str """PyMuPDF extracted text for this page.""" @@ -53,10 +49,9 @@ def render_document_pages( page_features: list[PageFeature] | None = None, ctx: ToolContext | None = None, dpi: int = 144, - thumb_dpi: int = 72, timeout: int = 300, ) -> list[PageRenderResult]: - """Render every page in *pdf_path* as PNG + thumb + raw text. + """Render every page in *pdf_path* as PNG + raw text. Parameters ---------- @@ -74,8 +69,6 @@ def render_document_pages( *None*, a lightweight context is constructed internally. dpi: Full-resolution DPI (default 144). - thumb_dpi: - Thumbnail DPI (default 72). timeout: PyMuPDF subprocess timeout in seconds. @@ -128,11 +121,6 @@ def render_document_pages( item["page"]: item["png_path"] for item in pngs if item.get("png_path") } - # ── thumbnails (downsample PNG → JPEG at thumb_dpi) ─────────────── - thumb_map = _generate_thumbnails( - png_map, output_dir=output_dir, dpi=dpi, thumb_dpi=thumb_dpi, - ) - # ── page dimensions from page_features ──────────────────────────── feature_map: dict[int, PageFeature] = {} if page_features: @@ -147,7 +135,6 @@ def render_document_pages( PageRenderResult( page_index=page, image_path=png_map.get(page, ""), - thumb_path=thumb_map.get(page, ""), raw_text=page_texts.get(page, ""), width=feat.width if feat else 0.0, height=feat.height if feat else 0.0, @@ -158,53 +145,8 @@ def render_document_pages( ) logger.info( - "[page_renderer] rendered {} pages ({} PNGs, {} thumbs)", + "[page_renderer] rendered {} pages ({} PNGs)", len(results), len(png_map), - len(thumb_map), ) return results - - -def _generate_thumbnails( - png_map: dict[int, str], - *, - output_dir: str, - dpi: int, - thumb_dpi: int, -) -> dict[int, str]: - """Downsample full-resolution PNGs to JPEG thumbnails.""" - thumb_dir = os.path.join(output_dir, "pages") - os.makedirs(thumb_dir, exist_ok=True) - - try: - from PIL import Image # type: ignore[import-untyped] - except ImportError: - logger.warning( - "[page_renderer] Pillow not available; skipping thumbnail generation" - ) - return {} - - scale = thumb_dpi / max(dpi, 1) - thumb_map: dict[int, str] = {} - - for page, png_path in png_map.items(): - if not os.path.exists(png_path): - continue - try: - with Image.open(png_path) as img: - new_size = ( - max(int(img.width * scale), 1), - max(int(img.height * scale), 1), - ) - thumb = img.resize(new_size, Image.LANCZOS) - thumb_name = f"thumb_{page}.jpg" - thumb_path = os.path.join(thumb_dir, thumb_name) - thumb.convert("RGB").save(thumb_path, "JPEG", quality=75) - thumb_map[page] = thumb_path - except Exception as exc: - logger.warning( - "[page_renderer] thumbnail failed for page {}: {}", page, exc, - ) - - return thumb_map diff --git a/apps/worker/app/services/page_memory/page_section_mapper.py b/apps/worker/app/services/page_memory/page_section_mapper.py index 06fc76b2f..23f275233 100644 --- a/apps/worker/app/services/page_memory/page_section_mapper.py +++ b/apps/worker/app/services/page_memory/page_section_mapper.py @@ -1,7 +1,7 @@ """Page-to-section mapper: assigns each page a section_path and role. -Given ``SectionSkeleton`` (from C4) and ``PageTagResult`` (from C3), -maps every page to one or more section paths with roles: +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 @@ -19,7 +19,6 @@ from loguru import logger -from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton @@ -45,7 +44,6 @@ def map_pages_to_sections( *, page_count: int, skeletons: list[SectionSkeleton], - tag_results: list[PageTagResult] | None = None, filename: str = "", ) -> list[PageSectionMapping]: """Map every page (1..page_count) to a section_path. @@ -57,9 +55,6 @@ def map_pages_to_sections( skeletons: Leaf skeletons from ``skeleton_extractor``. Each has ``section_path``, ``start_page``, ``end_page``. - tag_results: - Optional tagger output; ``observed_titles`` can refine primary - detection (future use — currently unused). filename: Source filename for the fallback root path. diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index cf9c1dc2d..d4a48eb58 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -1,8 +1,10 @@ -"""Page tagger: VLM per-page annotation for summary/status/observed_titles. +"""Page tagger: VLM per-page annotation for summary and keywords. For ``vlm_lite`` pages, sends the page PNG to the VLM and expects a JSON -response. ``text_only`` pages get a rules-based summary from raw text. -``skip_tagging`` pages are left empty. +response with ``summary`` and ``keywords``. +For ``text_only`` pages, calls the existing ``summary-full`` LLM prompt +to extract summary + keywords from raw text. +For ``skip_tagging`` pages, content is preserved but summary is omitted. Budget is drawn from the ``page_tagging`` stage envelope. """ @@ -23,21 +25,17 @@ from shared.utils.token_estimate import estimate_tokens -PageStatus = Literal["clear", "blurry", "skipped"] - - @dataclass class PageTagResult: """Tagging output for a single page.""" page_index: int summary: str = "" - status: PageStatus = "clear" - observed_titles: list[str] = field(default_factory=list) + keywords: list[str] = field(default_factory=list) strategy_used: str = "" -# ── prompt placeholder — NEEDS USER CONFIRMATION ───────────────────── +# ── VLM prompt: outputs summary + keywords only ───────────────────── _VLM_TAG_PROMPT = """\ You are annotating a single PDF page screenshot for a document memory system. @@ -45,16 +43,16 @@ class PageTagResult: { "summary": "<1-3 sentence summary of the page content>", - "status": "clear" | "blurry", - "observed_titles": ["", ...] + "keywords": ";;" } Rules: -- "summary" should describe the main content visible on the page. -- "status" should be "clear" if the page is legible, "blurry" if the text is - unreadable or the image is too low quality to summarize. -- "observed_titles" should list any headings, section titles, or chapter titles - visible on the page. Return an empty array if none are visible. +- "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. """ @@ -99,13 +97,7 @@ def tag_pages( strategy = plan.strategy if plan else PageProcessingStrategy.VLM_LITE if strategy == PageProcessingStrategy.SKIP_TAGGING: - results.append( - PageTagResult( - page_index=page.page_index, - status="skipped", - strategy_used="skip_tagging", - ) - ) + results.append(_tag_skip(page)) continue if strategy == PageProcessingStrategy.TEXT_ONLY: @@ -130,38 +122,85 @@ def tag_pages( len(results), vlm_calls, sum(1 for r in results if r.strategy_used == "text_only"), - sum(1 for r in results if r.status == "skipped"), + sum(1 for r in results if r.strategy_used == "skip_tagging"), ) return results +def _tag_skip(page: PageRenderResult) -> PageTagResult: + """Skip-tagging: preserve raw_text but no summary. + + If the page has no extractable text, mark as EMPTY. + """ + raw = page.raw_text.strip() + return PageTagResult( + page_index=page.page_index, + summary="" if raw else "EMPTY", + keywords=[], + strategy_used="skip_tagging", + ) + + def _tag_text_only(page: PageRenderResult) -> PageTagResult: - """Rules-based tag from raw text (no VLM).""" + """Use existing ``summary-full`` LLM prompt to extract summary + keywords. + + Falls back to raw text truncation if LLM is not available or fails. + """ raw = page.raw_text.strip() + if not raw: + return PageTagResult( + page_index=page.page_index, + summary="EMPTY", + keywords=[], + strategy_used="text_only", + ) + + # 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-chat") + 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] - if not summary: - summary = f"Page {page.page_index} (no extractable text)" - - # Simple heuristic: lines that look like headings (short, no trailing punct) - observed: list[str] = [] - for line in raw.splitlines()[:30]: - stripped = line.strip() - if ( - stripped - and 3 < len(stripped) < 100 - and not stripped.endswith((".", "。", ",", "，", ";", "；")) - and not stripped[0].isdigit() - and stripped[0].isupper() or not stripped[0].isascii() - ): - # Very rough heuristic; the real heading detection is in skeleton - pass - # For text_only mode, we don't attempt heading detection return PageTagResult( page_index=page.page_index, summary=summary, - status="clear", - observed_titles=observed, - strategy_used="text_only", + keywords=[], + strategy_used="text_only_fallback", ) @@ -184,8 +223,7 @@ def _tag_vlm_lite( result = PageTagResult( page_index=result.page_index, summary=result.summary, - status=result.status, - observed_titles=result.observed_titles, + keywords=result.keywords, strategy_used="text_only_budget_fallback", ) return result @@ -242,11 +280,12 @@ def _tag_vlm_lite( ) 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", "")), - status="clear" if data.get("status") != "blurry" else "blurry", - observed_titles=list(data.get("observed_titles") or []), + keywords=keywords, strategy_used="vlm_lite", ) except json.JSONDecodeError: @@ -256,19 +295,19 @@ def _tag_vlm_lite( page.page_index, attempt + 1, _MAX_JSON_RETRIES + 1, ) continue - # Final attempt failed: blurry degradation + # Final attempt failed: fallback to text_only logger.warning( - "[page_tagger] JSON retry exhausted for page {}; blurry fallback", + "[page_tagger] JSON retry exhausted for page {}; text_only fallback", page.page_index, ) - raw_summary = " ".join(page.raw_text.split())[:_RAW_TEXT_SUMMARY_LIMIT] - return PageTagResult( - page_index=page.page_index, - summary=raw_summary or f"Page {page.page_index} (VLM parse failed)", - status="blurry", - observed_titles=[], - strategy_used="vlm_lite_blurry_fallback", + 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 {}: {}", From c4689e78e619a5405fda4c2e12a599cd56214ff0 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 17 Jun 2026 14:52:05 +0800 Subject: [PATCH 4/5] feat: enhance Excel table parsing with hierarchical pathing, subtable support, and optimized evidence rendering with configurable character limits --- .../structure/hierarchy_locator.py | 29 ------ .../formats/excel/structure_parser.py | 14 +++ .../formats/excel/table_parser.py | 77 +++++++++++++--- .../tables/table_asset_writer.py | 10 ++- .../services/page_memory/memory_service.py | 61 +++++++++++-- .../app/services/page_memory/page_renderer.py | 10 ++- ...test_agentic_evidence_renderer_contract.py | 90 +++++++++++++++++++ .../contract/test_excel_parser_contract.py | 9 +- .../shared-python/shared/core/config/ai.py | 7 ++ .../retrieval/agentic/evidence/renderer.py | 71 +++++++++++++-- .../services/storage/zip_chunk_schema.py | 8 +- .../services/storage/zip_doc_navigation.py | 10 ++- 12 files changed, 326 insertions(+), 70 deletions(-) create mode 100644 apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index db2fd1c0c..3dabcdf02 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -18,7 +18,6 @@ ) TitleMatchSource = Literal[ - "exact", "anchored", "page_compact", "normalized", @@ -554,32 +553,6 @@ def _allowed_pages_between(start: int, end: int, allowed_pages: set[int]) -> lis return [page for page in range(start, end + 1) if page in allowed_pages] -def _find_exact_hits( - title: str, - scope_pages: list[int], - page_texts: dict[int, str], -) -> list[_LineHit]: - hits: list[_LineHit] = [] - needle = normalize_heading_text(title).casefold() - for page, line_index, line in _iter_lines(scope_pages, page_texts): - normalized_line = normalize_heading_text(line).casefold() - if needle and needle in normalized_line: - base = 1.0 - cleaned_line = normalize_heading_text(clean_toc_title(line)).casefold() - if normalized_line == needle or cleaned_line == needle: - base = 1.18 - hits.append( - _LineHit( - page=page, - line_index=line_index, - line=line.strip(), - source="exact", - score=_line_score(line=line, line_index=line_index, base=base), - ) - ) - return hits - - def _find_anchored_hits( title: str, scope_pages: list[int], @@ -715,7 +688,6 @@ def sort_key(hit: _LineHit) -> tuple[float, int, int, int]: best = ordered[0] pages = sorted({hit.page for hit in ordered}) confidence_by_source = { - "exact": 0.95, "anchored": 0.92, "page_compact": 0.9, "normalized": 0.84, @@ -747,7 +719,6 @@ def _preferred_source(hits: list[_LineHit]) -> TitleMatchSource: "normalized": 30, "token": 20, "printed_prior": 10, - "exact": 5, "h1_result": 60, "agent_vlm": 70, "agent_heuristic": 15, diff --git a/apps/worker/app/services/document_parser/formats/excel/structure_parser.py b/apps/worker/app/services/document_parser/formats/excel/structure_parser.py index bfe3553de..e8aa42e93 100644 --- a/apps/worker/app/services/document_parser/formats/excel/structure_parser.py +++ b/apps/worker/app/services/document_parser/formats/excel/structure_parser.py @@ -82,6 +82,14 @@ def parse_excel_structure( ) dataframe = result["df"] dataframe.attrs["row_header_cols"] = len(result["header_cols"]) + dataframe.attrs["sheet_name"] = selected_sheet_name + dataframe.attrs["subtable_index"] = index + 1 + dataframe.attrs["subtable_count"] = len(subtable_regions) + dataframe.attrs["subtable_title"] = ( + f"子表 {index + 1}" if len(subtable_regions) > 1 else "" + ) + dataframe.attrs["subtable_row_range"] = row_range + dataframe.attrs["subtable_col_range"] = col_range key = selected_sheet_name if index == 0 else f"{selected_sheet_name}_{index + 1}" logger.debug( f"Subtable '{key}': rows={row_range}, cols={col_range}, " @@ -94,6 +102,12 @@ def parse_excel_structure( result = _parse_subtable(worksheet, row_range, col_range, merged_ranges) dataframe = result["df"] dataframe.attrs["row_header_cols"] = len(result["header_cols"]) + dataframe.attrs["sheet_name"] = selected_sheet_name + dataframe.attrs["subtable_index"] = 1 + dataframe.attrs["subtable_count"] = 1 + dataframe.attrs["subtable_title"] = "" + dataframe.attrs["subtable_row_range"] = row_range + dataframe.attrs["subtable_col_range"] = col_range logger.debug( f"Sheet '{selected_sheet_name}': header_rows={result['header_rows']}, " f"header_cols={result['header_cols']}, " 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 3f5499ea3..3b2f1902a 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 @@ -27,7 +27,6 @@ from shared.core.exceptions.domain_exceptions import TableParsingException from shared.core.exceptions.knowhere_exception import KnowhereException -from shared.utils.chunk_refs import build_chunk_ref from app.services.common.file_loading import load_file_bytes from app.services.common.file_utils import path_handle from shared.utils.text_utils import tokenize2stw_remove @@ -150,6 +149,7 @@ def _parse_excel_sheet( time_stamp: str, ) -> list[ParsedRow]: parsed_rows: list[ParsedRow] = [] + sheet_attrs = dict(sheet_frame.attrs) try: table_frame = postprocess_tb(sheet_frame, drop=True) @@ -161,19 +161,22 @@ def _parse_excel_sheet( table_frame = _drop_source_row_columns(table_frame) row_header_cols = int(table_frame.attrs.get("row_header_cols", 0)) + logical_sheet_name = _sheet_name_from_attrs(sheet_attrs, fallback=sheet_name) + subtable_title = _subtable_title_from_attrs(sheet_attrs) _table_paths, table_html = parse_tb_contents( table_frame, - parent_dic={request.file_name: {sheet_name: {}}}, + parent_dic={request.file_name: {logical_sheet_name: {}}}, file_name=request.file_name, - sheet_name=sheet_name, + sheet_name=logical_sheet_name, row_header_cols=row_header_cols, ) parsed_rows.append( _write_excel_table_asset( request=request, - sheet_name=sheet_name, + sheet_name=logical_sheet_name, + subtable_title=subtable_title, table_frame=table_frame, table_html=table_html, time_stamp=time_stamp, @@ -208,6 +211,7 @@ def _write_excel_table_asset( *, request: ExcelWorkbookParseRequest, sheet_name: str, + subtable_title: str, table_frame: pd.DataFrame, table_html: str, time_stamp: str, @@ -218,9 +222,10 @@ def _write_excel_table_asset( sheet_name=sheet_name, llm_parameters=request.base_llm_paras, ) - table_index = f"table-{sheet_name}" + table_label = _table_label(sheet_name=sheet_name, subtable_title=subtable_title) + table_index = f"table-{table_label}" table_summary = f"{table_index}\n{summary}" if summary else table_index - effective_name = title or sheet_name + effective_name = title or table_label table_stem = path_handle( remove_spaces("table-" + effective_name), mode="clean_single", @@ -229,13 +234,9 @@ def _write_excel_table_asset( raise ValueError(f"Failed to sanitize Excel table name: {effective_name}") table_name = table_stem + ".html" table_html_string = BeautifulSoup(table_html, features="html.parser").prettify() - know_id = gen_str_codes(table_html + str(sheet_name)) - table_ref = build_chunk_ref(f"tables/{table_name}") - table_content = ( - f"{table_ref}\nTable summary:\n{table_summary}\nMain columns:\n{keywords}" - ) + know_id = gen_str_codes(table_html + str(table_label)) table_tokens = tokenize2stw_remove( - [table_content], + [table_summary, keywords], request.base_llm_paras["stopwords"], ) @@ -248,13 +249,61 @@ def _write_excel_table_asset( keywords=keywords, know_id=know_id, addtime=time_stamp, - content=table_content, + content=table_html_string, tokens=table_tokens, - length=len(table_html), + length=len(table_html_string), + path=_build_excel_table_path( + request=request, + sheet_name=sheet_name, + subtable_title=subtable_title, + ), + asset_path=f"tables/{table_name}", ) ) +def _sheet_name_from_attrs(attrs: dict[str, Any], *, fallback: str) -> str: + sheet_name = str(attrs.get("sheet_name") or fallback).strip() + return sheet_name or fallback + + +def _subtable_title_from_attrs(attrs: dict[str, Any]) -> str: + try: + subtable_count = int(attrs.get("subtable_count") or 1) + subtable_index = int(attrs.get("subtable_index") or 1) + except (TypeError, ValueError): + subtable_count = 1 + subtable_index = 1 + if subtable_count <= 1: + return "" + title = str(attrs.get("subtable_title") or "").strip() + return title or f"子表 {subtable_index}" + + +def _table_label(*, sheet_name: str, subtable_title: str) -> str: + return f"{sheet_name}-{subtable_title}" if subtable_title else sheet_name + + +def _build_excel_table_path( + *, + request: ExcelWorkbookParseRequest, + sheet_name: str, + subtable_title: str, +) -> str: + document_root = str(request.relative_root or request.file_name).strip() + parts = [ + _clean_path_segment(document_root), + _clean_path_segment(sheet_name), + ] + if subtable_title: + parts.append(_clean_path_segment(subtable_title)) + return "/".join(part for part in parts if part) + + +def _clean_path_segment(value: str) -> str: + return str(value).strip().replace("/", "_").replace("\\", "_") + + def _summarize_excel_table( *, table_frame: pd.DataFrame, 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 c8e5196f5..6ddeb57ec 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 @@ -4,6 +4,7 @@ from dataclasses import dataclass from app.services.document_parser.support.parser_rows import ParsedRow +from shared.utils.chunk_refs import build_chunk_ref @dataclass(frozen=True) @@ -18,6 +19,8 @@ class TableAssetInput: content: str | None = None tokens: str = "" length: int | None = None + path: str | None = None + asset_path: str | None = None def write_table_asset(table_input: TableAssetInput) -> ParsedRow: @@ -28,10 +31,13 @@ def write_table_asset(table_input: TableAssetInput) -> ParsedRow: 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 + row_type = "table" + if table_input.asset_path: + row_type = f"table\n{build_chunk_ref(table_input.asset_path)}" return ParsedRow( content=row_content, - path=f"tables/{table_filename}", - type="table", + path=table_input.path or f"tables/{table_filename}", + type=row_type, keywords=table_input.keywords, summary=table_input.summary, know_id=table_input.know_id, diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 12852e882..20d8e08c7 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -121,14 +121,21 @@ def _build_page_dataframe( if page_count <= 0: return pd.DataFrame(columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) - # ── budget (page_tagging envelope for VLM calls) ────────────────── + # ── unified budget (page_locate + page_tagging share one tracker) ── 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))) + ) budget = BudgetTracker( plan_budget=0, - visual_budget=page_tagging_budget, + 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, @@ -136,22 +143,28 @@ def _build_page_dataframe( }, ) + # ── build ToolContext for sub-agent VLM calls ───────────────────── + ctx = _build_page_ctx( + pdf_path=pdf_path, + job_id=filename, + output_dir=output_dir, + page_count=page_count, + budget=budget, + ) + # ── C4: skeleton (from profile anatomy) ─────────────────────────── + page_texts = read_page_texts( + pdf_path, list(range(1, page_count + 1)), timeout=300, + ) if anatomy is not None: - page_texts = read_page_texts( - pdf_path, list(range(1, page_count + 1)), timeout=300, - ) skeletons = extract_section_skeletons( anatomy=anatomy, filename=filename, page_texts=page_texts, - ctx=None, # no VLM for skeleton in page mode + ctx=ctx, ) else: skeletons = [] - page_texts = read_page_texts( - pdf_path, list(range(1, page_count + 1)), timeout=300, - ) logger.info( "[page_memory] C4 skeleton: {} sections from anatomy", len(skeletons), @@ -164,6 +177,7 @@ def _build_page_dataframe( page_count=page_count, output_dir=output_dir, page_features=page_features, + page_texts=page_texts, ) # ── C2: page plan ───────────────────────────────────────────────── @@ -260,6 +274,35 @@ def _build_page_dataframe( return pd.DataFrame(rows, columns=pd.Index([*PARSER_ROW_COLUMNS, "extra_metadata"])) +def _build_page_ctx( + *, + pdf_path: str, + job_id: str, + output_dir: str, + page_count: int, + budget: Any, +) -> ToolContext: + """Construct a ToolContext for C4 sub-agent and C3 tagger VLM calls.""" + blackboard = AgentBlackboard() + blackboard.page_count = page_count + vlm_model = os.environ.get("IMAGE_MODEL") + reason_model = ( + os.environ.get("PAGE_LOCATE_REASON_MODEL") + or os.environ.get("NORMOL_MODEL") + ) + return ToolContext( + pdf_path=pdf_path, + job_id=job_id, + blackboard=blackboard, + budget=budget, + trace=None, + output_dir=output_dir, + settings={ + "vlm_model": vlm_model, + "model": reason_model, + }, + ) + # ── whole_doc builder (PR3, unchanged) ──────────────────────────────── diff --git a/apps/worker/app/services/page_memory/page_renderer.py b/apps/worker/app/services/page_memory/page_renderer.py index ca6a02371..85d6442b3 100644 --- a/apps/worker/app/services/page_memory/page_renderer.py +++ b/apps/worker/app/services/page_memory/page_renderer.py @@ -47,6 +47,7 @@ def render_document_pages( page_count: int, output_dir: str, page_features: list[PageFeature] | None = None, + page_texts: dict[int, str] | None = None, ctx: ToolContext | None = None, dpi: int = 144, timeout: int = 300, @@ -64,6 +65,10 @@ def render_document_pages( page_features: If available, dimensions are read from here (avoiding a second PyMuPDF open). Otherwise falls back to 0/0/False. + page_texts: + Pre-read page texts ``{page_index: text}``. If provided, skips + the internal ``read_page_texts`` call (avoids a redundant + PyMuPDF subprocess). ctx: An optional ``ToolContext`` forwarded to ``render_pages``. If *None*, a lightweight context is constructed internally. @@ -81,8 +86,9 @@ def render_document_pages( if not pages: return [] - # ── raw text ────────────────────────────────────────────────────── - page_texts = read_page_texts(pdf_path, pages, timeout=timeout) + # ── raw text (reuse caller's data if provided) ──────────────────── + if page_texts is None: + page_texts = read_page_texts(pdf_path, pages, timeout=timeout) # ── full-resolution PNGs ────────────────────────────────────────── if ctx is not None: diff --git a/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py b/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py new file mode 100644 index 000000000..a3de92fad --- /dev/null +++ b/apps/worker/tests/contract/test_agentic_evidence_renderer_contract.py @@ -0,0 +1,90 @@ +from shared.services.retrieval.agentic.evidence.renderer import render_leaf_chunks + + +def test_render_direct_small_table_chunk_includes_asset_url() -> None: + parts: list[str] = [] + render_leaf_chunks( + parts, + [ + { + "chunk_id": "table-1", + "chunk_type": "table", + "content": "
企业名称
", + "file_path": "tables/table-企业入驻信息表.html", + } + ], + " ", + asset_lookup={"table-1": "http://localhost:4566/table.html?signature=test"}, + ) + + rendered = "\n".join(parts) + assert "[Table: http://localhost:4566/table.html?signature=test]" in rendered + assert "
企业名称
" in rendered + + +def test_render_direct_large_table_chunk_uses_asset_stub(monkeypatch) -> None: + monkeypatch.setenv("RETRIEVAL_AGENTIC_INLINE_TABLE_CHAR_LIMIT", "10") + + parts: list[str] = [] + render_leaf_chunks( + parts, + [ + { + "chunk_id": "table-1", + "chunk_type": "table", + "content": "
企业名称
", + "file_path": "tables/table-企业入驻信息表.html", + "source_chunk_path": "企业信息汇总260509 (1).xlsx/企业批量录入", + "chunk_metadata": { + "summary": "table-企业批量录入\n企业入驻信息登记模板", + "keywords": ["企业信息", "入驻管理"], + }, + } + ], + " ", + asset_lookup={"table-1": "http://localhost:4566/table.html?signature=test"}, + ) + + rendered = "\n".join(parts) + 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 + assert "
企业名称
" not in rendered + + +def test_render_connected_table_chunk_includes_asset_url() -> None: + parts: list[str] = [] + render_leaf_chunks( + parts, + [ + { + "chunk_id": "text-1", + "chunk_type": "text", + "content": "见表 [tables/table-1.html]", + "chunk_metadata": { + "connect_to": [ + { + "target": "table-1", + "ref": "[tables/table-1.html]", + } + ] + }, + }, + { + "chunk_id": "table-1", + "chunk_type": "table", + "content": "
A
", + "file_path": "tables/table-1.html", + }, + ], + " ", + asset_lookup={"table-1": "http://localhost:4566/table-1.html?signature=test"}, + ) + + rendered = "\n".join(parts) + assert "[Table: http://localhost:4566/table-1.html?signature=test]" in rendered + assert "
A
" in rendered diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py index 6f4935cf2..694992d6b 100644 --- a/apps/worker/tests/contract/test_excel_parser_contract.py +++ b/apps/worker/tests/contract/test_excel_parser_contract.py @@ -51,9 +51,12 @@ def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( parsed_df = parse_output.parsed_df assert full_output_dir.endswith("budget.xlsx") assert parsed_df is not None - assert parsed_df["type"].tolist() == ["table"] - assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] + assert parsed_df["type"].iloc[0].startswith("table") + 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(" list[str]: + 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: + lines.append(f"Table path: {table_path}") + 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 "" + if summary: + lines.append("Table summary:") + lines.extend(str(summary).split("\n")) + + keywords = metadata.get("keywords") if isinstance(metadata, dict) else [] + if isinstance(keywords, list) and keywords: + keyword_text = ";".join(str(keyword) for keyword in keywords if str(keyword).strip()) + if keyword_text: + lines.append("Main columns:") + lines.append(keyword_text) + return lines + + +def _inline_table_char_limit() -> int: + return int(os.getenv(INLINE_TABLE_CHAR_LIMIT_ENV, str(DEFAULT_INLINE_TABLE_CHAR_LIMIT))) def _infer_child_sort_order(child: DocTreeNode) -> float: @@ -215,4 +273,3 @@ def _infer_child_sort_order(child: DocTreeNode) -> float: grandchild_order = _infer_child_sort_order(grandchild) min_order = min(min_order, grandchild_order) return min_order - 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 5b4ea35c5..ece59d38a 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -20,6 +20,8 @@ def calculate_statistics(self, chunks: list[dict[str, Any]]) -> dict[str, Any]: image_chunks = 0 table_chunks = 0 + page_chunks = 0 + for chunk in chunks: chunk_type = chunk.get("type", "") normalized_type = _normalize_chunk_type(chunk_type) @@ -27,6 +29,8 @@ def calculate_statistics(self, chunks: list[dict[str, Any]]) -> dict[str, Any]: image_chunks += 1 elif normalized_type == "table": table_chunks += 1 + elif normalized_type == "page": + page_chunks += 1 else: text_chunks += 1 @@ -35,9 +39,7 @@ def calculate_statistics(self, chunks: list[dict[str, Any]]) -> dict[str, Any]: "text_chunks": text_chunks, "image_chunks": image_chunks, "table_chunks": table_chunks, - "page_chunks": sum( - 1 for c in chunks if _normalize_chunk_type(c.get("type", "")) == "page" - ), + "page_chunks": page_chunks, "total_pages": None, } 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 2cd721a33..ae22f47d6 100644 --- a/packages/shared-python/shared/services/storage/zip_doc_navigation.py +++ b/packages/shared-python/shared/services/storage/zip_doc_navigation.py @@ -39,6 +39,7 @@ def build_doc_nav( source_file_name: str, ) -> dict[str, Any]: text_chunks: list[dict[str, Any]] = [] + table_section_candidates: list[dict[str, Any]] = [] image_resources: list[dict[str, Any]] = [] table_resources: list[dict[str, Any]] = [] @@ -77,6 +78,12 @@ def build_doc_nav( "summary": summary or content_preview, } ) + table_section_candidates.append( + { + "path": path, + "summary": summary or content_preview, + } + ) elif chunk_type == "page": stats["page_chunks"] += 1 # Page chunks participate in section tree like text chunks @@ -95,8 +102,9 @@ def build_doc_nav( } ) + section_chunks = text_chunks or table_section_candidates sections = self._build_section_tree( - text_chunks, + section_chunks, source_file_name=source_file_name, ) stats["max_depth"] = _max_depth(sections) From ec49b4e8547a4deba9c9e28967a08033969ccd66 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Wed, 17 Jun 2026 14:58:18 +0800 Subject: [PATCH 5/5] fix: resolve lint and type errors (unused imports, unused variable, optional ctx/last_verify narrowing) --- .../services/document_agent/structure/page_locate_subagent.py | 3 ++- apps/worker/app/services/page_memory/memory_service.py | 2 -- apps/worker/app/services/page_memory/page_renderer.py | 2 -- apps/worker/app/services/page_memory/page_section_mapper.py | 1 - apps/worker/app/services/page_memory/page_tagger.py | 2 +- 5 files changed, 3 insertions(+), 7 deletions(-) diff --git a/apps/worker/app/services/document_agent/structure/page_locate_subagent.py b/apps/worker/app/services/document_agent/structure/page_locate_subagent.py index b6c3cbfc9..72dbd8084 100644 --- a/apps/worker/app/services/document_agent/structure/page_locate_subagent.py +++ b/apps/worker/app/services/document_agent/structure/page_locate_subagent.py @@ -20,7 +20,6 @@ from __future__ import annotations import json -import time from dataclasses import dataclass, field from typing import Any, Callable, cast @@ -183,6 +182,7 @@ def locate(self, *, title: str, printed_page: int | None = None) -> SubAgentResu return SubAgentResult(None, transcript, len(transcript), "exhausted") def _dispatch(self, name: str, args: dict[str, Any]): + assert self.ctx is not None, "dispatch requires a non-None ToolContext" return REGISTRY.dispatch(name, self.ctx, args) def _coerce_pages(self, raw: Any) -> list[int]: @@ -211,6 +211,7 @@ def _build_match( grep_hit = seen_grep.get(page) or {} vlm_confirmed = bool(last_verify and last_verify.get("selected_page") == page) if vlm_confirmed: + assert last_verify is not None # guaranteed by vlm_confirmed check source = cast(Any, last_verify.get("source") or "agent_vlm") confidence = float(last_verify.get("confidence") or 0.7) reason = last_verify.get("reason") diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 20d8e08c7..06ad49220 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -17,7 +17,6 @@ from app.services.page_memory.normalizer import normalize_to_pdf from loguru import logger -from shared.core.exceptions.domain_exceptions import ValidationException @dataclass(frozen=True) @@ -207,7 +206,6 @@ def _build_page_dataframe( # ── C7: assemble DataFrame rows ────────────────────────────────── tag_map = {t.page_index: t for t in tags} render_map = {r.page_index: r for r in rendered} - plan_map = {p.page_index: p for p in plans} # Build page → PageLabel.kind lookup label_map: dict[int, str] = {} diff --git a/apps/worker/app/services/page_memory/page_renderer.py b/apps/worker/app/services/page_memory/page_renderer.py index 85d6442b3..fdada1ac2 100644 --- a/apps/worker/app/services/page_memory/page_renderer.py +++ b/apps/worker/app/services/page_memory/page_renderer.py @@ -7,9 +7,7 @@ from __future__ import annotations -import os from dataclasses import dataclass -from typing import Any from loguru import logger diff --git a/apps/worker/app/services/page_memory/page_section_mapper.py b/apps/worker/app/services/page_memory/page_section_mapper.py index 23f275233..776393fe1 100644 --- a/apps/worker/app/services/page_memory/page_section_mapper.py +++ b/apps/worker/app/services/page_memory/page_section_mapper.py @@ -15,7 +15,6 @@ from dataclasses import dataclass, field from enum import Enum -from typing import Any from loguru import logger diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index d4a48eb58..84cc48323 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -15,7 +15,7 @@ import json import os from dataclasses import dataclass, field -from typing import Any, Literal, cast +from typing import Any, cast from loguru import logger