diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index ecafa4a4..9ed83c8c 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -71,7 +71,7 @@ class RetrievalQueryRequest(BaseModel): ) use_agentic: bool | None = Field( None, - description="Deprecated mode hint retained for cache/request compatibility; retrieval always uses the agentic workflow.", + description="Set to true to enable agentic retrieval (LLM doc-select + navigation). Default (None/false) uses classic 3-channel top-K.", ) @field_validator("channels") diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index 1c2b23a9..a1e70474 100644 --- a/apps/api/tests/contract/test_demo_documents_contract.py +++ b/apps/api/tests/contract/test_demo_documents_contract.py @@ -279,7 +279,6 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( monkeypatch: MonkeyPatch, ) -> None: fake_result_storage = FakeResultStorage() - monkeypatch.setenv("RETRIEVAL_AGENTIC_ENABLED", "false") async with developer_api_client_factory() as api_client: import app.services.demo.source_materializer as source_materializer_module diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 803240dc..58de969c 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -358,34 +358,11 @@ async def test_should_return_empty_results_for_an_empty_query( @pytest.mark.asyncio -async def test_retrieval_should_ignore_false_agentic_hint_and_use_workflow( +async def test_retrieval_should_use_classic_topk_when_agentic_is_false( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], - monkeypatch: MonkeyPatch, ) -> None: - async def fake_run_request( - self: object, - db: AsyncSession, - *, - request: WorkflowRunRequest, - llm_fn: object | None = None, - ) -> WorkflowResult: - return WorkflowResult( - namespace=request.namespace, - query=request.query, - router_used="workflow_single_step", - answer_text="", - plan=QueryPlan.single_step(request.query), - referenced_chunks=[], - results=[], - ) - - monkeypatch.setattr( - "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator.run_request", - fake_run_request, - ) - async with developer_api_client_factory() as api_client: await _seed_retrieval_document( user_id="local-dev-user", @@ -414,8 +391,7 @@ async def fake_run_request( assert response.status_code == 200 response_json = cast(dict[str, object], response.json()) - assert response_json["router_used"] == "workflow_single_step" - assert response_json["results"] == [] + assert response_json["router_used"] == "classic_topk" @pytest.mark.asyncio diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py index 48c9653c..4630ff43 100644 --- a/apps/worker/app/services/document_agent/coordinator.py +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -164,7 +164,6 @@ def _run_structural(self) -> PageAnatomyMap: profile, initial_decision, _planner_result = self._propose_profile( actor="planner" ) - self._run_h1_boundary_pipeline() executor_result = ReActExecutor( self.ctx, registry=REGISTRY, @@ -209,7 +208,6 @@ def _run_lightweight_anatomy( ) else: self._ensure_disabled_toc_placeholder() - self._run_h1_boundary_pipeline() if skip_shard_plan: # Page-based track processes pages individually via VLM and never # consumes the shard plan; only build_anatomy_map's invariant needs @@ -385,8 +383,3 @@ def _run_toc_extraction_pipeline(self) -> None: actor=f"toc:{tool_name}", ) - def _run_h1_boundary_pipeline(self) -> None: - self._dispatch_profile_tool( - tool_name="match.h1_pages", - actor="toc:match.h1_pages", - ) diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 2518e282..8861f631 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -176,11 +176,9 @@ class Shard: page_start: int page_end: int page_offset: int - anchor_type: Literal["h1_boundary", "blank_separator", "forced_max_size"] + anchor_type: Literal["h1_boundary", "blank_separator", "forced_max_size", "toc_chapter_boundary"] anchor_evidence: str confidence: float - split_depth: int = 1 # 1=H1 cut, 2=H2 cut, etc. - is_continuation: bool = False # True for continuation shards that don't contain parent heading def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -218,10 +216,11 @@ class PageAnatomyMap: page_features: list[PageFeature] page_labels: list[PageLabel] toc_result: TocResult - h1_result: H1BoundaryResult shard_plan: ShardPlan + h1_result: H1BoundaryResult | None = None document_profile: DocumentProfile | None = None toc_hierarchies: list[dict[str, Any]] | None = None + toc_page_offset: int | None = None global_signals: dict[str, Any] = field(default_factory=dict) trace_summary: dict[str, Any] = field(default_factory=dict) created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) @@ -236,12 +235,13 @@ def to_dict(self) -> dict[str, Any]: "page_features": [feature.to_dict() for feature in self.page_features], "page_labels": [label.to_dict() for label in self.page_labels], "toc_result": self.toc_result.to_dict(), - "h1_result": self.h1_result.to_dict(), + "h1_result": self.h1_result.to_dict() if self.h1_result else None, "shard_plan": self.shard_plan.to_dict(), "document_profile": self.document_profile.to_dict() if self.document_profile else None, "toc_hierarchies": self.toc_hierarchies, + "toc_page_offset": self.toc_page_offset, "global_signals": dict(self.global_signals), "trace_summary": dict(self.trace_summary), "created_at": self.created_at.isoformat(), diff --git a/apps/worker/app/services/document_agent/registry.py b/apps/worker/app/services/document_agent/registry.py index 80a17531..9fd4966c 100644 --- a/apps/worker/app/services/document_agent/registry.py +++ b/apps/worker/app/services/document_agent/registry.py @@ -155,9 +155,5 @@ def has_toc_hierarchies(blackboard: AgentBlackboard) -> tuple[bool, str]: return bool(blackboard.toc_hierarchies), "toc_hierarchies missing; call extract_toc first" -def has_h1_result(blackboard: AgentBlackboard) -> tuple[bool, str]: - return blackboard.h1_result is not None, "h1_result missing; call match_h1 first" - - def has_shard_plan(blackboard: AgentBlackboard) -> tuple[bool, str]: return blackboard.shard_plan is not None, "shard_plan missing; call propose_shard first" diff --git a/apps/worker/app/services/document_agent/state.py b/apps/worker/app/services/document_agent/state.py index bcdb1a1c..c717cbf0 100644 --- a/apps/worker/app/services/document_agent/state.py +++ b/apps/worker/app/services/document_agent/state.py @@ -37,6 +37,7 @@ class AgentBlackboard: toc_result: TocResult | None = None toc_hierarchies: list[dict[str, Any]] | None = None h1_result: H1BoundaryResult | None = None + toc_page_offset: int | None = None shard_plan: ShardPlan | None = None validation_report: dict[str, Any] | None = None verdict: AgentVerdict | None = None diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index f322f6d6..ac5c2990 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -6,7 +6,6 @@ from . import find_toc_anchor_pages as find_toc_anchor_pages # noqa: F401 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 diff --git a/apps/worker/app/services/document_agent/tools/match_h1_pages.py b/apps/worker/app/services/document_agent/tools/match_h1_pages.py deleted file mode 100644 index 3e4d9ed0..00000000 --- a/apps/worker/app/services/document_agent/tools/match_h1_pages.py +++ /dev/null @@ -1,285 +0,0 @@ -"""Match TOC level-1 headings to body pages via PyMuPDF text search.""" - -from __future__ import annotations - -import base64 -import json -import os -import time -from typing import Any, cast - -from app.services.document_agent.manifest import ( - H1BoundaryResult, - H1Candidate, - ToolContext, - ToolResult, -) -from app.services.document_agent.pdf_text import read_page_texts -from app.services.document_agent.registry import has_toc_result, register_tool -from app.services.document_agent.visual import render_pages -from app.services.document_parser.structure.body_boundary import ( - clean_toc_title, - extract_level1_titles, - normalize_heading_text, -) -from loguru import logger - - -# ── C1: Unified grep matching ──────────────────────────────────────────── - - -def grep_titles_in_pages( - titles: list[str], - search_pages: list[int], - page_texts: dict[int, str], - *, - source: str = "toc_grep", - confidence: float = 0.88, -) -> tuple[list[H1Candidate], list[str]]: - """Grep a list of titles across specified pages, returning match results. - - H1/H2 share this function. Callers control scope via *titles* and - *search_pages*. - - Returns: - (matched_candidates, unmatched_titles) - """ - candidates: list[H1Candidate] = [] - unmatched: list[str] = [] - - for title in titles: - normalized_title = normalize_heading_text(title) - found = False - for page in search_pages: - text = page_texts.get(page, "") - if normalized_title in normalize_heading_text(text): - matched_line = "" - for line in text.splitlines(): - if normalized_title in normalize_heading_text(line): - matched_line = line.strip()[:100] - break - candidates.append( - H1Candidate( - title=title, - page=page, - confidence=confidence, - matched_line=matched_line, - source=source, # type: ignore[arg-type] - evidence={ - "normalized_needle": normalized_title, - "page_text_length": len(text), - }, - ) - ) - found = True - break # First match per title - if not found: - unmatched.append(title) - - # Deduplicate by page – keep first hit - seen: set[int] = set() - deduped: list[H1Candidate] = [] - for c in candidates: - if c.page not in seen: - seen.add(c.page) - deduped.append(c) - - return deduped, unmatched - - -def extract_children_titles( - toc_hierarchies: list[dict[str, Any]], - parent_title: str, -) -> list[str]: - """Extract level-2 titles under a given H1 parent from toc_with_level.""" - titles: list[str] = [] - for hier in toc_hierarchies or []: - entries = hier.get("toc_with_level", []) - in_scope = False - for entry in entries: - if entry.get("level") == 1: - cleaned = clean_toc_title(entry.get("heading", "")) - in_scope = normalize_heading_text(cleaned) == normalize_heading_text( - parent_title - ) - continue - if in_scope and entry.get("level") == 2: - cleaned = clean_toc_title(entry.get("heading", "")) - if cleaned and len(cleaned) >= 2: - titles.append(cleaned) - return titles - - -def _extract_level1_titles(toc_hierarchies: list[dict[str, Any]]) -> list[str]: - return extract_level1_titles(toc_hierarchies) - - -# ── C2: Lazy VLM verification ──────────────────────────────────────────── - - -def verify_section_start( - *, - page: int, - title: str, - ctx: ToolContext, -) -> bool: - """VLM-confirm whether *page* is the start of a section titled *title*. - - Used for lazy verification before committing a shard cut. - If VLM is unavailable (no model / budget exhausted / render fails), - returns ``True`` (trust GREP). - """ - model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") - if not model: - return True # No VLM → trust GREP - - # Render 1 page PNG - png_items = render_pages( - ctx, [page], folder_name="verify_pages", prefix="verify", timeout=60, - ) - if not png_items: - return True # Render failed → trust GREP - - prompt = ( - f"This is page {page} of a PDF document.\n" - f"Question: Is this page the START of a section titled '{title}'?\n" - "Criteria: The title appears as a prominent heading/title on this page, " - "not merely mentioned in body text.\n" - 'Return JSON: {"is_section_start": true/false, "reason": "brief"}' - ) - est = 800 # ~800 tokens for 1 image - stage = "structural_react" - if not ctx.budget.try_reserve("visual", est, stage=stage): - return True # Budget exhausted → trust GREP - - try: - png_path = str(png_items[0]["png_path"]) - with open(png_path, "rb") as f: - img_b64 = base64.b64encode(f.read()).decode() - content_parts: list[dict[str, Any]] = [ - {"type": "text", "text": prompt}, - {"type": "text", "text": f"\n--- Page {page} ---"}, - { - "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) - raw, usage = client.chat_completion_with_usage( - messages=cast(Any, [{"role": "user", "content": content_parts}]), - model=model, - temperature=0.0, - max_tokens=256, - response_format={"type": "json_object"}, - usage_task="document_agent.match_h1_pages", - ) - ctx.budget.commit( - "visual", - actual=usage.get("total_tokens", est), - est=est, - stage=stage, - ) - data = json.loads(raw) - result = bool(data.get("is_section_start", True)) - logger.info( - "[verify_section_start] page={} title='{}' → {} reason={}", - page, title[:30], result, data.get("reason", ""), - ) - return result - except Exception as exc: - ctx.budget.refund("visual", est=est, stage=stage) - logger.warning("[verify_section_start] VLM failed for page {}: {}", page, exc) - return True # VLM failure → trust GREP - - -# ── Tool registration ──────────────────────────────────────────────────── - - -@register_tool( - name="match.h1_pages", - description=( - "Match TOC level-1 headings to body pages using PyMuPDF substring search. " - "Produces H1Candidate list for downstream shard planning." - ), - preconditions=(has_toc_result,), -) -def match_h1_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: - start = time.monotonic() - - if not ctx.blackboard.toc_hierarchies: - logger.info("[match.h1_pages] no toc_hierarchies, skipping") - ctx.blackboard.h1_result = H1BoundaryResult( - method="none", - notes="No toc_hierarchies available for H1 matching", - ) - return ToolResult( - status="ok", - payload={"h1_count": 0}, - latency_ms=int((time.monotonic() - start) * 1000), - ) - - level1_titles = _extract_level1_titles(ctx.blackboard.toc_hierarchies) - if not level1_titles: - logger.info("[match.h1_pages] no level-1 titles in toc_hierarchies") - ctx.blackboard.h1_result = H1BoundaryResult( - method="toc_grep", - notes="toc_hierarchies contained no level-1 entries", - ) - return ToolResult( - status="ok", - payload={"h1_count": 0}, - latency_ms=int((time.monotonic() - start) * 1000), - output_summary={"level1_titles": level1_titles}, - ) - - # Build exclusion set: TOC pages should not be searched - toc_page_set: set[int] = set() - if ctx.blackboard.toc_result: - toc_page_set.update(ctx.blackboard.toc_result.toc_pages) - - # Read text for all non-TOC pages - search_pages = sorted( - p - for p in range(1, ctx.blackboard.page_count + 1) - if p not in toc_page_set - ) - page_texts = read_page_texts(ctx.pdf_path, search_pages, timeout=300) - - # Delegate to unified grep function - h1_candidates, unmatched_titles = grep_titles_in_pages( - level1_titles, search_pages, page_texts, source="toc_exact_top", - ) - - matched_titles = [c.title for c in h1_candidates] - - ctx.blackboard.h1_result = H1BoundaryResult( - h1_candidates=h1_candidates, - method="toc_grep", - notes=( - f"Matched {len(matched_titles)}/{len(level1_titles)} level-1 titles. " - f"Unmatched: {unmatched_titles[:5]}" - ), - ) - - logger.info( - "[match.h1_pages] matched {}/{} level-1 titles to body pages: {}", - len(matched_titles), - len(level1_titles), - [(c.title[:20], c.page) for c in h1_candidates], - ) - - return ToolResult( - status="ok", - payload={"h1_count": len(h1_candidates)}, - latency_ms=int((time.monotonic() - start) * 1000), - output_summary={ - "level1_titles": level1_titles, - "matched": [(c.title, c.page) for c in h1_candidates], - "unmatched": unmatched_titles, - }, - ) diff --git a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py index 7ae7f019..1e17a220 100644 --- a/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -21,7 +21,6 @@ def _artifact_dir(ctx: ToolContext) -> Path: def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: if not ( ctx.blackboard.toc_result - and ctx.blackboard.h1_result and ctx.blackboard.shard_plan ): raise ValueError("cannot build anatomy map from incomplete blackboard") @@ -36,6 +35,7 @@ def build_anatomy_map(ctx: ToolContext) -> PageAnatomyMap: shard_plan=ctx.blackboard.shard_plan, document_profile=ctx.blackboard.document_profile, toc_hierarchies=ctx.blackboard.toc_hierarchies, + toc_page_offset=ctx.blackboard.toc_page_offset, global_signals=ctx.blackboard.global_signals, trace_summary={ "budget": ctx.budget.snapshot(), diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 7b50c98d..1ca3fb80 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -8,24 +8,316 @@ from typing import Any from app.services.document_agent.manifest import ( - H1Candidate, Shard, ShardPlan, ToolContext, ToolResult, ) -from app.services.document_agent.pdf_text import read_page_texts -from app.services.document_agent.registry import has_doc_stats, has_h1_result, has_toc_result, register_tool -from app.services.document_agent.tools.match_h1_pages import ( - extract_children_titles, - grep_titles_in_pages, - verify_section_start, -) +from app.services.document_agent.registry import has_doc_stats, has_toc_result, register_tool from app.services.document_agent.validators import single_shard_plan, validate_shard_plan from loguru import logger from shared.utils.token_estimate import estimate_tokens +def derive_leaf_cut_pages( + toc_hierarchies: list[dict[str, Any]] | None, + *, + offset_override: int | None = None, +) -> list[int]: + """Derive physical page numbers of TOC leaf nodes for shard splitting. + + Leaf nodes are entries in toc_with_level whose next sibling has level <= theirs + (i.e. they have no children). The offset from printed page to physical page is + either provided via offset_override (VLM-calibrated) or computed arithmetically + from toc_range and the first entry's page_number as a fallback. + """ + if not toc_hierarchies: + return [] + + all_pages: list[int] = [] + for hier in toc_hierarchies: + if hier.get("toc_range_unit") != "page": + continue + toc_range = hier.get("toc_range") + entries = hier.get("toc_with_level") + if not toc_range or not entries: + continue + if isinstance(entries, str): + entries = _parse_toc_with_level_entries(entries) + if not entries: + continue + + if offset_override is not None: + offset = offset_override + else: + toc_end_page = toc_range[1] if isinstance(toc_range, list) else toc_range + first_printed = next( + (e.get("page_number") for e in entries if e.get("page_number") is not None), + None, + ) + if first_printed is None: + continue + offset = (toc_end_page + 1) - first_printed + logger.warning( + "[propose_shard_plan] using arithmetic offset fallback: " + "toc_end={} first_printed={} offset={}", + toc_end_page, + first_printed, + offset, + ) + + for i, entry in enumerate(entries): + pn = entry.get("page_number") + if pn is None: + continue + is_leaf = ( + i == len(entries) - 1 + or entries[i + 1].get("level", 1) <= entry.get("level", 1) + ) + if is_leaf: + all_pages.append(pn + offset) + + return sorted(set(all_pages)) + + +def derive_chapter_boundaries( + toc_hierarchies: list[dict[str, Any]] | None, + *, + offset_override: int | None = None, + page_count: int, +) -> list[dict[str, Any]]: + """Extract chapter entries with physical page ranges for shard planning. + + Returns a flat list sorted by page_start: + [{"title": str, "level": int, "page_start": int, "page_end": int, + "page_span": int, "sub_entries": [...]}, ...] + + Includes all L1 entries. For any L1 whose span exceeds 200 pages, + its direct L2 children are included as sub_entries so the LLM can + split within it. + """ + if not toc_hierarchies: + return [] + + all_entries: list[dict[str, Any]] = [] + for hier in toc_hierarchies: + if hier.get("toc_range_unit") != "page": + continue + toc_range = hier.get("toc_range") + entries = hier.get("toc_with_level") + if not toc_range or not entries: + continue + if isinstance(entries, str): + entries = _parse_toc_with_level_entries(entries) + if not entries: + continue + + if offset_override is not None: + offset = offset_override + else: + toc_end_page = toc_range[1] if isinstance(toc_range, list) else toc_range + first_printed = next( + (e.get("page_number") for e in entries if e.get("page_number") is not None), + None, + ) + if first_printed is None: + continue + offset = (toc_end_page + 1) - first_printed + + # Collect all entries with physical pages + phys_entries: list[dict[str, Any]] = [] + for entry in entries: + pn = entry.get("page_number") + if pn is None: + continue + physical = pn + offset + if physical < 1 or physical > page_count: + continue + phys_entries.append({ + "title": entry.get("heading", ""), + "level": entry.get("level", 1), + "page_start": physical, + }) + + if not phys_entries: + continue + + # Compute page_end for each entry: next entry's page_start - 1 + for i, item in enumerate(phys_entries): + if i + 1 < len(phys_entries): + item["page_end"] = phys_entries[i + 1]["page_start"] - 1 + else: + item["page_end"] = page_count + item["page_span"] = item["page_end"] - item["page_start"] + 1 + + all_entries.extend(phys_entries) + + if not all_entries: + return [] + + # Build chapter-level structure: group by L1 with L2 sub_entries + min_level = min(e["level"] for e in all_entries) + chapters: list[dict[str, Any]] = [] + current_l1: dict[str, Any] | None = None + + for entry in all_entries: + if entry["level"] == min_level: + if current_l1 is not None: + chapters.append(current_l1) + current_l1 = {**entry, "sub_entries": []} + elif current_l1 is not None and entry["level"] == min_level + 1: + current_l1["sub_entries"].append(entry) + + if current_l1 is not None: + chapters.append(current_l1) + + # Recompute L1 page_end from the next L1's page_start - 1 + for i, chapter in enumerate(chapters): + if i + 1 < len(chapters): + chapter["page_end"] = chapters[i + 1]["page_start"] - 1 + else: + chapter["page_end"] = page_count + chapter["page_span"] = chapter["page_end"] - chapter["page_start"] + 1 + # Recompute sub_entry page_end within the L1's range + subs = chapter["sub_entries"] + for j, sub in enumerate(subs): + if j + 1 < len(subs): + sub["page_end"] = subs[j + 1]["page_start"] - 1 + else: + sub["page_end"] = chapter["page_end"] + sub["page_span"] = sub["page_end"] - sub["page_start"] + 1 + + return chapters + + +def split_toc_for_shard( + toc_hierarchies: list[dict[str, Any]] | None, + shard_page_start: int, + shard_page_end: int, + *, + offset_override: int | None = None, +) -> list[dict[str, Any]] | None: + """Build per-shard toc_hierarchies filtered to the shard's page range. + + For continuation shards (not starting at page 1), the ancestor chain of + the first entry is prepended so downstream heading prediction has the + full structural context. + """ + if not toc_hierarchies: + return None + + result: list[dict[str, Any]] = [] + for hier in toc_hierarchies: + if hier.get("toc_range_unit") != "page": + result.append(hier) + continue + toc_range = hier.get("toc_range") + entries = hier.get("toc_with_level") + if not toc_range or not entries: + continue + if isinstance(entries, str): + entries = _parse_toc_with_level_entries(entries) + if not entries: + continue + + if offset_override is not None: + offset = offset_override + else: + toc_end_page = toc_range[1] if isinstance(toc_range, list) else toc_range + first_printed = next( + (e.get("page_number") for e in entries if e.get("page_number") is not None), + None, + ) + if first_printed is None: + continue + offset = (toc_end_page + 1) - first_printed + + shard_entries: list[dict[str, Any]] = [] + first_idx: int | None = None + for idx, entry in enumerate(entries): + pn = entry.get("page_number") + if pn is None: + continue + physical = pn + offset + if shard_page_start <= physical <= shard_page_end: + if first_idx is None: + first_idx = idx + shard_entries.append(entry) + + if not shard_entries or first_idx is None: + continue + + # Prepend ancestor chain for continuation shards. Walk forward through + # every entry preceding the shard's first entry, maintaining a + # monotonic stack of "open" ancestors: an incoming entry closes out + # (pops) any stack entries at the same or deeper level before being + # pushed itself. A final pop against first_entry_level removes a + # trailing sibling that shares the same level as the shard's first + # entry (siblings are not ancestors). This is robust to non-monotonic + # level sequences (e.g. [L1, L2, L1, L3]), unlike a simple + # "smallest-unseen-level" scan. + first_entry_level = shard_entries[0].get("level", 1) + ancestors: list[dict[str, Any]] = [] + if first_entry_level > 1: + stack: list[dict[str, Any]] = [] + for ancestor in entries[:first_idx]: + ancestor_level = ancestor.get("level", 1) + while stack and stack[-1].get("level", 1) >= ancestor_level: + stack.pop() + stack.append(ancestor) + while stack and stack[-1].get("level", 1) >= first_entry_level: + stack.pop() + ancestors = [ + { + "heading": node.get("heading"), + "level": node.get("level", 1), + "page_number": None, + } + for node in stack + ] + + result.append({ + "toc_range": [shard_page_start, shard_page_end], + "toc_range_unit": "page", + "source": hier.get("source", "vlm_shard_split"), + "toc_with_level": ancestors + shard_entries, + }) + + return result if result else None + + +def _parse_toc_with_level_entries(markdown: str) -> list[dict[str, Any]]: + """Parse toc_with_level markdown table into list of dicts.""" + 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") + page_number = _safe_int(row.get("page_number")) + if heading and level: + entries.append({"heading": heading, "level": level, "page_number": page_number}) + return entries + + +def _safe_int(value: Any) -> int | None: + if value is None or value == "": + return None + try: + return int(value) + except (ValueError, TypeError): + return None + + def _thresholds(ctx: ToolContext) -> tuple[int, int, int]: threshold = int( ctx.settings.get("shard_threshold") @@ -45,18 +337,9 @@ def _thresholds(ctx: ToolContext) -> tuple[int, int, int]: def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> list[Shard]: shards: list[Shard] = [] previous = 0 - # Track which cuts came from H2 refinement to mark continuation shards - h2_cut_pages: set[int] = set() - for cut_page, _anchor_type, evidence, _confidence in cuts: - if evidence.startswith("H2 refine:"): - h2_cut_pages.add(cut_page) - for cut_page, anchor_type, evidence, confidence in cuts: if cut_page <= previous: continue - # A shard is continuation if it starts AFTER an H2 cut (previous cut was H2) - _is_continuation = previous in h2_cut_pages - _split_depth = 2 if (evidence.startswith("H2 refine:") or _is_continuation) else 1 shards.append( Shard( shard_index=len(shards), @@ -66,13 +349,10 @@ def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> anchor_type=anchor_type, # type: ignore[arg-type] anchor_evidence=evidence, confidence=confidence, - split_depth=_split_depth, - is_continuation=_is_continuation, ) ) previous = cut_page if previous < page_count: - _is_continuation = previous in h2_cut_pages shards.append( Shard( shard_index=len(shards), @@ -82,8 +362,6 @@ def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> anchor_type="forced_max_size", anchor_evidence="final shard", confidence=1.0, - split_depth=2 if _is_continuation else 1, - is_continuation=_is_continuation, ) ) return shards @@ -97,7 +375,7 @@ def _build_prompt( doc_stats: dict[str, Any], page_kind_counts: dict[str, int], toc_pages: list[int], - h1_pages: list[dict[str, Any]], + leaf_pages: list[int], profile: dict[str, Any] | None, visual_evidence: list[dict[str, Any]], grep_history: list[dict[str, Any]], @@ -109,19 +387,18 @@ def _build_prompt( "page_kind_counts": page_kind_counts, "doc_stats": doc_stats, "toc_pages": toc_pages, - "h1_pages": h1_pages, + "leaf_cut_pages": leaf_pages, "document_profile": profile, "visual_evidence": visual_evidence[-3:], "grep_history": grep_history[-3:], } return ( "You are a senior document parsing architect. Decide whether to split a PDF " - "and where to split it using document-scale features, TOC/H1 evidence, and " - "recent agent observations.\n" + "and where to split it using document-scale features and TOC leaf-node evidence.\n" "Rules:\n" "- Return strict JSON only.\n" - "- Prefer H1 start pages as semantic boundaries, cutting at page-1 when possible.\n" - "- Do not blindly split on every H1. Consider total page_count, spacing, min/max " + "- Prefer TOC leaf-node pages as semantic boundaries, cutting at page-1 when possible.\n" + "- Do not blindly split on every leaf node. Consider total page_count, spacing, min/max " "shard sizes, and over-fragmentation.\n" "- Prefer fewer, semantically coherent shards over many tiny shards.\n" "- Keep each cut rationale under 120 characters.\n" @@ -145,6 +422,69 @@ def _build_prompt( ) +def _build_chapter_prompt( + *, + page_count: int, + max_pages: int, + chapters: list[dict[str, Any]], +) -> str: + """Build LLM prompt for TOC-based shard planning using chapter boundaries.""" + chapter_list = [] + for ch in chapters: + item: dict[str, Any] = { + "title": ch["title"], + "level": ch["level"], + "page_start": ch["page_start"], + "page_end": ch["page_end"], + "page_span": ch["page_span"], + } + if ch.get("sub_entries"): + item["sub_entries"] = [ + { + "title": s["title"], + "level": s["level"], + "page_start": s["page_start"], + "page_end": s["page_end"], + "page_span": s["page_span"], + } + for s in ch["sub_entries"] + ] + chapter_list.append(item) + + payload = { + "page_count": page_count, + "max_pages_per_shard": max_pages, + "chapters": chapter_list, + } + return ( + "You are a document splitting architect. Given a PDF's chapter structure, " + "decide how to split it into shards for downstream parsing.\n" + "Rules:\n" + "- Return strict JSON only.\n" + "- Each shard must be <= max_pages_per_shard pages.\n" + "- Group adjacent chapters into shards to fill each shard as evenly as possible.\n" + "- Cut points must align with chapter boundaries (use the page_end of the last " + "chapter in the shard as cut_after_page).\n" + "- If a single chapter exceeds max_pages_per_shard, split it at one of its " + "sub_entries boundaries (use that sub_entry's page_end as cut_after_page).\n" + "- Prefer fewer shards over many small ones.\n" + "- Keep each cut rationale under 120 characters.\n" + "- If the total page_count <= max_pages_per_shard, return enabled=false.\n" + "Output schema:\n" + "{\n" + ' "enabled": boolean,\n' + ' "cuts": [\n' + ' {"cut_after_page": number, "anchor_type": "toc_chapter_boundary", ' + '"confidence": number, "rationale": string}\n' + " ],\n" + ' "reason": "llm_boundary_decision" | "not_needed",\n' + ' "rationale": string\n' + "}\n" + "Payload:\n" + + json.dumps(payload, ensure_ascii=False) + ) + + def _sanitize_rationale(text: str, max_length: int = 120) -> str: # Truncate overlong rationales but preserve H1 title references # which provide valuable semantic context for shard boundaries. @@ -204,7 +544,7 @@ def _parse_llm_plan( if not 1 <= cut_page < page_count: continue anchor_type = str(item.get("anchor_type") or "forced_max_size") - if anchor_type not in {"h1_boundary", "blank_separator", "forced_max_size"}: + if anchor_type not in {"h1_boundary", "blank_separator", "forced_max_size", "toc_chapter_boundary"}: anchor_type = "forced_max_size" confidence = float(item.get("confidence") or 0.5) cuts.append((cut_page, anchor_type, _sanitize_rationale(str(item.get("rationale") or rationale)), confidence)) @@ -219,188 +559,130 @@ def _deterministic_guardrail_plan( page_count: int, min_pages: int, max_pages: int, - h1_pages: list[int], + leaf_pages: list[int], ) -> tuple[list[tuple[int, str, str, float]], str]: cuts: list[tuple[int, str, str, float]] = [] previous = 0 while page_count - previous > max_pages: target = previous + max_pages eligible = [ - page for page in h1_pages if previous + 1 < page <= target + page for page in leaf_pages if previous + min_pages < page <= target ] if eligible: chosen = max(eligible) cut_page = chosen - 1 - cuts.append((cut_page, "h1_boundary", f"guardrail H1 start page {chosen}", 0.35)) + cuts.append((cut_page, "h1_boundary", f"guardrail leaf node at page {chosen}", 0.35)) previous = cut_page else: - break # No more H1 in range → leave oversized shard for H2 refinement + cut_page = previous + max_pages + cuts.append((cut_page, "forced_max_size", "no leaf node in range", 0.2)) + previous = cut_page return cuts, "too_large" -# ── C3: H2-aware shard refinement ──────────────────────────────────────── - - -def _find_h1_for_range( - h1_candidates: list[H1Candidate], - range_start: int, - range_end: int, -) -> str | None: - """Find the H1 title whose start page falls in [range_start+1, range_end].""" - for c in h1_candidates: - if range_start < c.page <= range_end: - return c.title - # Fallback: the H1 whose page is closest to and <= range_start+1 - best: H1Candidate | None = None - for c in h1_candidates: - if c.page <= range_start + 1: - if best is None or c.page > best.page: - best = c - return best.title if best else None - - -def _pick_and_verify_best_cut( - h2_candidates: list[H1Candidate], - shard_start: int, - shard_end: int, - min_pages: int, +def _deterministic_chapter_plan( + *, + chapters: list[dict[str, Any]], max_pages: int, - ctx: ToolContext, -) -> tuple[int, str, str, float] | None: - """Pick the H2 candidate that produces the most balanced sub-shards. - - Candidates are ranked by how close they split the shard to the midpoint. - Each candidate is VLM-verified before acceptance. - """ - if not h2_candidates: - return None - - shard_length = shard_end - shard_start - midpoint = shard_start + shard_length // 2 - - # Sort by distance to midpoint (most balanced first) - ranked = sorted(h2_candidates, key=lambda c: abs(c.page - midpoint)) - - for candidate in ranked: - cut_page = candidate.page - 1 # Cut *before* the H2 start page - left_len = cut_page - shard_start - right_len = shard_end - cut_page - if left_len < min_pages or right_len < min_pages: - continue - if left_len > max_pages or right_len > max_pages: - continue - # VLM verification - if not verify_section_start(page=candidate.page, title=candidate.title, ctx=ctx): - logger.info( - "[h2_refine] VLM rejected H2 cut at page {} ('{}')", - candidate.page, candidate.title[:30], - ) - continue - logger.info( - "[h2_refine] accepted H2 cut at page {} ('{}'), left={} right={}", - candidate.page, candidate.title[:30], left_len, right_len, - ) - return ( - cut_page, - "h1_boundary", - f"H2 refine: '{candidate.title[:60]}' at page {candidate.page}", - candidate.confidence * 0.9, # Slightly lower confidence than H1 - ) + page_count: int, + leaf_pages: list[int], +) -> tuple[list[tuple[int, str, str, float]], str]: + """Greedy chapter grouping when LLM is unavailable.""" + cuts: list[tuple[int, str, str, float]] = [] + shard_start = 0 + + for i, chapter in enumerate(chapters): + chapter_end = chapter["page_end"] + shard_span = chapter_end - shard_start + + if shard_span > max_pages: + # Current chapter alone exceeds max_pages; split within its sub_entries + subs = chapter.get("sub_entries") or [] + if subs: + for sub in subs: + sub_end = sub["page_end"] + if sub_end - shard_start > max_pages: + # cut before this sub_entry + cut_page = sub["page_start"] - 1 + if cut_page > shard_start: + cuts.append(( + cut_page, + "toc_chapter_boundary", + f"split within chapter at sub-entry: {sub['title'][:60]}", + 0.7, + )) + shard_start = cut_page + else: + # No sub_entries; fall back to leaf pages within this chapter + ch_leaf_pages = [ + p for p in leaf_pages + if chapter["page_start"] <= p <= chapter_end + ] + sub_previous = shard_start + while chapter_end - sub_previous > max_pages: + target = sub_previous + max_pages + eligible = [p for p in ch_leaf_pages if sub_previous + 20 < p <= target] + if eligible: + chosen = max(eligible) + cut_page = chosen - 1 + else: + cut_page = sub_previous + max_pages + cuts.append((cut_page, "forced_max_size", "oversized chapter, leaf fallback", 0.3)) + shard_start = cut_page + sub_previous = cut_page + + elif i + 1 < len(chapters): + next_chapter_end = chapters[i + 1]["page_end"] + next_shard_span = next_chapter_end - shard_start + if next_shard_span > max_pages: + # Adding next chapter would overflow; cut after current chapter + cuts.append(( + chapter_end, + "toc_chapter_boundary", + f"chapter boundary: {chapter['title'][:60]}", + 0.85, + )) + shard_start = chapter_end - return None + return cuts, "too_large" -def _refine_with_h2( - cuts: list[tuple[int, str, str, float]], +def _deterministic_no_toc_plan( + *, page_count: int, - min_pages: int, max_pages: int, - ctx: ToolContext, - h1_candidates: list[H1Candidate], -) -> list[tuple[int, str, str, float]]: - """Post-process cuts: split any shard that exceeds max_pages using H2 boundaries.""" - if not ctx.blackboard.toc_hierarchies: - return cuts - - refined: list[tuple[int, str, str, float]] = [] + low_content_pages: list[int], +) -> tuple[list[tuple[int, str, str, float]], str]: + """Deterministic shard plan using low-content pages as split candidates.""" + cuts: list[tuple[int, str, str, float]] = [] previous = 0 + while page_count - previous > max_pages: + target = previous + max_pages + # Look for a low-content page near the max boundary + eligible = [ + p for p in low_content_pages if previous + (max_pages - 20) < p <= target + ] + if eligible: + chosen = max(eligible) + cuts.append((chosen, "blank_separator", f"low-content page at {chosen}", 0.5)) + previous = chosen + else: + cut_page = previous + max_pages + cuts.append((cut_page, "forced_max_size", "no separator in range", 0.2)) + previous = cut_page + return cuts, "too_large" - # Build endpoints: each cut + the implicit final boundary - endpoints = [(cp, at, ev, cf) for cp, at, ev, cf in cuts] + [ - (page_count, "final", "", 1.0) - ] - - for cut_page, anchor_type, evidence, confidence in endpoints: - shard_length = cut_page - previous - if shard_length > max_pages: - # Try H2 refinement for this oversized shard - h1_title = _find_h1_for_range(h1_candidates, previous, cut_page) - h2_cut_found = False - if h1_title: - h2_titles = extract_children_titles( - ctx.blackboard.toc_hierarchies, h1_title, - ) - if h2_titles: - search_pages = list(range(previous + 1, cut_page + 1)) - page_texts = read_page_texts( - ctx.pdf_path, search_pages, timeout=120, - ) - h2_candidates, _ = grep_titles_in_pages( - h2_titles, search_pages, page_texts, - source="h2_refine", - ) - best = _pick_and_verify_best_cut( - h2_candidates, previous, cut_page, - min_pages, max_pages, ctx, - ) - if best: - refined.append(best) - h2_cut_found = True - logger.info( - "[h2_refine] split oversized shard [{}-{}] at page {}", - previous + 1, cut_page, best[0], - ) - else: - logger.warning( - "[h2_refine] no valid H2 cut for shard [{}-{}]", - previous + 1, cut_page, - ) - else: - logger.info( - "[h2_refine] no H2 titles found under H1 '{}' for shard [{}-{}]", - h1_title[:30], previous + 1, cut_page, - ) - else: - logger.info( - "[h2_refine] no H1 found for oversized shard [{}-{}]", - previous + 1, cut_page, - ) - - # Ultimate fallback: forced_max_size - if not h2_cut_found: - fallback_page = previous + max_pages - if fallback_page < cut_page: - refined.append(( - fallback_page, "forced_max_size", - "H2 refine fallback: forced max size", 0.2, - )) - logger.warning( - "[h2_refine] forced_max_size fallback at page {} for shard [{}-{}]", - fallback_page, previous + 1, cut_page, - ) - - # Append the original cut (skip the synthetic "final" endpoint) - if anchor_type != "final": - refined.append((cut_page, anchor_type, evidence, confidence)) - previous = cut_page - return refined +def _get_low_content_pages(ctx: ToolContext) -> list[int]: + """Extract low-content page numbers from page labels.""" + labels = ctx.blackboard.page_labels or [] + return sorted(label.page for label in labels if label.kind == "low_content") @register_tool( name="propose.shard_plan", - description="Ask the LLM to decide whether and where to split using profile, TOC, and H1 evidence.", - preconditions=(has_doc_stats, has_toc_result, has_h1_result), + description="Decide whether and where to split a long PDF using TOC chapter boundaries.", + preconditions=(has_doc_stats, has_toc_result), ) def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() @@ -415,93 +697,98 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - h1_candidates = ( - ctx.blackboard.h1_result.h1_candidates if ctx.blackboard.h1_result else [] - ) - h1_pages = [{"title": item.title, "page": item.page} for item in h1_candidates] - model = ctx.settings.get("model") - prompt = _build_prompt( + offset_hint: int | None = None + if ctx.blackboard.toc_hierarchies: + from app.services.document_agent.structure.hierarchy_locator import extract_toc_nodes + from app.services.page_memory.skeleton_extractor import _calibrate_offset_via_vlm + + nodes = extract_toc_nodes(ctx.blackboard.toc_hierarchies) + offset_hint, _ = _calibrate_offset_via_vlm( + nodes=nodes, + toc_hierarchies=ctx.blackboard.toc_hierarchies, + ctx=ctx, + page_texts={}, + page_count=page_count, + ) + ctx.blackboard.toc_page_offset = offset_hint + + # Try TOC chapter-based planning first + chapters = derive_chapter_boundaries( + ctx.blackboard.toc_hierarchies, + offset_override=offset_hint, page_count=page_count, - min_pages=min_pages, - max_pages=max_pages, - doc_stats=ctx.blackboard.doc_stats, - page_kind_counts=ctx.blackboard.global_signals.get("page_kind_counts", {}), - toc_pages=ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else [], - h1_pages=h1_pages, - profile=ctx.blackboard.document_profile.to_dict() - if ctx.blackboard.document_profile - else None, - visual_evidence=ctx.blackboard.global_signals.get("visual_inspections", []), - grep_history=ctx.blackboard.global_signals.get("grep_history", []), - ) - prompt_tokens_est = estimate_tokens(prompt) + ) if ctx.blackboard.toc_hierarchies else [] + warnings: list[str] = [] raw_response = "" rationale = "" llm_attempted = False - if model and ctx.budget.try_reserve("plan", prompt_tokens_est): - try: - llm_attempted = True - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - client = get_openai_client(model=model) - raw_response, usage = client.chat_completion_with_usage( - messages=[{"role": "user", "content": prompt}], - model=model, - temperature=0.0, - max_tokens=1600, - response_format={"type": "json_object"}, - usage_task="document_agent.propose_shard_plan", - ) - ctx.budget.commit("plan", actual=usage.get("total_tokens", prompt_tokens_est), est=prompt_tokens_est) - enabled, cuts, reason, rationale = _parse_llm_plan(raw_response, page_count, min_pages, max_pages) - if not enabled: - cuts = [] - reason = "not_needed" - except Exception as exc: - ctx.budget.refund("plan", est=prompt_tokens_est) - warnings.append(f"LLM shard decision failed; using guardrail plan: {exc}") - ctx.blackboard.global_signals.setdefault("degraded_reasons", []).append( - "shard_plan: llm_parse_failed" - ) - cuts, reason = _deterministic_guardrail_plan( - page_count=page_count, - min_pages=min_pages, + leaf_pages = derive_leaf_cut_pages(ctx.blackboard.toc_hierarchies, offset_override=offset_hint) + + if chapters: + # Path A: TOC chapter-based LLM decision + model = ctx.settings.get("model") + prompt = _build_chapter_prompt( + page_count=page_count, + max_pages=max_pages, + chapters=chapters, + ) + prompt_tokens_est = estimate_tokens(prompt) + + if model and ctx.budget.try_reserve("plan", prompt_tokens_est): + try: + llm_attempted = True + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw_response, usage = client.chat_completion_with_usage( + messages=[{"role": "user", "content": prompt}], + model=model, + temperature=0.0, + max_tokens=1600, + response_format={"type": "json_object"}, + usage_task="document_agent.propose_shard_plan", + ) + ctx.budget.commit("plan", actual=usage.get("total_tokens", prompt_tokens_est), est=prompt_tokens_est) + enabled, cuts, reason, rationale = _parse_llm_plan(raw_response, page_count, min_pages, max_pages) + if not enabled: + cuts = [] + reason = "not_needed" + except Exception as exc: + ctx.budget.refund("plan", est=prompt_tokens_est) + warnings.append(f"LLM chapter shard decision failed; using deterministic plan: {exc}") + ctx.blackboard.global_signals.setdefault("degraded_reasons", []).append( + "shard_plan: llm_parse_failed" + ) + cuts, reason = _deterministic_chapter_plan( + chapters=chapters, + max_pages=max_pages, + page_count=page_count, + leaf_pages=leaf_pages, + ) + rationale = "Deterministic chapter plan after LLM failure." + else: + if not model: + warnings.append("No model configured; using deterministic chapter plan.") + ctx.blackboard.global_signals.setdefault("degraded_reasons", []).append( + "shard_plan: no model" + ) + cuts, reason = _deterministic_chapter_plan( + chapters=chapters, max_pages=max_pages, - h1_pages=[item["page"] for item in h1_pages], - ) - rationale = "Guardrail plan after malformed LLM shard decision." - else: - if not model: - warnings.append("No model configured for shard decision; using guardrail plan.") - ctx.blackboard.global_signals.setdefault("degraded_reasons", []).append( - "shard_plan: no model" - ) - cuts, reason = _deterministic_guardrail_plan( page_count=page_count, - min_pages=min_pages, - max_pages=max_pages, - h1_pages=[item["page"] for item in h1_pages], - ) - rationale = "Guardrail plan without configured shard model." - else: - return ToolResult( - status="error", - error="Insufficient plan budget for shard decision.", - latency_ms=int((time.monotonic() - start) * 1000), - warnings=warnings, - debug={ - "prompt_excerpt": prompt[:4000], - "raw_response_excerpt": raw_response[:4000], - "llm_attempted": llm_attempted, - }, + leaf_pages=leaf_pages, ) - - # C3: H2 refinement – split any shard that still exceeds max_pages - if cuts: - cuts = _refine_with_h2( - cuts, page_count, min_pages, max_pages, ctx, h1_candidates, + rationale = "Deterministic chapter plan (no LLM)." + else: + # Path B: No TOC — purely deterministic using low-content pages + low_content_pages = _get_low_content_pages(ctx) + cuts, reason = _deterministic_no_toc_plan( + page_count=page_count, + max_pages=max_pages, + low_content_pages=low_content_pages, ) + rationale = "Deterministic plan from low-content page boundaries (no TOC)." shards = _cuts_to_shards(cuts, page_count) enabled = len(shards) > 1 @@ -528,11 +815,12 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: "valid": plan.validation.valid, }, latency_ms=int((time.monotonic() - start) * 1000), - tokens_used=ctx.budget.snapshot()["plan"]["used"], + tokens_used=ctx.budget.snapshot()["plan"]["used"] if llm_attempted else 0, input_summary={ "page_count": page_count, - "h1_count": len(h1_pages), - "model": model, + "chapter_count": len(chapters), + "leaf_page_count": len(leaf_pages), + "model": ctx.settings.get("model"), }, output_summary={ "enabled": plan.enabled, @@ -542,8 +830,7 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: }, warnings=warnings, debug={ - "prompt_excerpt": prompt[:4000], - "raw_response_excerpt": raw_response[:4000], + "raw_response_excerpt": raw_response[:4000] if raw_response else "", "llm_attempted": llm_attempted, }, ) diff --git a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py index c9e8273c..2e75594c 100644 --- a/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py +++ b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py @@ -8,7 +8,6 @@ from app.services.document_agent.manifest import PageAnatomyMap, ToolContext, ToolResult from app.services.document_agent.registry import ( - has_h1_result, has_shard_plan, has_toc_result, register_tool, @@ -31,13 +30,12 @@ def _thresholds(ctx: ToolContext) -> tuple[int, int]: @register_tool( name="validate.anatomy_map", description="Validate page anatomy, hierarchy hints, and shard coverage.", - preconditions=(has_toc_result, has_h1_result, has_shard_plan), + preconditions=(has_toc_result, has_shard_plan), ) def validate_current_anatomy(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: start = time.monotonic() if not ( ctx.blackboard.toc_result - and ctx.blackboard.h1_result and ctx.blackboard.shard_plan ): return ToolResult( diff --git a/apps/worker/app/services/document_agent/validators.py b/apps/worker/app/services/document_agent/validators.py index 12ea930b..fe26b432 100644 --- a/apps/worker/app/services/document_agent/validators.py +++ b/apps/worker/app/services/document_agent/validators.py @@ -86,11 +86,12 @@ def validate_anatomy_map( if label_pages != expected_pages: errors.append("page_labels do not cover every page") toc_pages = set(anatomy.toc_result.toc_pages) - for candidate in anatomy.h1_result.h1_candidates: - if candidate.page in toc_pages: - errors.append(f"h1 candidate points to toc page {candidate.page}") - if candidate.page < 1 or candidate.page > page_count: - errors.append(f"h1 candidate page {candidate.page} out of range") + if anatomy.h1_result: + for candidate in anatomy.h1_result.h1_candidates: + if candidate.page in toc_pages: + errors.append(f"h1 candidate points to toc page {candidate.page}") + if candidate.page < 1 or candidate.page > page_count: + errors.append(f"h1 candidate page {candidate.page} out of range") shard_report = validate_shard_plan( anatomy.shard_plan, page_count=page_count, diff --git a/apps/worker/app/services/document_parser/formats/pdf/parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py index f1a82911..fb4b8e03 100755 --- a/apps/worker/app/services/document_parser/formats/pdf/parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -117,6 +117,7 @@ def _parse_pdf_via_shards( bin_pack_shards, split_pdf, ) + from app.services.document_agent.tools.propose_shard_plan import split_toc_for_shard work_dir: str | None = None temp_shard_s3_keys: list[str] = [] @@ -261,7 +262,14 @@ def _predict_shard_headings( md_lines = merge_html_tables(md_lines) is_first_shard = shard_idx == 0 - shard_toc = toc_hierarchies + shard = merged_shards[shard_idx] + shard_toc = ( + toc_hierarchies if is_first_shard + else split_toc_for_shard( + toc_hierarchies, shard.page_start, shard.page_end, + offset_override=getattr(anatomy, "toc_page_offset", None), + ) + ) lines_with_heading = eval_md_headings( md_lines, @@ -313,15 +321,8 @@ def _predict_shard_headings( raise RuntimeError(f"Missing heading result for shard_{index}") complete_heading_results.append(result) - # Compute level offsets: continuation shards get shifted deeper. - shard_offsets: list[int] = [] - for shard in agent_shards[: len(complete_heading_results)]: - if shard.is_continuation: - shard_offsets.append(max(shard.split_depth - 1, 0)) - else: - shard_offsets.append(0) - if any(offset > 0 for offset in shard_offsets): - logger.info(f"📐 Shard level offsets: {shard_offsets}") + # No level offsets needed — leaf-node splitting produces self-contained shards + shard_offsets: list[int] = [0] * len(complete_heading_results) all_lines_with_heading: list[str] = merge_shard_lines( [result.lines_with_heading for result in complete_heading_results], diff --git a/apps/worker/app/services/page_memory/_utils.py b/apps/worker/app/services/page_memory/_utils.py index 27d9d6c5..219889bf 100644 --- a/apps/worker/app/services/page_memory/_utils.py +++ b/apps/worker/app/services/page_memory/_utils.py @@ -75,11 +75,13 @@ def build_hierarchy_scopes( filename: str, page_count: int, ) -> list[CoarseScope]: - """Split skeleton list into coarse scopes for independent processing. + """Split skeleton list into per-leaf scopes for independent processing. - Groups skeletons by top-level section_path prefix. Skeletons not claimed - by any top-node are grouped by their root ancestor (first path segment - after filename) as orphan scopes. + Each TOC leaf skeleton becomes its own scope. ``skeletons`` is leaf-level + (produced by ``extract_section_skeletons`` → ``resolve_hierarchy_page_ranges``, + which only emits leaf nodes). When multiple leaves share the same page range + (e.g. unlocated leaves that fell back to the same parent scope), they are + merged into a single scope to keep ``scope_id`` unique. """ if not skeletons: return [] @@ -101,119 +103,24 @@ def build_hierarchy_scopes( ) ] - min_level = min(int(getattr(item, "level", 0) or 0) for item in skeletons) - top_nodes = [ - item - for item in skeletons - if int(getattr(item, "level", 0) or 0) == min_level - or not str(getattr(item, "parent_path", "") or "").startswith(f"{filename}/") - ] - seen_top_paths: set[str] = set() - unique_top_nodes: list[Any] = [] - for item in sort_skeletons(top_nodes): - if item.section_path in seen_top_paths: - continue - seen_top_paths.add(item.section_path) - unique_top_nodes.append(item) - - scopes: list[CoarseScope] = [] - for index, top in enumerate(unique_top_nodes): - prefix = f"{top.section_path}/" - members = [ - item - for item in skeletons - if item.section_path == top.section_path - or str(item.section_path).startswith(prefix) - ] - if not members: + scopes_by_range: dict[tuple[int, int], CoarseScope] = {} + for skel in sort_skeletons(skeletons): + start_page = max(1, int(getattr(skel, "start_page", 1) or 1)) + end_page = min(page_count, int(getattr(skel, "end_page", start_page) or start_page)) + if end_page < start_page: continue - start_page = max(1, int(getattr(top, "start_page", 1) or 1)) - next_top_start = ( - int(getattr(unique_top_nodes[index + 1], "start_page", page_count + 1) or page_count + 1) - if index + 1 < len(unique_top_nodes) - else page_count + 1 - ) - end_page = min( - page_count, - max( - start_page, - min( - int(getattr(top, "end_page", page_count) or page_count), - next_top_start - 1, - ), - ), - ) - bounded_members = [ - replace( - item, - start_page=max( - start_page, - int(getattr(item, "start_page", start_page) or start_page), - ), - end_page=min( - end_page, - int(getattr(item, "end_page", end_page) or end_page), - ), - ) - for item in members - if int(getattr(item, "start_page", 1) or 1) <= end_page - and int(getattr(item, "end_page", end_page) or end_page) >= start_page - ] - bounded_members = sort_skeletons(bounded_members) - scopes.append( - CoarseScope( + key = (start_page, end_page) + existing = scopes_by_range.get(key) + if existing is None: + scopes_by_range[key] = CoarseScope( scope_id=scope_id_for_pages(start_page, end_page), - skeletons=bounded_members, - strategy=f"coarse_scope_{index + 1}", + skeletons=[skel], + strategy="leaf_scope", start_page=start_page, end_page=end_page, ) - ) - - if scopes: - claimed_paths: set[str] = set() - for scope in scopes: - for item in scope.skeletons: - claimed_paths.add(item.section_path) - orphans = [item for item in skeletons if item.section_path not in claimed_paths] - if orphans: - root_groups: dict[str, list[Any]] = {} - for item in orphans: - parts = str(getattr(item, "section_path", "") or "").split("/") - root_key = "/".join(parts[:2]) if len(parts) >= 2 else item.section_path - root_groups.setdefault(root_key, []).append(item) - for root_key, group in sorted(root_groups.items()): - group = sort_skeletons(group) - sp = max(1, min(int(getattr(g, "start_page", 1) or 1) for g in group)) - ep = min(page_count, max(int(getattr(g, "end_page", page_count) or page_count) for g in group)) - scopes.append( - CoarseScope( - scope_id=scope_id_for_pages(sp, ep), - skeletons=group, - strategy="orphan_scope", - start_page=sp, - end_page=ep, - ) - ) - return scopes - - ordered = sort_skeletons(skeletons) - all_pages: set[int] = set() - for item in ordered: - sp = max(1, int(getattr(item, "start_page", 1) or 1)) - ep = min(page_count, int(getattr(item, "end_page", sp) or sp)) - if ep >= sp: - all_pages.update(range(sp, ep + 1)) - pages = sorted(all_pages) or list(range(1, page_count + 1)) - return [ - CoarseScope( - scope_id=scope_id_for_pages( - pages[0] if pages else 1, - pages[-1] if pages else page_count, - ), - skeletons=ordered, - strategy="full_coarse_hierarchy", - start_page=pages[0] if pages else 1, - end_page=pages[-1] if pages else page_count, - ) - ] + else: + scopes_by_range[key] = replace( + existing, skeletons=existing.skeletons + [skel] + ) + return list(scopes_by_range.values()) diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index d0778137..4b917cc2 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -11,7 +11,6 @@ from typing import Any from app.services.document_agent.manifest import ( - H1Candidate, PageAnatomyMap, ToolContext, ) @@ -28,7 +27,6 @@ ) from app.services.document_parser.structure.body_boundary import ( clean_toc_title, - normalize_heading_text, ) from loguru import logger @@ -130,16 +128,19 @@ def extract_section_skeletons( 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", - ) - ] + if not toc_nodes: + # TODO: explore lightweight hierarchy inference for no-TOC documents + # (e.g. heading font-size clustering, visual layout analysis). + # For now, no TOC → flat page tagging + asset extraction only. + return [ + _root_skeleton( + root_path=root_path, + filename=filename, + page_count=page_count, + reason="no_toc", + ) + ] + nodes = toc_nodes # Collapse degenerate single-child intermediate chains before locate. # Rule: only merge a parent with its only child when that child is NOT a @@ -465,18 +466,6 @@ def _toc_range_end(hierarchy: dict[str, Any]) -> int | None: 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) diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index 01f569eb..b45a26af 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -273,15 +273,10 @@ def fake_toc_extraction() -> None: {"toc_range": [17, 17], "toc_range_unit": "page", "toc_tree": {}} ] - def fake_h1_boundary() -> None: - calls.append("h1") - coordinator.blackboard.h1_result = H1BoundaryResult(method="toc_grep") - def fake_persist(_anatomy): calls.append("persist") monkeypatch.setattr(coordinator, "_run_toc_extraction_pipeline", fake_toc_extraction) - monkeypatch.setattr(coordinator, "_run_h1_boundary_pipeline", fake_h1_boundary) monkeypatch.setattr(coordinator, "_persist_ready_anatomy", fake_persist) monkeypatch.setattr( @@ -324,7 +319,7 @@ def run(self): anatomy = coordinator.run_structural() - assert calls[:2] == ["toc", "h1"] + assert calls == ["toc", "persist"] assert anatomy.toc_result.toc_pages == [17] @@ -369,15 +364,10 @@ def fake_toc_extraction() -> None: calls: list[str] = [] - def fake_h1_boundary() -> None: - calls.append("h1") - coordinator.blackboard.h1_result = H1BoundaryResult(method="none") - def fake_persist(_anatomy): calls.append("persist") monkeypatch.setattr(coordinator, "_run_toc_extraction_pipeline", fake_toc_extraction) - monkeypatch.setattr(coordinator, "_run_h1_boundary_pipeline", fake_h1_boundary) monkeypatch.setattr(coordinator, "_persist_ready_anatomy", fake_persist) monkeypatch.setattr( @@ -420,7 +410,7 @@ def run(self): anatomy = coordinator.run_structural() - assert calls == ["h1", "persist"] + assert calls == ["persist"] assert anatomy.toc_result.failure_kind == "rejected_all" assert anatomy.toc_result.toc_pages == [] @@ -454,15 +444,10 @@ def fake_toc_extraction() -> None: {"toc_range": [17, 17], "toc_range_unit": "page", "toc_tree": {}} ] - def fake_h1_boundary() -> None: - calls.append("h1") - coordinator.blackboard.h1_result = H1BoundaryResult(method="toc_grep") - def fake_persist(_anatomy): calls.append("persist") monkeypatch.setattr(coordinator, "_run_toc_extraction_pipeline", fake_toc_extraction) - monkeypatch.setattr(coordinator, "_run_h1_boundary_pipeline", fake_h1_boundary) monkeypatch.setattr(coordinator, "_persist_ready_anatomy", fake_persist) def fake_propose(_self): @@ -510,7 +495,7 @@ def run(self): coordinator.run_coarse() anatomy = coordinator.run_structural() - assert calls == ["toc", "planner", "h1", "persist"] + assert calls == ["toc", "planner", "persist"] assert anatomy.toc_result.toc_pages == [17] diff --git a/packages/shared-python/shared/models/schemas/page_memory_config.py b/packages/shared-python/shared/models/schemas/page_memory_config.py index 17944439..fd71f49b 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -11,7 +11,7 @@ class PageMemoryConfig: """Resolved page-memory defaults used by worker execution.""" max_pages: int = 1500 - scope_concurrency: int = 4 + scope_concurrency: int = 5 tag_concurrency: int = 4 tag_mode: Literal["vlm", "text"] = "vlm" fine_min_pages: int = 4 diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 2a5b55e3..08b33a2a 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -2,6 +2,7 @@ from loguru import logger +from shared.services.retrieval.agentic.discovery.tools import bottom_discovery from shared.services.retrieval.execution.reference_resolver import resolve_workflow_references from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results from shared.services.retrieval.execution.response_projection import ( @@ -12,6 +13,7 @@ RetrievalRouteContext, RetrievalRouteOutcome, ) +from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( count_scoped_chunks, load_all_scoped_chunks, @@ -25,7 +27,10 @@ async def run_retrieval_route( if small_corpus_outcome is not None: return small_corpus_outcome - return await _run_agentic_route(context) + if context.use_agentic is True: + return await _run_agentic_route(context) + + return await _run_classic_topk_route(context) async def _try_run_small_corpus_route( @@ -91,6 +96,65 @@ async def _try_run_small_corpus_route( ) +async def _run_classic_topk_route( + context: RetrievalRouteContext, +) -> RetrievalRouteOutcome: + discovery_result = await bottom_discovery( + context.db, + user_id=context.user_id, + namespace=context.namespace, + query=context.query, + top_k=context.effective_recall_k, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + chunk_types=context.allowed_chunk_types, + signal_paths=context.signal_paths, + filter_mode=context.filter_mode, + channels=context.channels, + channel_weights=context.channel_weights, + internal_recall_k=context.internal_recall_k, + ) + + fused_rows = ( + discovery_result.payload.get("fused_rows", []) + if discovery_result.status != "error" + else [] + ) + + ranked_rows = await rank_retrieval_candidates( + context.db, + user_id=context.user_id, + namespace=context.namespace, + discovery_rows=fused_rows, + routed_rows=[], + top_k=context.top_k, + ) + + assembled_rows = await assemble_retrieval_results( + db=context.db, + rows=ranked_rows, + exclude_document_ids=context.exclude_document_ids, + exclude_sections=context.exclude_sections, + allowed_chunk_types=context.allowed_chunk_types, + ) + results = [attach_citation(row) for row in assembled_rows] + response = { + "namespace": context.namespace, + "query": context.query, + "router_used": "classic_topk", + "evidence_text": render_legacy_evidence_text(results), + "answer_text": "", + "results": results, + } + return RetrievalRouteOutcome( + response=response, + hit_stats_results=results, + completion_label="CLASSIC TOP-K", + completion_count=len(results), + completion_detail="results", + ) + + async def _run_agentic_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome: