From ff4751c730f72aaeeabd8084340532f5f37b2adb Mon Sep 17 00:00:00 2001 From: Chengke <167045138+EricNGOntos@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:21:41 +0800 Subject: [PATCH 1/5] feat: remove H1 match pipeline, use TOC leaf-based shard planning (#203) * feat: implement node assembler for page memory hierarchy to support shared page content and VLM-based node summarization * feat: implement page asset integration into node assembly and chunk connections with support for custom probes * feat: enhance document profiling and page memory integration with skip shard plan option - Added `skip_shard_plan` parameter to `run_lightweight_anatomy` and `profile_document` to allow bypassing LLM shard decisions during profiling. - Updated `memory_service` to utilize `skip_shard_plan` for page memory processing. - Introduced asset annotation functionality for visualizing extracted assets on pages. - Refactored page section mapping and node assembly to support internal section body pages. - Improved environment configuration for Java dependencies in Docker setup. * feat: enhance retrieval and document ingestion with new chunk types and effective parse track handling - Updated `data_type` in `RetrievalQueryRequest` to support new chunk types: page (7) and text+image+table (8). - Refactored document ingestion service to apply effective parse track based on file extension and user-defined settings. - Introduced new validation for parse track handling in oversized PDF processing. - Enhanced tests to cover new chunk types and parse track logic, ensuring robust validation and functionality. * refactor: decouple page memory from image URIs and simplify node metadata structure * feat: integrate summary engine for enhanced document parsing and asset handling - Updated document parser to utilize the new summary engine for images, tables, and text, improving the extraction of titles, summaries, and entities. - Refactored image and table handling to streamline the summarization process, replacing legacy methods with unified calls to the summary engine. - Enhanced markdown parsing to support new entity and asset title fields, ensuring comprehensive data capture during document processing. - Introduced serialization for typed entities to maintain structured data integrity across parsed documents. * feat: implement concurrent page tagging and enrich image and table asset metadata with LLM-extracted entities and titles * feat: standardize page_memory track for PDF/PPTX, implement page-count limits, and refactor section summary publication to initialization. * chore: align parser env defaults * refactor: replace legacy data_type integer with flexible chunk_types set across retrieval services * feat: implement intermediate node chain collapse and add utilities for cross-page table continuity analysis * Add thread-safe trace logging, remove unused heading dataframes, and implement a text-only tagging mode for page memory processing. * refactor: remove budget tracking dependency from page memory services * chore: fix lint and type errors caught by ruff and pyright - Remove unused pandas import in docx/parser.py (F401) - Remove unused collapse_page_ranges import in fine_hierarchy.py (F401) - Rename unused ctx variable to _ctx in debug_page_memory.py (F841) - Remove unused tempfile import in test_page_memory_cross_page_table_merge.py (F401) - Fix greenlet value type: filter None from scope_results list in memory_service.py (reportAssignmentType) - Fix BeautifulSoup Tag narrowing in page_assets.py: use isinstance(x, Tag) guards for find_all/append calls (reportAttributeAccessIssue, reportOptionalMemberAccess) * fix: resolve alembic revision ID conflict causing cycle detection failure The feature branch introduced a migration with revision ID 'a1b2c3d4e5f6' (replace_data_type_with_chunk_types), which collided with an existing main branch migration sharing the same ID (add_qstash_tracking_columns). Alembic detected this as a cycle and blocked all migration tests. Fix: Assigned a unique revision ID 'f0d85d209e68' to the chunk_types migration and rebased its down_revision onto 'f9d0e1f2a3b4' (add_document_metadata_to_documents), which is the current main head. The Alembic graph now has a single clean head with no cycles. * test: fix failing tests from page-memory native hierarchy branch - test_agentic_evidence_renderer_contract.py: Update expected string format to match 'Pages 225-226' instead of 'Pages: 225, 226'. - test_page_memory_cross_page_table_merge.py: Remove unused 'budget' kwarg from merge_cross_page_tables() test calls. - test_doc_profile_anatomy_contract.py: Mock 'PDF_PROFILE_TOC_ENABLED' = True since the feature branch relies on it to test TOC profile attempting logic. * refactor: remove page_memory_work_summary_20260624.md file and streamline page memory processing documentation * feat: enhance skeleton extractor with native hierarchy support and agentic navigation improvements - Add utility functions for page memory processing in _utils.py - Refactor skeleton_extractor with enhanced hierarchy building and cross-page table merge - Improve agentic navigation with structured state management and tool enhancements - Add contract tests for agentic discovery selection - Integrate hierarchy locator improvements for document structure analysis Co-Authored-By: Claude Opus 4.6 * fix: remove dead commented code and propagate PageLocateConfig to pending TOC fallback - Remove old _estimate_page_offset code block (superseded by _calibrate_offset_via_vlm) - Pass page_memory_config through to _resolve_pending_tocs so its PageLocateResidualAgent fallback uses the same PageLocateConfig as the primary path Co-Authored-By: Claude Opus 4.6 * refactor: remove PageLocateResidualAgent fallback from skeleton extractor The offset-guided bulk anchoring with binary search handles all breakpoint scenarios. The expensive grep+VLM residual agent fallback is no longer needed. When offset anchoring is unavailable or returns empty, we fall through with calibration_overrides only and let the hierarchy resolver handle the rest via physical hints. Co-Authored-By: Claude Opus 4.6 * fix: resolve PR #200 lint and typecheck failures - Remove unused imports flagged by ruff F401 across page_memory and document_agent modules - Fix 6 pyright errors in skeleton_extractor.py: - min() on list[int|None]: use explicit loop with type-narrowed append - arithmetic on TitleNode.printed_page after filter: extract local int variable with explicit None guard Co-authored-by: Cursor * feat: add toc_page_offset to PageAnatomyMap and AgentBlackboard - Introduced toc_page_offset to both PageAnatomyMap and AgentBlackboard to facilitate better handling of table of contents page offsets. - Updated relevant functions to utilize the new toc_page_offset for improved shard planning and TOC processing. - Removed the match_h1_pages tool as it was deemed unnecessary for current functionality. Co-authored-by: Cursor * feat: prune out-of-scope TOC nodes before offset-guided anchoring When a PDF is truncated (physical pages < TOC page references), leaves with printed_page + offset > page_count are pruned from the tree before anchoring. This eliminates unnecessary bisect/recalibrate cycles and prevents unlocated fallback noise. Also adjusts _verify_offset_tail to prefer non-boundary leaves (printed_page + offset < page_count) for more reliable VLM verification. Co-Authored-By: Claude Opus 4.6 * feat: add chapter boundary detection and toc_chapter_boundary anchor type Add derive_chapter_boundaries() for L1/L2-aware shard planning with physical page ranges. Extend Shard.anchor_type to accept "toc_chapter_boundary" alongside existing types. Co-Authored-By: Claude Opus 4.6 * fix: use private _calibrate_offset_via_vlm import after merge rename Co-Authored-By: Claude Opus 4.6 * feat: restore use_agentic switch with classic top-K retrieval route Default retrieval now uses the classic 3-channel BM25 + RRF top-K path. Only explicit use_agentic=True triggers the full agentic workflow (LLM doc-select + navigation). Removes dead RETRIEVAL_AGENTIC_ENABLED env var reference from contract test. Co-Authored-By: Claude Opus 4.6 * fix: update contract test to expect classic_topk when use_agentic=False The test previously asserted that use_agentic=False was ignored and still routed to the workflow. Now that the switch is functional, the test verifies that use_agentic=False correctly routes to classic_topk. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: Cursor --- apps/api/app/api/v1/routes/retrieval.py | 2 +- .../contract/test_demo_documents_contract.py | 1 - .../tests/contract/test_retrieval_contract.py | 28 +- .../services/document_agent/coordinator.py | 7 - .../app/services/document_agent/manifest.py | 10 +- .../app/services/document_agent/registry.py | 4 - .../app/services/document_agent/state.py | 1 + .../services/document_agent/tools/__init__.py | 1 - .../document_agent/tools/match_h1_pages.py | 285 ------ .../tools/persist_anatomy_map.py | 2 +- .../tools/propose_shard_plan.py | 819 ++++++++++++------ .../tools/validate_anatomy_map.py | 4 +- .../app/services/document_agent/validators.py | 11 +- .../document_parser/formats/pdf/parser.py | 21 +- .../worker/app/services/page_memory/_utils.py | 137 +-- .../page_memory/skeleton_extractor.py | 37 +- .../test_doc_profile_anatomy_contract.py | 21 +- .../models/schemas/page_memory_config.py | 2 +- .../services/retrieval/execution/routes.py | 66 +- 19 files changed, 685 insertions(+), 774 deletions(-) delete mode 100644 apps/worker/app/services/document_agent/tools/match_h1_pages.py diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index ecafa4a42..9ed83c8c8 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 1c2b23a93..a1e704747 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 803240dcf..58de969c7 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 48c9653c0..4630ff43c 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 2518e2823..8861f631e 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 80a17531f..9fd4966ca 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 bcdb1a1c5..c717cbf07 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 f322f6d6a..ac5c29909 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 3e4d9ed03..000000000 --- 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 7ae7f0195..1e17a2202 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 7b50c98d0..1ca3fb805 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 c9e8273c4..2e75594c4 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 12ea930b9..fe26b4328 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 f1a829114..fb4b8e031 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 27d9d6c5b..219889bfc 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 d07781371..4b917cc27 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 01f569eb1..b45a26af1 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 179444392..fd71f49b0 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 2a5b55e3e..08b33a2a4 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: From 16392289dc02f8e57a97d2c7075babf14af41896 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 16:26:22 +0800 Subject: [PATCH 2/5] Add page-memory VLM capacity gate --- .../services/page_memory/memory_service.py | 19 +- .../services/page_memory/node_assembler.py | 134 +++++++---- .../app/services/page_memory/page_tagger.py | 70 ++++-- ...est_page_memory_node_assembler_contract.py | 143 +++++++++++ .../test_page_memory_page_tagger_contract.py | 153 ++++++++++++ .../shared-python/shared/core/config/ai.py | 20 ++ .../models/schemas/page_memory_config.py | 23 +- .../ai/openai_compatible_client_sync.py | 100 +++++--- .../services/ai/page_memory_vlm_limiter.py | 203 ++++++++++++++++ .../shared/services/ai/summary/engine.py | 7 + .../tests/test_page_memory_vlm_limiter.py | 225 ++++++++++++++++++ 11 files changed, 992 insertions(+), 105 deletions(-) create mode 100644 apps/worker/tests/contract/test_page_memory_page_tagger_contract.py create mode 100644 packages/shared-python/shared/services/ai/page_memory_vlm_limiter.py create mode 100644 packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index d1397bc38..f34cc42c0 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -445,6 +445,7 @@ def _build_page_dataframe( vlm_model=vlm_model, page_assets_by_page=page_assets_by_page, node_summary_max_pages=page_memory_config.node_summary_max_pages, + node_assembly_concurrency=page_memory_config.node_assembly_concurrency, ) logger.info( "[page_memory] C7 assembled {} node rows (verdict={})", @@ -625,14 +626,15 @@ def _run_hierarchy_scope( fat_leaf_pages = compute_fat_leaf_pages(scope_skeletons, min_pages=fine_min) if fat_leaf_pages: title_pages = sorted(fat_leaf_pages) - title_rendered = render_document_pages( - pdf_path=pdf_path, - page_count=page_count, - output_dir=output_dir, - pages=title_pages, - page_features=page_features, - page_texts=page_texts, - ) + with stage_timer("page_memory.title_render", page_count=len(title_pages)): + title_rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=output_dir, + pages=title_pages, + page_features=page_features, + page_texts=page_texts, + ) title_tags = [ PageTagResult( page_index=page, @@ -649,6 +651,7 @@ def _run_hierarchy_scope( fat_leaf_pages=fat_leaf_pages, budget=None, vlm_model=vlm_model, + max_concurrent=page_memory_config.title_detection_concurrency, ) _record_trace_stage( trace_recorder, diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index ff4191a25..fdc7bf2b1 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -25,10 +25,11 @@ import shutil from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, cast from loguru import logger +from app.services.document_parser.support.stage_profiler import stage_timer from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time from app.services.document_parser.support.parser_rows import serialize_entities from app.services.page_memory.page_assets import ( @@ -416,63 +417,114 @@ def build_node_rows( vlm_model: str | None = None, page_assets_by_page: dict[int, list[PageAsset]] | None = None, node_summary_max_pages: int = _NODE_SUMMARY_MAX_PAGES_DEFAULT, + node_assembly_concurrency: int = 2, ) -> list[dict[str, Any]]: """Assemble one row per leaf section node (node-granularity chunks).""" available_pages = set(raw_text_by_page.keys()) leaves = identify_leaf_nodes(skeletons) views, page_owner = assign_pages_to_leaves(leaves, available_pages=available_pages) page_to_leaves = pages_by_leaf_count(views) + resolved_concurrency = max(1, node_assembly_concurrency) # Resolve body text once per owned page (PyMuPDF, OCR fallback for scanned). resolved_text: dict[int, str] = {} - for view in views: - for page in view.owned_pages: - resolved_text[page] = resolve_page_text( + owned_pages = sorted({page for view in views for page in view.owned_pages}) + if owned_pages: + import gevent + from gevent.pool import Pool as GeventPool + + def _resolve_one(page: int) -> tuple[int, str]: + return page, resolve_page_text( page=page, raw_text=raw_text_by_page.get(page, ""), image_path=image_path_by_page.get(page), vlm_model=vlm_model, ) + with stage_timer( + "page_memory.node_ocr", + page_count=len(owned_pages), + concurrency=resolved_concurrency, + ): + pool = GeventPool(size=min(resolved_concurrency, len(owned_pages))) + greenlets = [pool.spawn(_resolve_one, page) for page in owned_pages] + gevent.joinall(greenlets, raise_error=True) + resolved_pairs = [ + cast(tuple[int, str], greenlet.value) + for greenlet in greenlets + ] + resolved_text = {page: text for page, text in resolved_pairs} + + summaries: dict[int, tuple[str, list[str], list[dict[str, str]]]] = {} + if views: + import gevent + from gevent.pool import Pool as GeventPool + + def _summarize_one( + index: int, + ) -> tuple[int, tuple[str, list[str], list[dict[str, str]]]]: + return index, compute_node_summary( + view=views[index], + page_to_leaves=page_to_leaves, + tag_by_page=tag_by_page, + image_path_by_page=image_path_by_page, + vlm_model=vlm_model, + node_summary_max_pages=node_summary_max_pages, + ) + + with stage_timer( + "page_memory.node_summary", + node_count=len(views), + concurrency=resolved_concurrency, + ): + pool = GeventPool(size=min(resolved_concurrency, len(views))) + greenlets = [ + pool.spawn(_summarize_one, index) + for index in range(len(views)) + ] + gevent.joinall(greenlets, raise_error=True) + summary_pairs = [ + cast( + tuple[int, tuple[str, list[str], list[dict[str, str]]]], + greenlet.value, + ) + for greenlet in greenlets + ] + summaries = {index: result for index, result in summary_pairs} + rows: list[dict[str, Any]] = [] rows_by_path: dict[str, dict[str, Any]] = {} - for view in views: - leaf = view.leaf - content = build_node_content( - view, - page_owner=page_owner, - page_text=resolved_text, - ) - summary, keywords, entities = compute_node_summary( - view=view, - page_to_leaves=page_to_leaves, - tag_by_page=tag_by_page, - image_path_by_page=image_path_by_page, - vlm_model=vlm_model, - node_summary_max_pages=node_summary_max_pages, - ) - know_id = f"node_{gen_str_codes(f'{filename}::{leaf.section_path}')}" - row = { - "content": content, - "path": leaf.section_path, - "type": "page", - "length": len(content), - "keywords": ";".join(keywords), - "summary": summary, - "know_id": know_id, - "tokens": "", - "connectto": "", - "addtime": get_str_time(), - "page_nums": ",".join(str(page) for page in view.pages), - "entities": serialize_entities(entities), - "asset_title": "", - "extra_metadata": _build_page_extra_metadata( - pages=view.pages, - image_path_by_page=image_path_by_page, - ), - } - rows.append(row) - rows_by_path[leaf.section_path] = row + with stage_timer("page_memory.node_rows", node_count=len(views)): + for index, view in enumerate(views): + leaf = view.leaf + content = build_node_content( + view, + page_owner=page_owner, + page_text=resolved_text, + ) + summary, keywords, entities = summaries.get(index, ("", [], [])) + know_id = f"node_{gen_str_codes(f'{filename}::{leaf.section_path}')}" + row = { + "content": content, + "path": leaf.section_path, + "type": "page", + "length": len(content), + "keywords": ";".join(keywords), + "summary": summary, + "know_id": know_id, + "tokens": "", + "connectto": "", + "addtime": get_str_time(), + "page_nums": ",".join(str(page) for page in view.pages), + "entities": serialize_entities(entities), + "asset_title": "", + "extra_metadata": _build_page_extra_metadata( + pages=view.pages, + image_path_by_page=image_path_by_page, + ), + } + rows.append(row) + rows_by_path[leaf.section_path] = row asset_rows: list[dict[str, Any]] = [] if page_assets_by_page: diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index fcee9da5f..aa5be87d9 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -26,6 +26,7 @@ from app.services.page_memory.page_renderer import PageRenderResult from shared.services.ai.prompt_service import build_prompt from shared.services.ai.summary.engine import summarize +from shared.core.exceptions.domain_exceptions import UnavailableException @dataclass @@ -90,6 +91,8 @@ def tag_pages( resolved_max_concurrent = max_concurrent or int( getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 4) ) + if not pages: + return [] def _tag_one(page: PageRenderResult) -> PageTagResult: plan = plan_map.get(page.page_index) @@ -112,9 +115,9 @@ def _tag_one(page: PageRenderResult) -> PageTagResult: pool = GeventPool(size=min(resolved_max_concurrent, len(pages))) greenlets = [pool.spawn(_tag_one, page) for page in pages] - gevent.joinall(greenlets) + gevent.joinall(greenlets, raise_error=True) - results = [g.value for g in greenlets if g.value is not None] + results = [cast(PageTagResult, g.value) for g in greenlets] vlm_calls = sum(1 for r in results if r.strategy_used == "vlm_lite") logger.info( "[page_tagger] tagged {} pages ({} VLM calls, {} text_only, {} skipped, {} failed) concurrency={}", @@ -182,6 +185,8 @@ def _tag_text_only( response_format={"type": "json_object"}, usage_task="page_memory.text_tag", ) + except UnavailableException: + raise except Exception as exc: logger.warning( "[page_tagger] text_only LLM failed for page {}: {}", @@ -262,6 +267,7 @@ def tag_page_titles( fat_leaf_pages: set[int], budget: Any | None = None, vlm_model: str | None = None, + max_concurrent: int | None = None, ) -> list[PageTagResult]: """Run independent VLM title detection on fat-leaf pages. @@ -278,6 +284,8 @@ def tag_page_titles( Deprecated, ignored. Kept for call-site compatibility. vlm_model: VLM model name; falls back to ``$IMAGE_MODEL``. + max_concurrent: + Maximum concurrent title-detection calls. Returns ------- @@ -294,29 +302,57 @@ def tag_page_titles( tag_map = {t.page_index: t for t in tag_results} page_map = {p.page_index: p for p in pages} - vlm_calls = 0 - titles_found = 0 + work_items = [ + (page_idx, page, tag_map[page_idx]) + for page_idx in sorted(fat_leaf_pages) + if (page := page_map.get(page_idx)) is not None + and page_idx in tag_map + and page.image_path + and os.path.exists(page.image_path) + ] + if not work_items: + return tag_results - for page_idx in sorted(fat_leaf_pages): - page = page_map.get(page_idx) - tag = tag_map.get(page_idx) - if page is None or tag is None: - continue + from shared.core.config import settings + + resolved_max_concurrent = max_concurrent or int( + getattr(settings, "PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY", 2) + ) - # Skip pages without images (text_only / skip) - if not page.image_path or not os.path.exists(page.image_path): - continue + def _detect_one( + page_idx: int, + page: PageRenderResult, + ) -> tuple[int, list[dict[str, Any]]]: + return page_idx, _tag_vlm_titles(page, model=model) - observed = _tag_vlm_titles(page, model=model) + import gevent + from gevent.pool import Pool as GeventPool + + pool = GeventPool(size=min(resolved_max_concurrent, len(work_items))) + greenlets = [ + pool.spawn(_detect_one, page_idx, page) + for page_idx, page, _tag in work_items + ] + gevent.joinall(greenlets, raise_error=True) + + title_pairs = [ + cast(tuple[int, list[dict[str, Any]]], greenlet.value) + for greenlet in greenlets + ] + titles_by_page: dict[int, list[dict[str, Any]]] = dict(title_pairs) + for page_idx, _page, tag in work_items: + observed = titles_by_page.get(page_idx, []) tag.observed_titles = observed - vlm_calls += 1 - titles_found += len(observed) + + vlm_calls = len(work_items) + titles_found = sum(len(titles) for titles in titles_by_page.values()) logger.info( - "[page_tagger] title detection: {} VLM calls on {} fat-leaf pages, {} titles found", + "[page_tagger] title detection: {} VLM calls on {} fat-leaf pages, {} titles found concurrency={}", vlm_calls, len(fat_leaf_pages), titles_found, + resolved_max_concurrent, ) return tag_results @@ -413,6 +449,8 @@ def _tag_vlm_titles( page.page_index, ) return [] + except UnavailableException: + raise except Exception as exc: logger.warning( "[page_tagger] title VLM failed for page {}: {}", diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index c66fc217b..a502ae3be 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -9,6 +9,9 @@ os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") os.environ.setdefault("S3_TEMP_PATH", "/tmp") +import pytest + +import app.services.page_memory.node_assembler as node_assembler from app.services.page_memory.node_assembler import ( SAME_AS_PREFIX, assign_pages_to_leaves, @@ -21,6 +24,7 @@ from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks +from shared.core.exceptions.domain_exceptions import UnavailableException import pandas as pd from PIL import Image @@ -54,6 +58,20 @@ def _same_page_sibling_skeletons() -> list[SectionSkeleton]: return [parent, child_a, child_b] +def _ordered_page_skeletons() -> list[SectionSkeleton]: + return [ + SectionSkeleton( + section_path=f"demo.pdf/{page_index}", + level=1, + start_page=page_index, + end_page=page_index, + title=f"Section {page_index}", + parent_path="demo.pdf", + ) + for page_index in [1, 2, 3] + ] + + def test_identify_leaf_nodes_drops_internal_parents() -> None: leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) assert [leaf.title for leaf in leaves] == ["3.1 职责", "3.2 管理规定"] @@ -125,6 +143,131 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: assert leaf_b["extra_metadata"] == {} +def test_build_node_rows_preserves_order_under_ocr_and_summary_concurrency( + monkeypatch, +) -> None: + import gevent + + def _fake_resolve_page_text(**kwargs) -> str: + page = int(kwargs["page"]) + gevent.sleep(0.01 * (4 - page)) + return f"text-{page}" + + def _fake_compute_node_summary(**kwargs): + view = kwargs["view"] + gevent.sleep(0.01 * view.leaf.start_page) + return f"summary-{view.leaf.start_page}", [f"k{view.leaf.start_page}"], [] + + monkeypatch.setattr(node_assembler, "resolve_page_text", _fake_resolve_page_text) + monkeypatch.setattr( + node_assembler, + "compute_node_summary", + _fake_compute_node_summary, + ) + + rows = build_node_rows( + skeletons=_ordered_page_skeletons(), + raw_text_by_page={1: "", 2: "", 3: ""}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={}, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + node_assembly_concurrency=2, + ) + + assert [row["path"] for row in rows] == [ + "demo.pdf/1", + "demo.pdf/2", + "demo.pdf/3", + ] + assert [row["content"] for row in rows] == ["text-1", "text-2", "text-3"] + assert [row["summary"] for row in rows] == [ + "summary-1", + "summary-2", + "summary-3", + ] + + +def test_build_node_rows_failed_ocr_greenlet_fails_stage(monkeypatch) -> None: + def _fake_resolve_page_text(**kwargs) -> str: + if int(kwargs["page"]) == 2: + raise RuntimeError("ocr failed") + return "ok" + + monkeypatch.setattr(node_assembler, "resolve_page_text", _fake_resolve_page_text) + + with pytest.raises(RuntimeError): + build_node_rows( + skeletons=_ordered_page_skeletons()[:2], + raw_text_by_page={1: "", 2: ""}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={}, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + node_assembly_concurrency=2, + ) + + +def test_build_node_rows_unavailable_propagates_from_ocr(monkeypatch) -> None: + def _fake_resolve_page_text(**kwargs) -> str: + raise UnavailableException( + internal_message="ocr capacity busy", + retry_after=5, + ) + + monkeypatch.setattr(node_assembler, "resolve_page_text", _fake_resolve_page_text) + + with pytest.raises(UnavailableException): + build_node_rows( + skeletons=_ordered_page_skeletons()[:1], + raw_text_by_page={1: ""}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={}, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + node_assembly_concurrency=1, + ) + + +def test_build_node_rows_unavailable_propagates_from_node_summary( + monkeypatch, +) -> None: + def _fake_compute_node_summary(**kwargs): + raise UnavailableException( + internal_message="summary capacity busy", + retry_after=5, + ) + + monkeypatch.setattr( + node_assembler, + "compute_node_summary", + _fake_compute_node_summary, + ) + + with pytest.raises(UnavailableException): + build_node_rows( + skeletons=_ordered_page_skeletons()[:1], + raw_text_by_page={1: "text-1"}, + image_path_by_page={}, + kind_by_page={}, + tag_by_page={}, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model="fake-vlm", + node_assembly_concurrency=1, + ) + + def test_build_node_rows_attaches_page_citation_assets_for_rendered_pages(tmp_path) -> None: page_image = tmp_path / "pages" / "page-231.png" page_image.parent.mkdir() diff --git a/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py b/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py new file mode 100644 index 000000000..8eeba6360 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_page_tagger_contract.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import os + +import pytest + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.page_memory.page_renderer import PageRenderResult +from app.services.page_memory.page_tagger import PageTagResult, tag_page_titles +from shared.core.exceptions.domain_exceptions import UnavailableException + + +def _write_page_image(tmp_path, page_index: int) -> str: + image_path = tmp_path / f"page-{page_index}.png" + image_path.write_bytes(b"png") + return str(image_path) + + +def test_title_detection_preserves_page_index_assignment_under_concurrency( + monkeypatch, + tmp_path, +) -> None: + import gevent + + def _fake_tag_vlm_titles( + page: PageRenderResult, + *, + model: str, + ) -> list[dict[str, object]]: + gevent.sleep(0.01 * (4 - page.page_index)) + return [{"text": f"title-{page.page_index}", "prominence": 0.8}] + + monkeypatch.setitem( + tag_page_titles.__globals__, + "_tag_vlm_titles", + _fake_tag_vlm_titles, + ) + pages = [ + PageRenderResult( + page_index=page_index, + image_path=_write_page_image(tmp_path, page_index), + raw_text="", + width=100, + height=200, + is_landscape=False, + ) + for page_index in [1, 2, 3] + ] + tag_results = [ + PageTagResult(page_index=3), + PageTagResult(page_index=1), + PageTagResult(page_index=2), + ] + + results = tag_page_titles( + pages=pages, + tag_results=tag_results, + fat_leaf_pages={1, 2, 3}, + vlm_model="fake-vlm", + max_concurrent=2, + ) + + observed_by_page = { + tag.page_index: tag.observed_titles[0]["text"] + for tag in results + } + assert observed_by_page == { + 1: "title-1", + 2: "title-2", + 3: "title-3", + } + + +def test_title_detection_failed_greenlet_fails_stage(monkeypatch, tmp_path) -> None: + def _fake_tag_vlm_titles( + page: PageRenderResult, + *, + model: str, + ) -> list[dict[str, object]]: + if page.page_index == 2: + raise RuntimeError("title detection failed") + return [{"text": f"title-{page.page_index}"}] + + monkeypatch.setitem( + tag_page_titles.__globals__, + "_tag_vlm_titles", + _fake_tag_vlm_titles, + ) + pages = [ + PageRenderResult( + page_index=page_index, + image_path=_write_page_image(tmp_path, page_index), + raw_text="", + width=100, + height=200, + is_landscape=False, + ) + for page_index in [1, 2] + ] + tag_results = [PageTagResult(page_index=1), PageTagResult(page_index=2)] + + with pytest.raises(RuntimeError): + tag_page_titles( + pages=pages, + tag_results=tag_results, + fat_leaf_pages={1, 2}, + vlm_model="fake-vlm", + max_concurrent=2, + ) + + +def test_title_detection_unavailable_exception_propagates( + monkeypatch, + tmp_path, +) -> None: + def _fake_tag_vlm_titles( + page: PageRenderResult, + *, + model: str, + ) -> list[dict[str, object]]: + raise UnavailableException( + internal_message="capacity busy", + retry_after=5, + ) + + monkeypatch.setitem( + tag_page_titles.__globals__, + "_tag_vlm_titles", + _fake_tag_vlm_titles, + ) + page = PageRenderResult( + page_index=1, + image_path=_write_page_image(tmp_path, 1), + raw_text="", + width=100, + height=200, + is_landscape=False, + ) + + with pytest.raises(UnavailableException): + tag_page_titles( + pages=[page], + tag_results=[PageTagResult(page_index=1)], + fat_leaf_pages={1}, + vlm_model="fake-vlm", + max_concurrent=1, + ) diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index 173992f5f..d58315e34 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -96,6 +96,26 @@ class AIConfig(BaseModel): default=4, description="Max concurrent local DOCX image summary VLM calls per parse job.", ) + PAGE_MEMORY_VLM_MAX_INFLIGHT: int = Field( + default=16, + description="Global Redis-backed in-flight limit for page-memory VLM calls.", + ) + PAGE_MEMORY_VLM_LEASE_TTL_SECONDS: int = Field( + default=600, + description="Safety TTL for page-memory VLM in-flight leases.", + ) + PAGE_MEMORY_VLM_WAIT_TIMEOUT_SECONDS: int = Field( + default=120, + description="Max wait before page-memory VLM capacity pressure becomes retryable.", + ) + PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY: int = Field( + default=2, + description="Local per-job title detection concurrency for page-memory.", + ) + PAGE_MEMORY_NODE_ASSEMBLY_CONCURRENCY: int = Field( + default=2, + description="Local per-job node OCR and summary concurrency for page-memory.", + ) TOKEN_PRICING_TABLE_JSON: str = Field( default="", description=( 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 fd71f49b0..b895f1476 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -13,6 +13,8 @@ class PageMemoryConfig: max_pages: int = 1500 scope_concurrency: int = 5 tag_concurrency: int = 4 + title_detection_concurrency: int = 2 + node_assembly_concurrency: int = 2 tag_mode: Literal["vlm", "text"] = "vlm" fine_min_pages: int = 4 hierarchy_model: str | None = None @@ -35,7 +37,18 @@ class PageMemoryConfig: @classmethod def default(cls) -> Self: - return cls() + from shared.core.config import settings + + return cls( + title_detection_concurrency=_as_int( + getattr(settings, "PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY", 2), + 2, + ), + node_assembly_concurrency=_as_int( + getattr(settings, "PAGE_MEMORY_NODE_ASSEMBLY_CONCURRENCY", 2), + 2, + ), + ) @classmethod def from_mapping(cls, value: object) -> Self: @@ -57,6 +70,14 @@ def from_mapping(cls, value: object) -> Self: value.get("tag_concurrency"), default.tag_concurrency, ), + title_detection_concurrency=_as_int( + value.get("title_detection_concurrency"), + default.title_detection_concurrency, + ), + node_assembly_concurrency=_as_int( + value.get("node_assembly_concurrency"), + default.node_assembly_concurrency, + ), tag_mode=resolved_tag_mode, fine_min_pages=_as_int( value.get("fine_min_pages"), diff --git a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py index f3bf53ebb..6a1cf590a 100644 --- a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py +++ b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py @@ -16,11 +16,15 @@ from openai.types.chat import ChatCompletionMessageParam from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import LLMServiceException +from shared.core.exceptions.domain_exceptions import LLMServiceException, UnavailableException from shared.services.http.client_pool import get_sync_client from shared.services.ai.llm_mock import build_mock_chat_completion_response from shared.utils.security_utils import mask_api_key from shared.services.ai.token_tracking import record_tokens +from shared.services.ai.page_memory_vlm_limiter import ( + PageMemoryVlmLease, + get_page_memory_vlm_limiter, +) LOCAL_DEBUG = os.getenv("LOCAL_DEBUG", "0") == "1" LLMUsage = dict[str, int] @@ -39,6 +43,10 @@ def _should_mock_llm_calls() -> bool: return bool(getattr(settings, "LLM_MOCK_ENABLED", False)) +def _is_page_memory_usage_task(usage_task: str | None) -> bool: + return bool(usage_task and usage_task.startswith("page_memory.")) + + def _empty_usage() -> LLMUsage: return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} @@ -171,6 +179,19 @@ def _resolve_direct_api_key( # Ali token-pool helpers # ------------------------------------------------------------------ + def _acquire_page_memory_vlm_lease( + self, *, usage_task: str | None + ) -> PageMemoryVlmLease | None: + if not _is_page_memory_usage_task(usage_task): + return None + return get_page_memory_vlm_limiter().acquire(usage_task=usage_task or "") + + def _release_page_memory_vlm_lease( + self, lease: PageMemoryVlmLease | None + ) -> None: + if lease is not None: + get_page_memory_vlm_limiter().release(lease) + def _should_use_ali_pool(self) -> bool: """Whether to route through the AliQuotaManager instead of a fixed key.""" if self._explicit_api_key: @@ -327,23 +348,24 @@ def chat_completion_raw_with_usage( ) return {"mock_content": content}, _empty_usage() - if self._should_use_ali_pool(): - return self._make_ali_pool_raw_call( - model=effective_model, - all_messages=all_messages, - temperature=temperature, - max_tokens=max_tokens, - api_kwargs=api_kwargs, - usage_task=usage_task, - ) - - client = self._client - if client is None: - raise LLMServiceException( - internal_message="OpenAI client is not initialized for direct provider requests", - provider=self.default_model, - ) + lease = self._acquire_page_memory_vlm_lease(usage_task=usage_task) try: + if self._should_use_ali_pool(): + return self._make_ali_pool_raw_call( + model=effective_model, + all_messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + api_kwargs=api_kwargs, + usage_task=usage_task, + ) + + client = self._client + if client is None: + raise LLMServiceException( + internal_message="OpenAI client is not initialized for direct provider requests", + provider=self.default_model, + ) response = client.chat.completions.create( model=effective_model, messages=all_messages, @@ -361,11 +383,14 @@ def chat_completion_raw_with_usage( return response, usage except LLMServiceException: raise + except UnavailableException: + raise except Exception as exc: + base_url = getattr(self._client, "base_url", None) logger.error( "LLM raw request failed: model={model}, base_url={base_url}, error_chain={error_chain}", model=effective_model, - base_url=client.base_url, + base_url=base_url, error_chain=_summarize_exception_chain(exc), ) raise LLMServiceException( @@ -373,6 +398,8 @@ def chat_completion_raw_with_usage( provider=self.default_model, original_exception=exc, ) from exc + finally: + self._release_page_memory_vlm_lease(lease) def chat_completion_with_usage( self, @@ -428,9 +455,10 @@ def chat_completion_with_usage( model_name=effective_model, ), _empty_usage() - # Route through Ali token pool when applicable - if self._should_use_ali_pool(): - try: + lease = self._acquire_page_memory_vlm_lease(usage_task=usage_task) + try: + # Route through Ali token pool when applicable + if self._should_use_ali_pool(): return self._make_ali_pool_call( model=effective_model, all_messages=all_messages, @@ -439,25 +467,14 @@ def chat_completion_with_usage( api_kwargs=api_kwargs, usage_task=usage_task, ) - except LLMServiceException: - raise - except Exception as exc: - logger.error(f"LLM request failed (Ali pool): model={effective_model}, error={exc}") + + # Non-Ali path: use the single pre-configured client + client = self._client + if client is None: raise LLMServiceException( - internal_message=f"API request failed: {str(exc)}", + internal_message="OpenAI client is not initialized for direct provider requests", provider=self.default_model, - original_exception=exc, - ) from exc - - # Non-Ali path: use the single pre-configured client - client = self._client - if client is None: - raise LLMServiceException( - internal_message="OpenAI client is not initialized for direct provider requests", - provider=self.default_model, - ) - - try: + ) response = client.chat.completions.create( model=effective_model, messages=all_messages, @@ -479,11 +496,14 @@ def chat_completion_with_usage( return content, usage except LLMServiceException: raise + except UnavailableException: + raise except Exception as exc: + base_url = getattr(self._client, "base_url", None) logger.error( "LLM request failed: model={model}, base_url={base_url}, error_chain={error_chain}", model=effective_model, - base_url=client.base_url, + base_url=base_url, error_chain=_summarize_exception_chain(exc), ) raise LLMServiceException( @@ -491,6 +511,8 @@ def chat_completion_with_usage( provider=self.default_model, original_exception=exc, ) from exc + finally: + self._release_page_memory_vlm_lease(lease) def chat_completion( self, diff --git a/packages/shared-python/shared/services/ai/page_memory_vlm_limiter.py b/packages/shared-python/shared/services/ai/page_memory_vlm_limiter.py new file mode 100644 index 000000000..0c04f6524 --- /dev/null +++ b/packages/shared-python/shared/services/ai/page_memory_vlm_limiter.py @@ -0,0 +1,203 @@ +"""Redis-backed global in-flight gate for page-memory VLM calls.""" + +from __future__ import annotations + +import random +import time +from dataclasses import dataclass +from typing import Optional + +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import UnavailableException +from shared.services.redis.redis_sync_service import ( + SyncRedisService, + SyncRedisServiceFactory, +) + + +@dataclass(frozen=True) +class PageMemoryVlmLease: + """A reserved page-memory VLM capacity slot.""" + + usage_task: str + acquired_at: float + current_inflight: int + + +class PageMemoryVlmLimiter: + """Global Redis counter for page-memory VLM in-flight capacity.""" + + INFLIGHT_KEY = "page_memory:vlm:inflight" + user_message = "Document parsing VLM capacity is busy. Please retry shortly." + + _ACQUIRE_SCRIPT: str = """ +local key = KEYS[1] +local max_inflight = tonumber(ARGV[1]) +local ttl_seconds = tonumber(ARGV[2]) + +local current = tonumber(redis.call('GET', key) or '0') +if current >= max_inflight then + redis.call('EXPIRE', key, ttl_seconds) + return {0, current} +end + +local new_count = redis.call('INCR', key) +redis.call('EXPIRE', key, ttl_seconds) +return {1, new_count} +""" + + _RELEASE_SCRIPT: str = """ +local key = KEYS[1] +local current = tonumber(redis.call('GET', key) or '0') +if current <= 1 then + redis.call('DEL', key) + return 0 +end +return redis.call('DECR', key) +""" + + def __init__( + self, + redis_service: SyncRedisService, + *, + max_inflight: int, + lease_ttl_seconds: int, + wait_timeout_seconds: int, + ) -> None: + self.redis = redis_service + self.max_inflight = max(1, max_inflight) + self.lease_ttl_seconds = max(1, lease_ttl_seconds) + self.wait_timeout_seconds = max(1, wait_timeout_seconds) + + @classmethod + def from_settings( + cls, + redis_service: Optional[SyncRedisService] = None, + ) -> "PageMemoryVlmLimiter": + """Create a limiter from internal page-memory settings.""" + return cls( + redis_service or SyncRedisServiceFactory.get_service(), + max_inflight=int(getattr(settings, "PAGE_MEMORY_VLM_MAX_INFLIGHT", 16)), + lease_ttl_seconds=int( + getattr(settings, "PAGE_MEMORY_VLM_LEASE_TTL_SECONDS", 600) + ), + wait_timeout_seconds=int( + getattr(settings, "PAGE_MEMORY_VLM_WAIT_TIMEOUT_SECONDS", 120) + ), + ) + + def acquire(self, *, usage_task: str) -> PageMemoryVlmLease: + """Wait for and reserve one global VLM capacity slot.""" + started_at = time.monotonic() + deadline = started_at + self.wait_timeout_seconds + last_inflight = 0 + + while time.monotonic() < deadline: + acquired, current_inflight = self._try_acquire(usage_task=usage_task) + last_inflight = current_inflight + if acquired: + wait_ms = int((time.monotonic() - started_at) * 1000) + logger.bind( + event="page_memory.vlm_gate.acquired", + usage_task=usage_task, + current_inflight=current_inflight, + max_inflight=self.max_inflight, + wait_ms=wait_ms, + ).info("page_memory.vlm_gate.acquired") + return PageMemoryVlmLease( + usage_task=usage_task, + acquired_at=time.monotonic(), + current_inflight=current_inflight, + ) + + remaining = deadline - time.monotonic() + if remaining <= 0: + break + self._sleep(min(self._jittered_backoff(last_inflight), remaining)) + + raise UnavailableException( + internal_message=( + "Page-memory VLM capacity gate timed out " + f"after {self.wait_timeout_seconds}s " + f"({last_inflight}/{self.max_inflight} in-flight)" + ), + retry_after=max(1, min(self.wait_timeout_seconds, 60)), + limit=self.max_inflight, + period="minute", + user_message=self.user_message, + ) + + def release(self, lease: PageMemoryVlmLease) -> None: + """Release a previously acquired page-memory VLM capacity slot.""" + try: + self.redis.eval( + self._RELEASE_SCRIPT, + keys=[self.INFLIGHT_KEY], + args=[], + ) + except Exception: + logger.opt(exception=True).warning( + "Failed to release page-memory VLM in-flight slot for {}", + lease.usage_task, + ) + + def get_inflight_count(self) -> int: + """Return the current Redis in-flight count.""" + try: + raw = self.redis.get(self.INFLIGHT_KEY, 0) + return max(0, int(raw or 0)) + except Exception: + return 0 + + def _try_acquire(self, *, usage_task: str) -> tuple[bool, int]: + try: + result = self.redis.eval( + self._ACQUIRE_SCRIPT, + keys=[self.INFLIGHT_KEY], + args=[self.max_inflight, self.lease_ttl_seconds], + ) + except Exception as exc: + logger.opt(exception=True).warning( + "Failed to check page-memory VLM in-flight counter for {}", + usage_task, + ) + raise UnavailableException( + internal_message="Failed to check page-memory VLM capacity gate", + retry_after=5, + user_message=self.user_message, + original_exception=exc, + ) from exc + + if not isinstance(result, list) or len(result) < 2: + raise UnavailableException( + internal_message=f"Unexpected page-memory VLM gate result: {result!r}", + retry_after=5, + user_message=self.user_message, + ) + return int(result[0]) == 1, max(0, int(result[1])) + + def _jittered_backoff(self, current_inflight: int) -> float: + pressure = max(current_inflight - self.max_inflight + 1, 1) + base = min(0.25 * pressure, 2.0) + return base + random.uniform(0.0, 0.2) + + def _sleep(self, seconds: float) -> None: + try: + import gevent + + gevent.sleep(seconds) + except ImportError: + time.sleep(seconds) + + +_page_memory_vlm_limiter: Optional[PageMemoryVlmLimiter] = None + + +def get_page_memory_vlm_limiter() -> PageMemoryVlmLimiter: + """Return a singleton page-memory VLM limiter for worker processes.""" + global _page_memory_vlm_limiter + if _page_memory_vlm_limiter is None: + _page_memory_vlm_limiter = PageMemoryVlmLimiter.from_settings() + return _page_memory_vlm_limiter diff --git a/packages/shared-python/shared/services/ai/summary/engine.py b/packages/shared-python/shared/services/ai/summary/engine.py index 7a71d1f1b..eaedd05c2 100644 --- a/packages/shared-python/shared/services/ai/summary/engine.py +++ b/packages/shared-python/shared/services/ai/summary/engine.py @@ -29,6 +29,7 @@ from loguru import logger +from shared.core.exceptions.domain_exceptions import UnavailableException from shared.services.ai import openai_compatible_client_sync as _client_mod from shared.services.ai.prompt_service import _detect_text_language, build_prompt from shared.services.ai.response_process_service import eval_response @@ -177,6 +178,12 @@ def _call_llm( ) continue return None + except UnavailableException: + if budget is not None: + budget.refund(budget_pool, est=est, stage=budget_stage) + if usage_task.startswith("page_memory."): + raise + return None except Exception as exc: logger.warning("[summary] LLM call failed for {}: {}", usage_task, exc) if budget is not None: diff --git a/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py b/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py new file mode 100644 index 000000000..49c82acd2 --- /dev/null +++ b/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py @@ -0,0 +1,225 @@ +from __future__ import annotations + +import os +from typing import Any + +import pytest + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.core.exceptions.domain_exceptions import ( # noqa: E402 + LLMServiceException, + UnavailableException, +) +from shared.services.ai.openai_compatible_client_sync import ( # noqa: E402 + OpenAICompatibleClientSync, +) +from shared.services.ai.page_memory_vlm_limiter import ( # noqa: E402 + PageMemoryVlmLease, + PageMemoryVlmLimiter, +) +from shared.services.ai.summary.engine import summarize, transcribe # noqa: E402 + + +class _FakeRedis: + def __init__(self, *, initial_count: int = 0, fail_eval: bool = False) -> None: + self.count = initial_count + self.fail_eval = fail_eval + + def eval(self, script: str, keys: list[str], args: list[Any]) -> list[int] | int: + if self.fail_eval: + raise RuntimeError("redis unavailable") + if "INCR" in script: + max_inflight = int(args[0]) + if self.count >= max_inflight: + return [0, self.count] + self.count += 1 + return [1, self.count] + self.count = max(self.count - 1, 0) + return self.count + + def get(self, key: str, default: Any = None) -> int: + return self.count + + +class _DelayedCapacityRedis(_FakeRedis): + def __init__(self) -> None: + super().__init__(initial_count=1) + self.acquire_attempts = 0 + + def eval(self, script: str, keys: list[str], args: list[Any]) -> list[int] | int: + if "INCR" not in script: + return super().eval(script, keys, args) + self.acquire_attempts += 1 + if self.acquire_attempts == 1: + return [0, 1] + self.count = 1 + return [1, 1] + + +def _build_limiter(redis: _FakeRedis) -> PageMemoryVlmLimiter: + limiter = PageMemoryVlmLimiter( + redis, + max_inflight=1, + lease_ttl_seconds=60, + wait_timeout_seconds=1, + ) + limiter._jittered_backoff = lambda current_inflight: 0.001 # type: ignore[method-assign] + return limiter + + +def test_page_memory_vlm_limiter_acquire_release_updates_inflight_count() -> None: + redis = _FakeRedis() + limiter = _build_limiter(redis) + + lease = limiter.acquire(usage_task="page_memory.tag") + assert lease.current_inflight == 1 + assert limiter.get_inflight_count() == 1 + + limiter.release(lease) + assert limiter.get_inflight_count() == 0 + + +def test_page_memory_vlm_limiter_waits_then_succeeds() -> None: + redis = _DelayedCapacityRedis() + limiter = _build_limiter(redis) + + lease = limiter.acquire(usage_task="page_memory.title_detection") + + assert lease.current_inflight == 1 + assert redis.acquire_attempts == 2 + + +def test_page_memory_vlm_limiter_timeout_raises_unavailable() -> None: + redis = _FakeRedis(initial_count=1) + limiter = PageMemoryVlmLimiter( + redis, + max_inflight=1, + lease_ttl_seconds=60, + wait_timeout_seconds=1, + ) + limiter._jittered_backoff = lambda current_inflight: 1.1 # type: ignore[method-assign] + + with pytest.raises(UnavailableException): + limiter.acquire(usage_task="page_memory.node_ocr") + + +def test_page_memory_vlm_limiter_redis_failure_raises_unavailable() -> None: + limiter = _build_limiter(_FakeRedis(fail_eval=True)) + + with pytest.raises(UnavailableException): + limiter.acquire(usage_task="page_memory.node_summary") + + +def test_page_memory_provider_exception_still_releases_lease(monkeypatch) -> None: + class _FakeLimiter: + def __init__(self) -> None: + self.release_count = 0 + + def acquire(self, *, usage_task: str) -> PageMemoryVlmLease: + return PageMemoryVlmLease( + usage_task=usage_task, + acquired_at=0.0, + current_inflight=1, + ) + + def release(self, lease: PageMemoryVlmLease) -> None: + self.release_count += 1 + + class _FakeCompletions: + def create(self, **kwargs: Any) -> Any: + raise RuntimeError("provider failed") + + class _FakeChat: + completions = _FakeCompletions() + + class _FakeSdkClient: + chat = _FakeChat() + base_url = "http://provider.example" + + import shared.services.ai.openai_compatible_client_sync as client_mod + + fake_limiter = _FakeLimiter() + monkeypatch.setattr( + client_mod, + "get_page_memory_vlm_limiter", + lambda: fake_limiter, + ) + client = OpenAICompatibleClientSync( + api_key="test", + api_url="http://provider.example/v1", + default_model="deepseek-test", + ) + client._client = _FakeSdkClient() # type: ignore[assignment] + + with pytest.raises(LLMServiceException): + client.chat_completion_with_usage( + messages="hello", + usage_task="page_memory.tag", + ) + + assert fake_limiter.release_count == 1 + + +def test_page_memory_summary_unavailable_exception_propagates( + monkeypatch, + tmp_path, +) -> None: + class _FakeClient: + def chat_completion_with_usage(self, **kwargs: Any) -> Any: + raise UnavailableException( + internal_message="capacity busy", + retry_after=5, + ) + + image_path = tmp_path / "page.png" + image_path.write_bytes(b"png") + import shared.services.ai.summary.engine as summary_engine + + monkeypatch.setattr( + summary_engine._client_mod, + "get_openai_client", + lambda model=None: _FakeClient(), + ) + + with pytest.raises(UnavailableException): + summarize( + mode="page", + image_paths=[str(image_path)], + model="fake-vlm", + usage_task="page_memory.node_summary", + ) + + +def test_page_memory_transcription_unavailable_exception_propagates( + monkeypatch, + tmp_path, +) -> None: + class _FakeClient: + def chat_completion_with_usage(self, **kwargs: Any) -> Any: + raise UnavailableException( + internal_message="capacity busy", + retry_after=5, + ) + + image_path = tmp_path / "page.png" + image_path.write_bytes(b"png") + import shared.services.ai.summary.engine as summary_engine + + monkeypatch.setattr( + summary_engine._client_mod, + "get_openai_client", + lambda model=None: _FakeClient(), + ) + + with pytest.raises(UnavailableException): + transcribe( + image_paths=[str(image_path)], + model="fake-vlm", + usage_task="page_memory.node_ocr", + ) From e042948a5a2be0c5f75cff56ee80c19388dda6da Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 17:37:54 +0800 Subject: [PATCH 3/5] Tune page-memory concurrency defaults --- .../services/page_memory/node_assembler.py | 2 +- .../app/services/page_memory/page_tagger.py | 2 +- .../shared-python/shared/core/config/ai.py | 12 +++++++++-- .../models/schemas/page_memory_config.py | 20 ++++++++++++------ .../shared/tests/test_page_memory_config.py | 21 +++++++++++++++++++ 5 files changed, 47 insertions(+), 10 deletions(-) create mode 100644 packages/shared-python/shared/tests/test_page_memory_config.py diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index fdc7bf2b1..7f23d4998 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -417,7 +417,7 @@ def build_node_rows( vlm_model: str | None = None, page_assets_by_page: dict[int, list[PageAsset]] | None = None, node_summary_max_pages: int = _NODE_SUMMARY_MAX_PAGES_DEFAULT, - node_assembly_concurrency: int = 2, + node_assembly_concurrency: int = 3, ) -> list[dict[str, Any]]: """Assemble one row per leaf section node (node-granularity chunks).""" available_pages = set(raw_text_by_page.keys()) diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index aa5be87d9..4fe72ac23 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -316,7 +316,7 @@ def tag_page_titles( from shared.core.config import settings resolved_max_concurrent = max_concurrent or int( - getattr(settings, "PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY", 2) + getattr(settings, "PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY", 3) ) def _detect_one( diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index d58315e34..33d711fd2 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -108,12 +108,20 @@ class AIConfig(BaseModel): default=120, description="Max wait before page-memory VLM capacity pressure becomes retryable.", ) + PAGE_MEMORY_SCOPE_CONCURRENCY: int = Field( + default=5, + description="Local per-job hierarchy scope concurrency for page-memory.", + ) + PAGE_MEMORY_TAG_CONCURRENCY: int = Field( + default=4, + description="Local per-job page tagging concurrency for page-memory.", + ) PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY: int = Field( - default=2, + default=3, description="Local per-job title detection concurrency for page-memory.", ) PAGE_MEMORY_NODE_ASSEMBLY_CONCURRENCY: int = Field( - default=2, + default=3, description="Local per-job node OCR and summary concurrency for page-memory.", ) TOKEN_PRICING_TABLE_JSON: str = Field( 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 b895f1476..bbac71ed3 100644 --- a/packages/shared-python/shared/models/schemas/page_memory_config.py +++ b/packages/shared-python/shared/models/schemas/page_memory_config.py @@ -13,8 +13,8 @@ class PageMemoryConfig: max_pages: int = 1500 scope_concurrency: int = 5 tag_concurrency: int = 4 - title_detection_concurrency: int = 2 - node_assembly_concurrency: int = 2 + title_detection_concurrency: int = 3 + node_assembly_concurrency: int = 3 tag_mode: Literal["vlm", "text"] = "vlm" fine_min_pages: int = 4 hierarchy_model: str | None = None @@ -40,13 +40,21 @@ def default(cls) -> Self: from shared.core.config import settings return cls( + scope_concurrency=_as_int( + getattr(settings, "PAGE_MEMORY_SCOPE_CONCURRENCY", 5), + 5, + ), + tag_concurrency=_as_int( + getattr(settings, "PAGE_MEMORY_TAG_CONCURRENCY", 4), + 4, + ), title_detection_concurrency=_as_int( - getattr(settings, "PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY", 2), - 2, + getattr(settings, "PAGE_MEMORY_TITLE_DETECTION_CONCURRENCY", 3), + 3, ), node_assembly_concurrency=_as_int( - getattr(settings, "PAGE_MEMORY_NODE_ASSEMBLY_CONCURRENCY", 2), - 2, + getattr(settings, "PAGE_MEMORY_NODE_ASSEMBLY_CONCURRENCY", 3), + 3, ), ) diff --git a/packages/shared-python/shared/tests/test_page_memory_config.py b/packages/shared-python/shared/tests/test_page_memory_config.py new file mode 100644 index 000000000..ea97ad773 --- /dev/null +++ b/packages/shared-python/shared/tests/test_page_memory_config.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.models.schemas.page_memory_config import PageMemoryConfig + + +def test_page_memory_config_defaults_resolve_concurrency_settings() -> None: + config = PageMemoryConfig.default() + + assert config.scope_concurrency == 5 + assert config.tag_concurrency == 4 + assert config.title_detection_concurrency == 3 + assert config.node_assembly_concurrency == 3 From ba685d6bdfc68eb71dbb5924e604afafe8031331 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 17:51:24 +0800 Subject: [PATCH 4/5] Fix page-memory test import style --- ...est_page_memory_node_assembler_contract.py | 22 ++++++++++++------- .../tests/test_page_memory_vlm_limiter.py | 16 +++++--------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index a502ae3be..6b89292dd 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -11,7 +11,6 @@ import pytest -import app.services.page_memory.node_assembler as node_assembler from app.services.page_memory.node_assembler import ( SAME_AS_PREFIX, assign_pages_to_leaves, @@ -158,10 +157,12 @@ def _fake_compute_node_summary(**kwargs): gevent.sleep(0.01 * view.leaf.start_page) return f"summary-{view.leaf.start_page}", [f"k{view.leaf.start_page}"], [] - monkeypatch.setattr(node_assembler, "resolve_page_text", _fake_resolve_page_text) monkeypatch.setattr( - node_assembler, - "compute_node_summary", + "app.services.page_memory.node_assembler.resolve_page_text", + _fake_resolve_page_text, + ) + monkeypatch.setattr( + "app.services.page_memory.node_assembler.compute_node_summary", _fake_compute_node_summary, ) @@ -197,7 +198,10 @@ def _fake_resolve_page_text(**kwargs) -> str: raise RuntimeError("ocr failed") return "ok" - monkeypatch.setattr(node_assembler, "resolve_page_text", _fake_resolve_page_text) + monkeypatch.setattr( + "app.services.page_memory.node_assembler.resolve_page_text", + _fake_resolve_page_text, + ) with pytest.raises(RuntimeError): build_node_rows( @@ -221,7 +225,10 @@ def _fake_resolve_page_text(**kwargs) -> str: retry_after=5, ) - monkeypatch.setattr(node_assembler, "resolve_page_text", _fake_resolve_page_text) + monkeypatch.setattr( + "app.services.page_memory.node_assembler.resolve_page_text", + _fake_resolve_page_text, + ) with pytest.raises(UnavailableException): build_node_rows( @@ -248,8 +255,7 @@ def _fake_compute_node_summary(**kwargs): ) monkeypatch.setattr( - node_assembler, - "compute_node_summary", + "app.services.page_memory.node_assembler.compute_node_summary", _fake_compute_node_summary, ) diff --git a/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py b/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py index 49c82acd2..c69854395 100644 --- a/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py +++ b/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py @@ -16,14 +16,12 @@ LLMServiceException, UnavailableException, ) -from shared.services.ai.openai_compatible_client_sync import ( # noqa: E402 - OpenAICompatibleClientSync, -) +import shared.services.ai.openai_compatible_client_sync as client_mod # noqa: E402 from shared.services.ai.page_memory_vlm_limiter import ( # noqa: E402 PageMemoryVlmLease, PageMemoryVlmLimiter, ) -from shared.services.ai.summary.engine import summarize, transcribe # noqa: E402 +import shared.services.ai.summary.engine as summary_engine # noqa: E402 class _FakeRedis: @@ -142,15 +140,13 @@ class _FakeSdkClient: chat = _FakeChat() base_url = "http://provider.example" - import shared.services.ai.openai_compatible_client_sync as client_mod - fake_limiter = _FakeLimiter() monkeypatch.setattr( client_mod, "get_page_memory_vlm_limiter", lambda: fake_limiter, ) - client = OpenAICompatibleClientSync( + client = client_mod.OpenAICompatibleClientSync( api_key="test", api_url="http://provider.example/v1", default_model="deepseek-test", @@ -179,7 +175,6 @@ def chat_completion_with_usage(self, **kwargs: Any) -> Any: image_path = tmp_path / "page.png" image_path.write_bytes(b"png") - import shared.services.ai.summary.engine as summary_engine monkeypatch.setattr( summary_engine._client_mod, @@ -188,7 +183,7 @@ def chat_completion_with_usage(self, **kwargs: Any) -> Any: ) with pytest.raises(UnavailableException): - summarize( + summary_engine.summarize( mode="page", image_paths=[str(image_path)], model="fake-vlm", @@ -209,7 +204,6 @@ def chat_completion_with_usage(self, **kwargs: Any) -> Any: image_path = tmp_path / "page.png" image_path.write_bytes(b"png") - import shared.services.ai.summary.engine as summary_engine monkeypatch.setattr( summary_engine._client_mod, @@ -218,7 +212,7 @@ def chat_completion_with_usage(self, **kwargs: Any) -> Any: ) with pytest.raises(UnavailableException): - transcribe( + summary_engine.transcribe( image_paths=[str(image_path)], model="fake-vlm", usage_task="page_memory.node_ocr", From 3d806d1b8de639e8246e4f79b9dadf3e1ff47d31 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 18:02:36 +0800 Subject: [PATCH 5/5] Fix page-memory node assembler test monkeypatching --- ...est_page_memory_node_assembler_contract.py | 69 ++++++++++--------- 1 file changed, 37 insertions(+), 32 deletions(-) diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index 6b89292dd..21304323e 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -11,13 +11,7 @@ import pytest -from app.services.page_memory.node_assembler import ( - SAME_AS_PREFIX, - assign_pages_to_leaves, - build_node_content, - build_node_rows, - identify_leaf_nodes, -) +import app.services.page_memory.node_assembler as node_assembler from app.services.document_parser.support.parser_rows import PARSER_ROW_COLUMNS from app.services.page_memory.page_assets import PageAsset from app.services.page_memory.page_tagger import PageTagResult @@ -72,13 +66,15 @@ def _ordered_page_skeletons() -> list[SectionSkeleton]: def test_identify_leaf_nodes_drops_internal_parents() -> None: - leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) + leaves = node_assembler.identify_leaf_nodes(_same_page_sibling_skeletons()) assert [leaf.title for leaf in leaves] == ["3.1 职责", "3.2 管理规定"] def test_page_ownership_first_leaf_owns_shared_page() -> None: - leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) - views, page_owner = assign_pages_to_leaves(leaves, available_pages={231, 232}) + leaves = node_assembler.identify_leaf_nodes(_same_page_sibling_skeletons()) + views, page_owner = node_assembler.assign_pages_to_leaves( + leaves, available_pages={231, 232} + ) assert page_owner[231].title == "3.1 职责" assert page_owner[232].title == "3.2 管理规定" @@ -90,26 +86,30 @@ def test_page_ownership_first_leaf_owns_shared_page() -> None: def test_build_node_content_uses_same_as_for_shared_page() -> None: - leaves = identify_leaf_nodes(_same_page_sibling_skeletons()) - views, page_owner = assign_pages_to_leaves(leaves, available_pages={231, 232}) + leaves = node_assembler.identify_leaf_nodes(_same_page_sibling_skeletons()) + views, page_owner = node_assembler.assign_pages_to_leaves( + leaves, available_pages={231, 232} + ) by_title = {view.leaf.title: view for view in views} page_text = {231: "text-231", 232: "text-232"} - content_a = build_node_content( + content_a = node_assembler.build_node_content( by_title["3.1 职责"], page_owner=page_owner, page_text=page_text ) - content_b = build_node_content( + content_b = node_assembler.build_node_content( by_title["3.2 管理规定"], page_owner=page_owner, page_text=page_text ) assert content_a == "text-231" - assert content_b.startswith(f"[{SAME_AS_PREFIX} demo.pdf/3 基本规定/3.1 职责 p231]") + assert content_b.startswith( + f"[{node_assembler.SAME_AS_PREFIX} demo.pdf/3 基本规定/3.1 职责 p231]" + ) assert "text-232" in content_b assert "text-231" not in content_b def test_build_node_rows_reuses_tags_without_vlm() -> None: - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={}, @@ -137,7 +137,7 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: leaf_b = by_path["demo.pdf/3 基本规定/3.2 管理规定"] assert leaf_b["page_nums"] == "231,232" - assert SAME_AS_PREFIX in leaf_b["content"] + assert node_assembler.SAME_AS_PREFIX in leaf_b["content"] assert "text-232" in leaf_b["content"] assert leaf_b["extra_metadata"] == {} @@ -158,15 +158,17 @@ def _fake_compute_node_summary(**kwargs): return f"summary-{view.leaf.start_page}", [f"k{view.leaf.start_page}"], [] monkeypatch.setattr( - "app.services.page_memory.node_assembler.resolve_page_text", + node_assembler, + "resolve_page_text", _fake_resolve_page_text, ) monkeypatch.setattr( - "app.services.page_memory.node_assembler.compute_node_summary", + node_assembler, + "compute_node_summary", _fake_compute_node_summary, ) - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_ordered_page_skeletons(), raw_text_by_page={1: "", 2: "", 3: ""}, image_path_by_page={}, @@ -199,12 +201,13 @@ def _fake_resolve_page_text(**kwargs) -> str: return "ok" monkeypatch.setattr( - "app.services.page_memory.node_assembler.resolve_page_text", + node_assembler, + "resolve_page_text", _fake_resolve_page_text, ) with pytest.raises(RuntimeError): - build_node_rows( + node_assembler.build_node_rows( skeletons=_ordered_page_skeletons()[:2], raw_text_by_page={1: "", 2: ""}, image_path_by_page={}, @@ -226,12 +229,13 @@ def _fake_resolve_page_text(**kwargs) -> str: ) monkeypatch.setattr( - "app.services.page_memory.node_assembler.resolve_page_text", + node_assembler, + "resolve_page_text", _fake_resolve_page_text, ) with pytest.raises(UnavailableException): - build_node_rows( + node_assembler.build_node_rows( skeletons=_ordered_page_skeletons()[:1], raw_text_by_page={1: ""}, image_path_by_page={}, @@ -255,12 +259,13 @@ def _fake_compute_node_summary(**kwargs): ) monkeypatch.setattr( - "app.services.page_memory.node_assembler.compute_node_summary", + node_assembler, + "compute_node_summary", _fake_compute_node_summary, ) with pytest.raises(UnavailableException): - build_node_rows( + node_assembler.build_node_rows( skeletons=_ordered_page_skeletons()[:1], raw_text_by_page={1: "text-1"}, image_path_by_page={}, @@ -279,7 +284,7 @@ def test_build_node_rows_attaches_page_citation_assets_for_rendered_pages(tmp_pa page_image.parent.mkdir() Image.new("RGB", (2, 3), color=(255, 255, 255)).save(page_image) - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={231: str(page_image)}, @@ -328,7 +333,7 @@ def test_build_node_rows_keeps_internal_section_body_pages() -> None: parent_path="demo.pdf/4 风险辨识与分级管控", ) - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=[parent, child], raw_text_by_page={233: "parent body", 234: "child body"}, image_path_by_page={}, @@ -374,8 +379,8 @@ def test_boundary_page_belongs_to_next_sibling_start() -> None: parent_path="demo.pdf/安全类", ) - leaves = identify_leaf_nodes([coarse, next_sibling]) - views, page_owner = assign_pages_to_leaves( + leaves = node_assembler.identify_leaf_nodes([coarse, next_sibling]) + views, page_owner = node_assembler.assign_pages_to_leaves( leaves, available_pages={301, 302, 303}, ) @@ -411,7 +416,7 @@ def chat_completion_with_usage(self, **kwargs): img = tmp_path / "page-231.png" img.write_bytes(b"\x89PNG\r\n\x1a\n fake") - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={231: str(img), 232: str(img)}, @@ -452,7 +457,7 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: extraction_status="table_html_extracted", ) - rows = build_node_rows( + rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={},