diff --git a/.gitignore b/.gitignore index 0cc60df6b..8a56b4f57 100644 --- a/.gitignore +++ b/.gitignore @@ -69,6 +69,7 @@ test_*.csv !requirements.csv # Local debugging scripts +apps/worker/scripts/ apps/worker/start_celery_worker.py apps/worker/start_celery_debug.sh apps/worker/clear_celery_queues.sh diff --git a/Makefile b/Makefile index b6488b017..073def0e8 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: lint lint-fix typecheck check +.PHONY: lint lint-fix typecheck check test-doc-agent UV := uv REPO_UV_CACHE_DIR := $(CURDIR)/.uv-cache @@ -29,3 +29,6 @@ typecheck: $(PYRIGHT) --project pyproject.toml $(PYRIGHT_PATHS) check: lint typecheck + +test-doc-agent: + cd apps/worker && $(UV_RUN_ENV) $(UV) run pytest tests/document_agent diff --git a/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py b/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py new file mode 100644 index 000000000..5ae58232b --- /dev/null +++ b/apps/api/alembic/versions/f8a9b0c1d2e3_add_parse_agent_tables.py @@ -0,0 +1,83 @@ +"""add parse agent tables + +Revision ID: f8a9b0c1d2e3 +Revises: f7a8b9c0d1e2 +Create Date: 2026-05-22 10:45:00.000000 + +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "f8a9b0c1d2e3" +down_revision: Union[str, Sequence[str], None] = "f7a8b9c0d1e2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "parse_runs", + sa.Column("run_id", sa.String(length=36), nullable=False), + sa.Column("job_id", sa.String(length=36), nullable=False), + sa.Column("kind", sa.String(length=32), nullable=False, server_default="profile"), + sa.Column("final_status", sa.String(length=32), nullable=False), + sa.Column("rounds_count", sa.Integer(), nullable=False, server_default="0"), + sa.Column("total_tokens", sa.Integer(), nullable=False, server_default="0"), + sa.Column("total_latency_ms", sa.Integer(), nullable=False, server_default="0"), + sa.Column("summary", sa.JSON(), nullable=True), + sa.Column("started_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.Column("finished_at", sa.DateTime(), nullable=True), + sa.ForeignKeyConstraint(["job_id"], ["jobs.job_id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("run_id"), + ) + op.create_index("idx_parse_runs_job_kind", "parse_runs", ["job_id", "kind"]) + op.create_index("idx_parse_runs_started", "parse_runs", ["started_at"]) + + op.create_table( + "parse_steps", + sa.Column("step_id", sa.String(length=36), nullable=False), + sa.Column("run_id", sa.String(length=36), nullable=False), + sa.Column("round_index", sa.Integer(), nullable=False, server_default="0"), + sa.Column("actor", sa.String(length=64), nullable=False), + sa.Column("action_type", sa.String(length=64), nullable=False), + sa.Column("tool_name", sa.String(length=64), nullable=True), + sa.Column("tool_args", sa.JSON(), nullable=True), + sa.Column("observation", sa.JSON(), nullable=True), + sa.Column("tokens_used", sa.Integer(), nullable=False, server_default="0"), + sa.Column("latency_ms", sa.Integer(), nullable=False, server_default="0"), + sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.ForeignKeyConstraint(["run_id"], ["parse_runs.run_id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("step_id"), + ) + op.create_index("idx_parse_steps_run_round", "parse_steps", ["run_id", "round_index"]) + op.create_index("idx_parse_steps_tool", "parse_steps", ["tool_name"]) + + op.create_table( + "document_page_plan", + sa.Column("page_plan_id", sa.String(length=36), nullable=False), + sa.Column("job_id", sa.String(length=36), nullable=False), + sa.Column("page_count", sa.Integer(), nullable=False, server_default="0"), + sa.Column("shard_plan", sa.JSON(), nullable=True), + sa.Column("global_signals", sa.JSON(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.ForeignKeyConstraint(["job_id"], ["jobs.job_id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("page_plan_id"), + ) + op.create_index("idx_document_page_plan_job", "document_page_plan", ["job_id"]) + op.create_index("idx_document_page_plan_created", "document_page_plan", ["created_at"]) + + +def downgrade() -> None: + op.drop_index("idx_document_page_plan_created", table_name="document_page_plan") + op.drop_index("idx_document_page_plan_job", table_name="document_page_plan") + op.drop_table("document_page_plan") + op.drop_index("idx_parse_steps_tool", table_name="parse_steps") + op.drop_index("idx_parse_steps_run_round", table_name="parse_steps") + op.drop_table("parse_steps") + op.drop_index("idx_parse_runs_started", table_name="parse_runs") + op.drop_index("idx_parse_runs_job_kind", table_name="parse_runs") + op.drop_table("parse_runs") diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index d0b401839..6016aaeb6 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -622,7 +622,7 @@ async def test_should_list_current_document_chunks_by_document_id( "chunk_type": "text", "content": "First chunk content", "source_chunk_path": "Chapter 1/Intro", - "metadata": {"summary": "Intro", "page_nums": [1]}, + "metadata": {"summary": "Intro", "page_nums": []}, }, { "id": second_chunk_id, @@ -631,7 +631,7 @@ async def test_should_list_current_document_chunks_by_document_id( "content": "| A | B |", "source_chunk_path": "Chapter 1/Table", "file_path": "tables/table-1.html", - "metadata": {"summary": "Table", "page_nums": [2]}, + "metadata": {"summary": "Table", "page_nums": []}, }, ], ) @@ -666,7 +666,7 @@ async def test_should_list_current_document_chunks_by_document_id( "source_chunk_path": "Chapter 1/Intro", "file_path": None, "sort_order": 0, - "metadata": {"summary": "Intro", "page_nums": [1]}, + "metadata": {"summary": "Intro", "page_nums": []}, "created_at": chunks[0]["created_at"], } ] @@ -720,7 +720,7 @@ async def test_should_return_one_document_chunk_by_document_chunk_id( "content": "Figure summary", "source_chunk_path": "Chapter 1/Figure", "file_path": "images/figure-1.png", - "metadata": {"summary": "Figure", "page_nums": [3]}, + "metadata": {"summary": "Figure", "page_nums": []}, } ], ) @@ -745,7 +745,7 @@ async def test_should_return_one_document_chunk_by_document_chunk_id( assert chunk["section_path"] == "Chapter 1" assert chunk["source_chunk_path"] == "Chapter 1/Figure" assert chunk["file_path"] == "images/figure-1.png" - assert chunk["metadata"] == {"summary": "Figure", "page_nums": [3]} + assert chunk["metadata"] == {"summary": "Figure", "page_nums": []} assert chunk["created_at"] @@ -768,7 +768,7 @@ async def test_should_return_not_found_when_requesting_a_missing_document_chunk( "chunk_type": "text", "content": "First chunk content", "source_chunk_path": "Chapter 1/Intro", - "metadata": {"summary": "Intro", "page_nums": [1]}, + "metadata": {"summary": "Intro", "page_nums": []}, } ], ) diff --git a/apps/worker/app/services/document_agent/__init__.py b/apps/worker/app/services/document_agent/__init__.py index 2e36f6f2e..ee28ee9f6 100644 --- a/apps/worker/app/services/document_agent/__init__.py +++ b/apps/worker/app/services/document_agent/__init__.py @@ -1,15 +1,17 @@ -"""Phase 1 document split-agent utilities.""" +"""Page anatomy agent for hierarchy-first PDF profiling.""" from app.services.document_agent.manifest import ( - GlobalSignals, - ShardManifest, - ShardSignal, - SpecialPage, + PageAnatomyMap, + PageFeature, + PageLabel, + ShardPlan, ) +from app.services.document_agent.profile_agent import ProfileAgent __all__ = [ - "GlobalSignals", - "ShardManifest", - "ShardSignal", - "SpecialPage", + "PageAnatomyMap", + "PageFeature", + "PageLabel", + "ProfileAgent", + "ShardPlan", ] diff --git a/apps/worker/app/services/document_agent/bootstrap/__init__.py b/apps/worker/app/services/document_agent/bootstrap/__init__.py new file mode 100644 index 000000000..b8494b99f --- /dev/null +++ b/apps/worker/app/services/document_agent/bootstrap/__init__.py @@ -0,0 +1,7 @@ +"""Deterministic bootstrap steps for the document profile agent.""" + +from app.services.document_agent.bootstrap.aggregate_stats import aggregate_doc_stats +from app.services.document_agent.bootstrap.classify import classify_page_kinds +from app.services.document_agent.bootstrap.probe import probe_page_features + +__all__ = ["aggregate_doc_stats", "classify_page_kinds", "probe_page_features"] diff --git a/apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py b/apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py new file mode 100644 index 000000000..2177cbf6c --- /dev/null +++ b/apps/worker/app/services/document_agent/bootstrap/aggregate_stats.py @@ -0,0 +1,119 @@ +"""Aggregate page-feature statistics for VLM profile planning.""" + +from __future__ import annotations + +import statistics +import time +from typing import Any + +from app.services.document_agent.manifest import PageFeature, ToolContext, ToolResult + +PROFILE_METRICS = ( + "raw_text_length", + "text_density", + "image_coverage", + "table_count", + "drawings_count", +) + +EXTREMA_ROLES = { + "raw_text_length": ("min", "max"), + "text_density": ("min", "max"), + "image_coverage": ("max",), + "table_count": ("max",), + "drawings_count": ("max",), +} + +EXTREMA_LABELS = { + "raw_text_length": "text_length", + "text_density": "text_density", + "image_coverage": "image_heavy", + "table_count": "table_heavy", + "drawings_count": "drawing_heavy", +} + + +def _percentile(values: list[float], percentile: float) -> float: + if not values: + return 0.0 + if len(values) == 1: + return values[0] + ordered = sorted(values) + index = (len(ordered) - 1) * percentile + lower = int(index) + upper = min(lower + 1, len(ordered) - 1) + weight = index - lower + return ordered[lower] * (1 - weight) + ordered[upper] * weight + + +def _metric_value(feature: PageFeature, metric: str) -> float: + return float(getattr(feature, metric)) + + +def aggregate_doc_stats(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + features = list(ctx.blackboard.page_features) + stats: dict[str, Any] = {} + extrema_pages: list[int] = [] + extrema_samples: list[dict[str, Any]] = [] + for metric in PROFILE_METRICS: + pairs = [(feature.page, _metric_value(feature, metric)) for feature in features] + values = [value for _, value in pairs] + if not pairs: + stats[metric] = { + "mean": 0.0, + "p50": 0.0, + "p90": 0.0, + "min": {"page": None, "value": 0.0}, + "max": {"page": None, "value": 0.0}, + } + continue + min_page, min_value = min(pairs, key=lambda item: (item[1], item[0])) + max_page, max_value = max(pairs, key=lambda item: (item[1], -item[0])) + stats[metric] = { + "mean": round(statistics.fmean(values), 4), + "p50": round(_percentile(values, 0.5), 4), + "p90": round(_percentile(values, 0.9), 4), + "min": {"page": min_page, "value": round(min_value, 4)}, + "max": {"page": max_page, "value": round(max_value, 4)}, + } + extrema_by_role = { + "min": (min_page, min_value), + "max": (max_page, max_value), + } + for role in EXTREMA_ROLES[metric]: + page, value = extrema_by_role[role] + extrema_pages.append(page) + extrema_samples.append( + { + "page": page, + "metric": metric, + "label": EXTREMA_LABELS[metric], + "role": role, + "value": round(value, 4), + } + ) + + deduped_extrema = sorted(set(extrema_pages)) + ctx.blackboard.doc_stats = stats + ctx.blackboard.extrema_pages = deduped_extrema + ctx.blackboard.global_signals["doc_stats"] = stats + ctx.blackboard.global_signals["extrema_pages"] = deduped_extrema + ctx.blackboard.global_signals["extrema_samples"] = extrema_samples + return ToolResult( + status="ok", + payload={ + "metric_count": len(PROFILE_METRICS), + "extrema_pages": deduped_extrema, + "extrema_samples": extrema_samples, + }, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={ + "doc_stats": stats, + "extrema_pages": deduped_extrema, + "extrema_samples": extrema_samples, + }, + ) + + +__all__ = ["PROFILE_METRICS", "aggregate_doc_stats"] diff --git a/apps/worker/app/services/document_agent/bootstrap/classify.py b/apps/worker/app/services/document_agent/bootstrap/classify.py new file mode 100644 index 000000000..c8dc808c7 --- /dev/null +++ b/apps/worker/app/services/document_agent/bootstrap/classify.py @@ -0,0 +1,5 @@ +"""Bootstrap wrapper for deterministic page classification.""" + +from app.services.document_agent.tools.classify_page_kinds import classify_page_kinds + +__all__ = ["classify_page_kinds"] diff --git a/apps/worker/app/services/document_agent/bootstrap/probe.py b/apps/worker/app/services/document_agent/bootstrap/probe.py new file mode 100644 index 000000000..14877b1ba --- /dev/null +++ b/apps/worker/app/services/document_agent/bootstrap/probe.py @@ -0,0 +1,5 @@ +"""Bootstrap wrapper for deterministic page probing.""" + +from app.services.document_agent.tools.probe_page_features import probe_page_features + +__all__ = ["probe_page_features"] diff --git a/apps/worker/app/services/document_agent/budget.py b/apps/worker/app/services/document_agent/budget.py new file mode 100644 index 000000000..6edbc1309 --- /dev/null +++ b/apps/worker/app/services/document_agent/budget.py @@ -0,0 +1,71 @@ +"""Small synchronous budget tracker for parse-side agent planning.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class BudgetPool: + capacity: int + used: int = 0 + reserved: int = 0 + + @property + def remaining(self) -> int: + return max(self.capacity - self.used - self.reserved, 0) + + +class BudgetTracker: + """A minimal synchronous ledger with plan and visual pools.""" + + def __init__( + self, + *, + plan_budget: int = 5000, + visual_budget: int = 8000, + ) -> None: + self._plan = BudgetPool(capacity=max(int(plan_budget), 0)) + self._visual = BudgetPool(capacity=max(int(visual_budget), 0)) + + def try_reserve(self, pool: str, est: int) -> bool: + if pool not in {"plan", "visual"}: + return True + est = max(int(est), 0) + budget_pool = self._pool(pool) + if budget_pool.remaining < est: + return False + budget_pool.reserved += est + return True + + def commit(self, pool: str, *, actual: int, est: int) -> None: + if pool not in {"plan", "visual"}: + return + est = max(int(est), 0) + actual = max(int(actual), 0) + budget_pool = self._pool(pool) + budget_pool.reserved = max(budget_pool.reserved - est, 0) + budget_pool.used = min(budget_pool.capacity, budget_pool.used + actual) + + def refund(self, pool: str, *, est: int) -> None: + if pool not in {"plan", "visual"}: + return + budget_pool = self._pool(pool) + budget_pool.reserved = max(budget_pool.reserved - max(int(est), 0), 0) + + def _pool(self, pool: str) -> BudgetPool: + return self._visual if pool == "visual" else self._plan + + def _pool_snapshot(self, pool: BudgetPool) -> dict[str, int]: + return { + "capacity": pool.capacity, + "used": pool.used, + "reserved": pool.reserved, + "remaining": pool.remaining, + } + + def snapshot(self) -> dict[str, object]: + return { + "plan": self._pool_snapshot(self._plan), + "visual": self._pool_snapshot(self._visual), + } diff --git a/apps/worker/app/services/document_agent/coordinator.py b/apps/worker/app/services/document_agent/coordinator.py new file mode 100644 index 000000000..c32b05d46 --- /dev/null +++ b/apps/worker/app/services/document_agent/coordinator.py @@ -0,0 +1,155 @@ +"""ReAct-style coordinator for the document profile agent.""" + +from __future__ import annotations + +import os +from typing import Any + +from loguru import logger + +from app.services.document_agent.bootstrap import ( + aggregate_doc_stats, + classify_page_kinds, + probe_page_features, +) +from app.services.document_agent.budget import BudgetTracker +from app.services.document_agent.executor import ReActExecutor +from app.services.document_agent.manifest import PageAnatomyMap, ToolContext +from app.services.document_agent.persist import build_anatomy_map, persist_anatomy_map +from app.services.document_agent.planner import ProfilePlanner +from app.services.document_agent.registry import REGISTRY +from app.services.document_agent.state import AgentBlackboard, DocumentAgentState +from app.services.document_agent import tools as _registered_tools # noqa: F401 +from app.services.document_agent.trace import ParseRunRecorder + + +class ProfileCoordinator: + def __init__( + self, + *, + pdf_path: str, + job_id: str, + output_dir: str | None = None, + db: Any | None = None, + model: str | None = None, + settings: dict[str, Any] | None = None, + ) -> None: + self.state = DocumentAgentState.INIT + self.blackboard = AgentBlackboard() + self.budget = BudgetTracker( + plan_budget=int(os.environ.get("PARSE_AGENT_PLAN_BUDGET", "50000")), + visual_budget=int(os.environ.get("PARSE_AGENT_VISUAL_BUDGET", "80000")), + ) + effective_settings = settings or {} + if model: + effective_settings["model"] = model + self.ctx = ToolContext( + pdf_path=pdf_path, + job_id=job_id, + blackboard=self.blackboard, + budget=self.budget, + trace=None, + output_dir=output_dir, + settings=effective_settings, + ) + self.trace = ParseRunRecorder(job_id=job_id, db=db) + self.ctx.trace = self.trace + self.round_index = 0 + + def run(self) -> PageAnatomyMap: + try: + self.state = DocumentAgentState.RUNNING + self._run_bootstrap() + self._run_toc_pipeline() + profile, initial_decision, planner_result = ProfilePlanner(self.ctx).propose() + self.blackboard.document_profile = profile + self.blackboard.global_signals["document_profile"] = profile.to_dict() + self.trace.record_step( + round_index=self.round_index, + actor="planner", + action_type="plan", + result=planner_result, + tool_name=None, + tool_args={}, + ) + self.round_index += 1 + + executor_result = ReActExecutor( + self.ctx, + registry=REGISTRY, + max_rounds=int(self.ctx.settings.get("max_rounds", 30)), + initial_decision=initial_decision, + ).run() + if executor_result.verdict.status != "success": + raise RuntimeError( + f"profile aborted: {executor_result.verdict.rationale}" + ) + anatomy = build_anatomy_map(self.ctx) + persist_result = persist_anatomy_map(self.ctx, {}) + self.trace.record_step( + round_index=self.round_index, + actor="persist", + action_type="persist", + result=persist_result, + tool_name="persist.anatomy_map", + tool_args={}, + ) + self.state = DocumentAgentState.READY + self.trace.write_trace_artifact( + self.ctx.output_dir, + final_status="ready", + summary=anatomy.trace_summary | self.trace.summary(), + ) + self.trace.flush( + final_status="ready", + summary=anatomy.trace_summary | self.trace.summary(), + ) + return anatomy + except Exception as exc: + logger.error(f"[document_agent] profile failed: {exc}") + self.state = DocumentAgentState.FAILED + self.trace.write_trace_artifact( + self.ctx.output_dir, + final_status="failed", + summary={"error": str(exc), "budget": self.ctx.budget.snapshot()}, + ) + self.trace.flush(final_status="failed", summary={"error": str(exc)}) + raise + + def _run_bootstrap(self) -> None: + for tool_name, handler in ( + ("probe.page_features", probe_page_features), + ("classify.page_kinds", classify_page_kinds), + ("aggregate.doc_stats", aggregate_doc_stats), + ): + result = handler(self.ctx, {}) + self.trace.record_step( + round_index=self.round_index, + actor=f"bootstrap:{tool_name}", + action_type="bootstrap", + result=result, + tool_name=tool_name, + tool_args={}, + ) + if result.status != "ok": + raise RuntimeError(result.error or f"{tool_name} failed") + self.round_index += 1 + + def _run_toc_pipeline(self) -> None: + for tool_name in ( + "find.toc_anchor_pages", + "extract.toc_with_boundaries", + "match.h1_pages", + ): + result = REGISTRY.dispatch(tool_name, self.ctx, {}) + self.trace.record_step( + round_index=self.round_index, + actor=f"toc:{tool_name}", + action_type="toc", + result=result, + tool_name=tool_name, + tool_args={}, + ) + if result.status not in {"ok", "invalid"}: + raise RuntimeError(result.error or f"{tool_name} failed") + self.round_index += 1 diff --git a/apps/worker/app/services/document_agent/executor/__init__.py b/apps/worker/app/services/document_agent/executor/__init__.py new file mode 100644 index 000000000..e5656b258 --- /dev/null +++ b/apps/worker/app/services/document_agent/executor/__init__.py @@ -0,0 +1,13 @@ +"""ReAct executor for the document profile agent.""" + +from app.services.document_agent.executor.react_loop import ( + ExecutorResult, + ReActExecutor, + _parse_decision, +) + +__all__ = [ + "ExecutorResult", + "ReActExecutor", + "_parse_decision", +] diff --git a/apps/worker/app/services/document_agent/executor/prompts.py b/apps/worker/app/services/document_agent/executor/prompts.py new file mode 100644 index 000000000..8c225d2cc --- /dev/null +++ b/apps/worker/app/services/document_agent/executor/prompts.py @@ -0,0 +1,15 @@ +"""Prompts for executor reflexion.""" + +REFLEXION_INSTRUCTIONS = ( + "You are the executor of a document profiling agent. Decide the next action " + "from the blackboard facts and available tools. Return strict JSON with keys: " + "action (tool_call or verdict_now), rationale, optional tool_name/tool_args, " + "optional verdict {status, rationale}. Use inspect.pages when more visual " + "evidence is needed, grep.text when native-PDF text evidence is needed, " + "propose.shard_plan when evidence is sufficient to shard, validate.anatomy_map " + "after a shard plan exists, and verdict only after validation succeeds. If a " + "tool failed or validation is invalid, either gather targeted evidence and " + "retry the relevant tool or abort with a clear rationale." +) + +__all__ = ["REFLEXION_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/executor/react_loop.py b/apps/worker/app/services/document_agent/executor/react_loop.py new file mode 100644 index 000000000..ea2daf32d --- /dev/null +++ b/apps/worker/app/services/document_agent/executor/react_loop.py @@ -0,0 +1,305 @@ +"""ReAct-style executor for the document profile agent.""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass +from typing import Any + +from app.services.document_agent.manifest import ( + AgentVerdict, + ReflexionDecision, + ToolContext, + ToolResult, +) +from app.services.document_agent.executor.prompts import REFLEXION_INSTRUCTIONS +from app.services.document_agent.registry import ToolRegistry +from shared.utils.token_estimate import estimate_tokens + + +@dataclass +class ExecutorResult: + verdict: AgentVerdict + rounds: int + + +def _compact_blackboard(ctx: ToolContext) -> dict[str, Any]: + return { + "page_count": ctx.blackboard.page_count, + "page_kind_counts": ctx.blackboard.global_signals.get("page_kind_counts", {}), + "doc_stats": ctx.blackboard.doc_stats, + "extrema_pages": ctx.blackboard.extrema_pages, + "document_profile": ctx.blackboard.document_profile.to_dict() + if ctx.blackboard.document_profile + else None, + "toc_anchor_pages": [anchor.page for anchor in ctx.blackboard.toc_anchor_pages], + "toc_pages": ( + ctx.blackboard.toc_result.toc_pages if ctx.blackboard.toc_result else [] + ), + "toc_hierarchies_count": len(ctx.blackboard.toc_hierarchies or []), + "h1_count": ( + len(ctx.blackboard.h1_result.h1_candidates) + if ctx.blackboard.h1_result + else 0 + ), + "shard_plan": ctx.blackboard.shard_plan.to_dict() + if ctx.blackboard.shard_plan + else None, + "validation_report": ctx.blackboard.validation_report, + "verdict": ctx.blackboard.verdict.to_dict() + if ctx.blackboard.verdict + else None, + "visual_inspections": ctx.blackboard.global_signals.get("visual_inspections", [])[-3:], + "grep_history": ctx.blackboard.global_signals.get("grep_history", [])[-3:], + "budget": ctx.budget.snapshot(), + } + + +def _parse_decision(raw: str) -> ReflexionDecision: + data = json.loads(raw) + action = str(data.get("action") or "tool_call") + if action not in {"tool_call", "verdict_now"}: + action = "tool_call" + verdict = None + if isinstance(data.get("verdict"), dict): + verdict_data = data["verdict"] + status = str(verdict_data.get("status") or "abort") + if status not in {"success", "abort"}: + status = "abort" + verdict = AgentVerdict( + status=status, # type: ignore[arg-type] + rationale=str(verdict_data.get("rationale") or data.get("rationale") or ""), + ) + return ReflexionDecision( + action=action, # type: ignore[arg-type] + rationale=str(data.get("rationale") or ""), + tool_name=data.get("tool_name"), + tool_args=dict(data.get("tool_args") or {}), + verdict=verdict, + ) + + +class ReActExecutor: + def __init__( + self, + ctx: ToolContext, + *, + registry: ToolRegistry, + max_rounds: int = 30, + initial_decision: ReflexionDecision | None = None, + ) -> None: + self.ctx = ctx + self.registry = registry + self.max_rounds = max_rounds + self._initial_decision = initial_decision + + def run(self) -> ExecutorResult: + for round_index in range(self.max_rounds): + pending_recovery_verdict: AgentVerdict | None = None + decision, result = self._next_decision(round_index) + self.ctx.blackboard.global_signals.setdefault("reflexion_decisions", []).append( + decision.to_dict() + ) + if self.ctx.trace: + self.ctx.trace.record_step( + round_index=round_index, + actor=f"executor:r{round_index}", + action_type="reflexion", + result=result, + tool_name=decision.tool_name, + tool_args=decision.tool_args, + ) + + tool_name: str | None = None + tool_args: dict[str, Any] = {} + if decision.action == "verdict_now": + verdict = decision.verdict or AgentVerdict( + status="abort", + rationale=decision.rationale or "Executor stopped without verdict.", + ) + if verdict.status == "success" and not ( + self.ctx.blackboard.validation_report + and self.ctx.blackboard.validation_report.get("valid") is True + ): + decision = ReflexionDecision( + action="tool_call", + rationale=( + "Validate the anatomy map before accepting a success verdict." + ), + tool_name="validate.anatomy_map", + tool_args={}, + ) + tool_name, tool_args = self._resolve_tool_call(decision) + else: + self.ctx.blackboard.verdict = verdict + return ExecutorResult(verdict=verdict, rounds=round_index + 1) + else: + tool_name, tool_args = self._resolve_tool_call(decision) + + if not tool_name: + verdict = AgentVerdict( + status="abort", + rationale="Executor did not choose a tool.", + ) + self.ctx.blackboard.verdict = verdict + return ExecutorResult(verdict=verdict, rounds=round_index + 1) + + tool_result = self.registry.dispatch(tool_name, self.ctx, tool_args) + if self.ctx.trace: + self.ctx.trace.record_step( + round_index=round_index, + actor=f"tool:{tool_name}", + action_type="tool_call", + result=tool_result, + tool_name=tool_name, + tool_args=tool_args, + ) + self.ctx.blackboard.step_history.append( + { + "round": round_index, + "tool_name": tool_name, + "tool_args": tool_args, + "status": tool_result.status, + "error": tool_result.error, + } + ) + if tool_result.status == "error": + pending_recovery_verdict = AgentVerdict( + status="abort", + rationale=tool_result.error or f"{tool_name} failed", + ) + elif tool_result.status == "precondition_unmet": + pending_recovery_verdict = AgentVerdict( + status="abort", + rationale=tool_result.error or f"{tool_name} precondition unmet", + ) + + if self.ctx.blackboard.verdict is not None: + return ExecutorResult( + verdict=self.ctx.blackboard.verdict, + rounds=round_index + 1, + ) + if pending_recovery_verdict is not None and self._is_deterministic_mode(): + self.ctx.blackboard.verdict = pending_recovery_verdict + return ExecutorResult( + verdict=pending_recovery_verdict, + rounds=round_index + 1, + ) + + verdict = AgentVerdict(status="abort", rationale="Maximum executor rounds reached.") + self.ctx.blackboard.verdict = verdict + return ExecutorResult(verdict=verdict, rounds=self.max_rounds) + + def _resolve_tool_call( + self, + decision: ReflexionDecision, + ) -> tuple[str | None, dict[str, Any]]: + if decision.action == "tool_call" and decision.tool_name: + return decision.tool_name, decision.tool_args + return None, {} + + def _is_deterministic_mode(self) -> bool: + return not (self.ctx.settings.get("executor_model") or self.ctx.settings.get("model")) + + def _next_decision(self, round_index: int) -> tuple[ReflexionDecision, ToolResult]: + if round_index == 0 and self._initial_decision is not None: + decision = self._initial_decision + return decision, ToolResult(status="ok", payload=decision.to_dict()) + model = self.ctx.settings.get("executor_model") or self.ctx.settings.get("model") + if not model: + decision = self._deterministic_decision() + return decision, ToolResult(status="ok", payload=decision.to_dict()) + + payload = { + "blackboard": _compact_blackboard(self.ctx), + "history_tail": self.ctx.blackboard.step_history[-6:], + "available_tools": self.registry.openai_specs(self.ctx.blackboard), + "round_index": round_index, + } + prompt = REFLEXION_INSTRUCTIONS + "\nPayload:\n" + json.dumps( + payload, + ensure_ascii=False, + ) + est = estimate_tokens(prompt) + if not self.ctx.budget.try_reserve("plan", est): + decision = ReflexionDecision( + action="verdict_now", + rationale="Planner budget exhausted.", + verdict=AgentVerdict(status="abort", rationale="Planner budget exhausted."), + ) + return decision, ToolResult( + status="ok", + payload=decision.to_dict(), + input_summary=payload, + ) + start = time.monotonic() + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=[{"role": "user", "content": prompt}], + model=model, + temperature=0.0, + max_tokens=1200, + response_format={"type": "json_object"}, + ) + self.ctx.budget.commit( + "plan", + actual=usage.get("total_tokens", est), + est=est, + ) + decision = _parse_decision(raw) + return decision, ToolResult( + status="ok", + payload=decision.to_dict(), + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=usage.get("total_tokens", 0), + input_summary=payload, + debug={"prompt_text": prompt, "raw_response": raw}, + ) + except Exception: + self.ctx.budget.refund("plan", est=est) + raise + + def _deterministic_decision(self) -> ReflexionDecision: + if self.ctx.blackboard.shard_plan is None: + return ReflexionDecision( + action="tool_call", + rationale="Create a shard plan.", + tool_name="propose.shard_plan", + tool_args={}, + ) + if not self.ctx.blackboard.validation_report: + return ReflexionDecision( + action="tool_call", + rationale="Validate the current shard plan.", + tool_name="validate.anatomy_map", + tool_args={}, + ) + if self.ctx.blackboard.validation_report.get("valid") is True: + return ReflexionDecision( + action="tool_call", + rationale="Validation succeeded; finish profile run.", + tool_name="verdict", + tool_args={ + "status": "success", + "rationale": "Validation succeeded; finishing profile run.", + }, + ) + # Validation failed: fallback to single shard instead of aborting. + # Clear the invalid plan and re-propose as a single shard. + from app.services.document_agent.tools.propose_shard_plan import single_shard_plan + + self.ctx.blackboard.shard_plan = single_shard_plan( + self.ctx.blackboard.page_count + ) + self.ctx.blackboard.validation_report = None + return ReflexionDecision( + action="tool_call", + rationale="Validation failed; falling back to single shard plan.", + tool_name="validate.anatomy_map", + tool_args={}, + ) + diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index 7ed3f9cba..ebc3c4e3e 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -1,8 +1,4 @@ -"""Stable Phase 1 shard manifest contract. - -The manifest is intentionally independent from existing parser internals so it -can be produced and inspected before parser integration work begins. -""" +"""Contracts for the hierarchy-first document profile agent.""" from __future__ import annotations @@ -11,127 +7,247 @@ from typing import Any, Literal -SpecialKind = Literal[ - "toc", - "blank", - "sparse", - "table_heavy", - "image_heavy", - "landscape", - "single_image", - "normal", -] +PageKind = Literal["normal", "table_heavy", "image_heavy", "low_content", "landscape"] -PredominantKind = Literal[ - "text_dense", - "table_heavy", - "image_heavy", - "mixed", - "landscape_block", - "toc", - "sparse", -] +ReflexionAction = Literal["tool_call", "verdict_now"] +VerdictStatus = Literal["success", "abort"] @dataclass -class SpecialPage: +class PageFeature: page: int - kind: SpecialKind + raw_text_length: int + text_density: float + image_coverage: float + image_count: int + table_count: int + drawings_count: int + orientation: Literal["portrait", "landscape"] + width: float + height: float + is_blank_like: bool + text_lines_preview: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class PageLabel: + page: int + kind: PageKind confidence: float - note: str = "" + evidence: dict[str, Any] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass -class ShardSignal: - page_start: int - page_end: int - page_offset: int - predominant_kind: PredominantKind - special_pages: list[SpecialPage] = field(default_factory=list) - estimated_difficulty: str | None = None - parser_hint: str | None = None - cut_rationale: str = "" +class DocumentProfile: + is_scanned: bool + category: str + category_rationale: str = "" + language: str = "unknown" + rationale: str = "" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class AgentVerdict: + status: VerdictStatus + rationale: str + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class ReflexionDecision: + action: ReflexionAction + rationale: str + tool_name: str | None = None + tool_args: dict[str, Any] = field(default_factory=dict) + verdict: AgentVerdict | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "action": self.action, + "rationale": self.rationale, + "tool_name": self.tool_name, + "tool_args": dict(self.tool_args), + "verdict": self.verdict.to_dict() if self.verdict else None, + } + + +@dataclass +class TocCandidate: + title: str + normalized_title: str + source_page: int + line_index: int + numbering: str = "" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class TocAnchorPage: + """A candidate TOC start page identified by keyword scan, pending VLM confirmation.""" + + page: int # 1-based page number + png_path: str # local PNG path for VLM inspection + source: Literal["page_label", "text_scan", "visual_scan"] # how this anchor was discovered + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class TocResult: + toc_pages: list[int] = field(default_factory=list) + candidates: list[TocCandidate] = field(default_factory=list) + method: Literal["toc_marker", "vlm_progressive", "visual_scan", "none"] = "none" + notes: str = "" def to_dict(self) -> dict[str, Any]: data = asdict(self) - data["special_pages"] = [page.to_dict() for page in self.special_pages] + data["candidates"] = [candidate.to_dict() for candidate in self.candidates] return data @dataclass -class GlobalSignals: - has_toc: bool - toc_pages: list[int] - landscape_ratio: float - table_page_ratio: float - image_page_ratio: float - sample_size: int +class H1Candidate: + title: str + page: int + confidence: float + matched_line: str + source: Literal["toc_exact_top", "toc_fuzzy_top", "heading_grep", "none"] + evidence: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class H1BoundaryResult: + h1_candidates: list[H1Candidate] = field(default_factory=list) + method: Literal["toc_grep", "heading_grep", "none"] = "none" notes: str = "" + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + data["h1_candidates"] = [candidate.to_dict() for candidate in self.h1_candidates] + return data + + +@dataclass +class ValidationReport: + valid: bool + errors: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass +class Shard: + shard_index: int + page_start: int + page_end: int + page_offset: int + anchor_type: Literal["h1_boundary", "blank_separator", "forced_max_size"] + anchor_evidence: str + confidence: float + def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass -class ShardManifest: +class ShardPlan: + enabled: bool + reason: Literal[ + "too_large", + "not_needed", + "parser_stability", + "hierarchy_isolation", + "llm_boundary_decision", + ] + shards: list[Shard] = field(default_factory=list) + validation: ValidationReport = field( + default_factory=lambda: ValidationReport(valid=True) + ) + + def to_dict(self) -> dict[str, Any]: + return { + "enabled": self.enabled, + "reason": self.reason, + "shards": [shard.to_dict() for shard in self.shards], + "validation": self.validation.to_dict(), + } + + +@dataclass +class PageAnatomyMap: job_id: str - file_uri: str - file_sha: str + file_path: str page_count: int - shard_count: int - shards: list[ShardSignal] - global_signals: GlobalSignals - decision_log_ref: str + page_features: list[PageFeature] + page_labels: list[PageLabel] + toc_result: TocResult + h1_result: H1BoundaryResult + shard_plan: ShardPlan + document_profile: DocumentProfile | None = None + toc_hierarchies: list[dict[str, Any]] | 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)) version: str = "1.0" - def validate(self) -> None: - """Enforce the downstream coverage contract.""" - if self.page_count < 0: - raise ValueError("page_count must be non-negative") - if self.shard_count != len(self.shards): - raise ValueError("shard_count must match shards length") - if self.page_count == 0: - if self.shards: - raise ValueError("empty documents cannot contain shards") - return - - expected_start = 1 - for shard in self.shards: - if shard.page_start != expected_start: - raise ValueError( - f"non-contiguous shard coverage at page {expected_start}: " - f"got start={shard.page_start}" - ) - if shard.page_end < shard.page_start: - raise ValueError( - f"invalid shard range {shard.page_start}-{shard.page_end}" - ) - if shard.page_offset != shard.page_start - 1: - raise ValueError( - f"invalid page_offset for shard {shard.page_start}-{shard.page_end}" - ) - expected_start = shard.page_end + 1 - - if expected_start != self.page_count + 1: - raise ValueError( - f"shards must cover 1..{self.page_count}, stopped at {expected_start - 1}" - ) - def to_dict(self) -> dict[str, Any]: - self.validate() return { "version": self.version, "job_id": self.job_id, - "file_uri": self.file_uri, - "file_sha": self.file_sha, + "file_path": self.file_path, "page_count": self.page_count, - "shard_count": self.shard_count, - "shards": [shard.to_dict() for shard in self.shards], - "global_signals": self.global_signals.to_dict(), - "decision_log_ref": self.decision_log_ref, + "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(), + "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, + "global_signals": dict(self.global_signals), + "trace_summary": dict(self.trace_summary), "created_at": self.created_at.isoformat(), } + + +@dataclass +class ToolResult: + status: str + payload: dict[str, Any] = field(default_factory=dict) + latency_ms: int = 0 + error: str | None = None + tokens_used: int = 0 + input_summary: dict[str, Any] | None = None + output_summary: dict[str, Any] | None = None + warnings: list[str] = field(default_factory=list) + debug: dict[str, Any] | None = None + +@dataclass +class ToolContext: + pdf_path: str + job_id: str + blackboard: Any + budget: Any + trace: Any + output_dir: str | None = None + settings: dict[str, Any] = field(default_factory=dict) diff --git a/apps/worker/app/services/document_agent/pdf_text.py b/apps/worker/app/services/document_agent/pdf_text.py new file mode 100644 index 000000000..ca9684370 --- /dev/null +++ b/apps/worker/app/services/document_agent/pdf_text.py @@ -0,0 +1,60 @@ +"""PyMuPDF helpers used by document-agent tools.""" + +from __future__ import annotations + +import gc +from typing import Any + +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) + + +def normalize_spaces(text: str) -> str: + return " ".join((text or "").split()) + + +@worker +def _read_page_texts_worker(queue, pdf_path: str, pages: list[int]) -> None: + import pymupdf # type: ignore[import] + + texts: dict[int, str] = {} + try: + doc = pymupdf.open(pdf_path) + for page in pages: + idx = page - 1 + if 0 <= idx < doc.page_count: + texts[page] = str(doc[idx].get_text() or "") + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "texts": texts}) + + +def read_page_texts( + pdf_path: str, + pages: list[int], + *, + timeout: int = 180, +) -> dict[int, str]: + if not pages: + return {} + result = run_in_child_process(_read_page_texts_worker, pdf_path, pages, timeout=timeout) + return {int(k): str(v) for k, v in (result.get("texts") or {}).items()} + + +def meaningful_lines(text: str) -> list[str]: + return [normalize_spaces(line) for line in text.splitlines() if normalize_spaces(line)] + + +def top_lines(text: str, *, max_lines: int = 20) -> list[str]: + lines = meaningful_lines(text) + return lines[: max(max_lines, 0)] + + +def compact_payload_keys(payload: dict[str, Any]) -> list[str]: + return sorted(str(key) for key in payload.keys()) diff --git a/apps/worker/app/services/document_agent/persist/__init__.py b/apps/worker/app/services/document_agent/persist/__init__.py new file mode 100644 index 000000000..a34439602 --- /dev/null +++ b/apps/worker/app/services/document_agent/persist/__init__.py @@ -0,0 +1,8 @@ +"""Persist anatomy map artifacts.""" + +from app.services.document_agent.tools.persist_anatomy_map import ( + build_anatomy_map, + persist_anatomy_map, +) + +__all__ = ["build_anatomy_map", "persist_anatomy_map"] diff --git a/apps/worker/app/services/document_agent/planner/__init__.py b/apps/worker/app/services/document_agent/planner/__init__.py new file mode 100644 index 000000000..09de37f57 --- /dev/null +++ b/apps/worker/app/services/document_agent/planner/__init__.py @@ -0,0 +1,13 @@ +"""One-shot VLM profile planner.""" + +from app.services.document_agent.planner.planner import ( + PAGE_KIND_DEFINITIONS, + ProfilePlanner, + _sample_pages, +) + +__all__ = [ + "PAGE_KIND_DEFINITIONS", + "ProfilePlanner", + "_sample_pages", +] diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py new file mode 100644 index 000000000..61d63f9f5 --- /dev/null +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -0,0 +1,295 @@ +"""Initial VLM profile planner for the document profile agent.""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from typing import Any, cast + +from loguru import logger + +from app.services.document_agent.manifest import ( + DocumentProfile, + ReflexionDecision, + ToolContext, + ToolResult, +) +from app.services.document_agent.planner.prompts import PLANNER_INSTRUCTIONS +from app.services.document_agent.visual import render_pages +from shared.utils.token_estimate import estimate_tokens + +PAGE_KIND_DEFINITIONS = { + "normal": ( + "Page with enough extractable native text and no dominant table/image " + "structure." + ), + "table_heavy": ( + "Page with detected tables or many vector drawings, often financial " + "tables or dense tabular layout." + ), + "image_heavy": ( + "Page dominated by image coverage with little extractable native text; " + "may be scanned, infographic, photo, or rendered page." + ), + "low_content": ( + "Page with very little extractable text and little visual/table content; " + "may be blank, separator, short heading page, or sparse transition page." + ), + "landscape": "Landscape-oriented page, often wide tables, drawings, slides, or diagrams.", +} + + +def _feature_rows(ctx: ToolContext, pages: list[int]) -> list[dict[str, Any]]: + labels_by_page = {label.page: label for label in ctx.blackboard.page_labels} + selected = [] + for feature in ctx.blackboard.page_features: + if feature.page not in pages: + continue + label = labels_by_page.get(feature.page) + selected.append( + { + "page": feature.page, + "kind": label.kind if label else None, + "confidence": label.confidence if label else None, + "raw_text_length": feature.raw_text_length, + "text_density": feature.text_density, + "image_coverage": feature.image_coverage, + "image_count": feature.image_count, + "table_count": feature.table_count, + "drawings_count": feature.drawings_count, + "orientation": feature.orientation, + "is_blank_like": feature.is_blank_like, + } + ) + return selected + + +def _segment_sample(candidates: list[int], count: int) -> list[int]: + if count <= 0 or not candidates: + return [] + if len(candidates) <= count: + return candidates + if count == 1: + return [candidates[len(candidates) // 2]] + step = (len(candidates) - 1) / (count - 1) + return [candidates[round(index * step)] for index in range(count)] + + +def _sample_pages(page_count: int, extrema_pages: list[int]) -> list[int]: + if page_count <= 0: + return [] + extrema = [page for page in extrema_pages if 1 <= page <= page_count] + remaining = [page for page in range(1, page_count + 1) if page not in set(extrema)] + if not remaining: + return sorted(set(extrema)) + third = max(len(remaining) // 3, 1) + front = remaining[:third] + middle = remaining[third : third * 2] + back = remaining[third * 2 :] + sampled = ( + _segment_sample(front, 4) + + _segment_sample(middle or remaining, 3) + + _segment_sample(back or remaining, 3) + ) + ordered = [] + for page in extrema + sampled: + if page not in ordered: + ordered.append(page) + return ordered[:20] + + +def _parse_profile_and_decision(raw: str) -> tuple[DocumentProfile, ReflexionDecision]: + data = json.loads(raw) + if not isinstance(data, dict): + raise ValueError("planner output must be a JSON object") + category = " ".join(str(data.get("category") or "unknown document").split()[:5]) + raw_is_scanned = data.get("is_scanned") + if isinstance(raw_is_scanned, bool): + is_scanned = raw_is_scanned + elif isinstance(raw_is_scanned, str): + is_scanned = raw_is_scanned.strip().lower() in {"true", "yes", "1", "scanned"} + else: + is_scanned = bool(raw_is_scanned) + profile = DocumentProfile( + is_scanned=is_scanned, + category=category or "unknown document", + category_rationale=str(data.get("category_rationale") or ""), + language=str(data.get("language") or "unknown"), + rationale=str(data.get("rationale") or ""), + ) + next_action = str(data.get("next_action") or "ready_to_shard") + tool_name: str | None = None + tool_args: dict[str, Any] = {} + if next_action == "inspect_more": + pages = [int(page) for page in (data.get("inspect_pages") or [])] + tool_name = "inspect.pages" + tool_args = { + "pages": pages[:10], + "question": "Clarify the document structure and whether these pages change the profile or sharding strategy.", + } + elif next_action == "grep_text" and not profile.is_scanned: + query = str(data.get("grep_query") or "").strip() + if query: + tool_name = "grep.text" + tool_args = {"query": query, "max_results": 20} + elif next_action == "verdict_now": + return profile, ReflexionDecision( + action="verdict_now", + rationale=profile.rationale, + verdict=None, + ) + if tool_name: + return profile, ReflexionDecision( + action="tool_call", + rationale=profile.rationale, + tool_name=tool_name, + tool_args=tool_args, + ) + return profile, ReflexionDecision( + action="tool_call", + rationale=profile.rationale, + tool_name="propose.shard_plan", + tool_args={}, + ) + + +class ProfilePlanner: + """One-shot VLM planner that profiles the document and proposes the first action.""" + + def __init__(self, ctx: ToolContext) -> None: + self.ctx = ctx + + def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: + start = time.monotonic() + model = ( + self.ctx.settings.get("planner_model") + or self.ctx.settings.get("vlm_model") + or os.environ.get("IMAGE_MODEL") + ) + pages = _sample_pages(self.ctx.blackboard.page_count, self.ctx.blackboard.extrema_pages) + if not model: + profile = DocumentProfile( + is_scanned=False, + category="unknown document", + rationale="No planner model configured.", + ) + decision = ReflexionDecision( + action="tool_call", + rationale=profile.rationale, + tool_name="propose.shard_plan", + tool_args={}, + ) + return profile, decision, ToolResult( + status="ok", + payload={"source": "deterministic", "sampled_pages": pages}, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["No planner model configured; using conservative profile."], + input_summary={"page_count": self.ctx.blackboard.page_count}, + output_summary={"profile": profile.to_dict(), "decision": decision.to_dict()}, + ) + pngs = render_pages( + self.ctx, + pages, + folder_name="planner_pages", + prefix="planner", + timeout=180, + ) + feature_summary = _feature_rows(self.ctx, pages) + payload = { + "page_count": self.ctx.blackboard.page_count, + "page_kind_counts": self.ctx.blackboard.global_signals.get( + "page_kind_counts", + {}, + ), + "page_kind_definitions": PAGE_KIND_DEFINITIONS, + "doc_stats": self.ctx.blackboard.doc_stats, + "extrema_samples": self.ctx.blackboard.global_signals.get( + "extrema_samples", + [], + ), + "sampled_page_features": feature_summary, + "toc_pages": self.ctx.blackboard.toc_result.toc_pages + if self.ctx.blackboard.toc_result + else [], + "h1_pages": [ + {"title": item.title, "page": item.page} + for item in ( + self.ctx.blackboard.h1_result.h1_candidates + if self.ctx.blackboard.h1_result + else [] + ) + ], + "available_actions": [ + "inspect.pages", + "grep.text", + "propose.shard_plan", + "validate.anatomy_map", + "verdict", + ], + } + prompt_text = PLANNER_INSTRUCTIONS + "\nPayload:\n" + json.dumps( + payload, + ensure_ascii=False, + ) + prompt_tokens_est = estimate_tokens(prompt_text) + len(pngs) * 800 + if not self.ctx.budget.try_reserve("visual", prompt_tokens_est): + raise RuntimeError("Insufficient visual budget for profile planning.") + + content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt_text}] + for item in pngs: + try: + with open(str(item["png_path"]), "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + content_parts.append( + {"type": "text", "text": f"\n--- Page {item['page']} ---"} + ) + content_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + except Exception as exc: + logger.warning("[document_agent] planner png attach failed: {}", exc) + + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=0.0, + max_tokens=1800, + response_format={"type": "json_object"}, + ) + self.ctx.budget.commit( + "visual", + actual=usage.get("total_tokens", prompt_tokens_est), + est=prompt_tokens_est, + ) + profile, decision = _parse_profile_and_decision(raw) + return profile, decision, ToolResult( + status="ok", + payload={ + "source": "llm", + "sampled_pages": pages, + "first_action": decision.tool_name, + }, + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=usage.get("total_tokens", 0), + input_summary=payload, + output_summary={"profile": profile.to_dict(), "decision": decision.to_dict()}, + debug={ + "prompt_text": prompt_text, + "sampled_pngs": pngs, + "raw_response": raw, + }, + ) + except Exception: + self.ctx.budget.refund("visual", est=prompt_tokens_est) + raise + + diff --git a/apps/worker/app/services/document_agent/planner/prompts.py b/apps/worker/app/services/document_agent/planner/prompts.py new file mode 100644 index 000000000..49feab396 --- /dev/null +++ b/apps/worker/app/services/document_agent/planner/prompts.py @@ -0,0 +1,16 @@ +"""Prompts for the document profile planner.""" + +PLANNER_INSTRUCTIONS = ( + "You are a document profile agent. Use global page-feature statistics, " + "TOC/H1 evidence, and page screenshots to classify the document and decide " + "whether enough evidence exists to continue toward sharding. Return strict " + "JSON only with keys: is_scanned, category, category_rationale, language, " + "rationale, next_action, inspect_pages, grep_query. category must be at " + "most 5 English words. next_action must be one of inspect_more, grep_text, " + "ready_to_shard, verdict_now. Use inspect_more only when specific extra " + "page screenshots are needed. Use grep_text only for native PDFs when a " + "global text search would clarify structure. Do not output a fixed step " + "plan." +) + +__all__ = ["PLANNER_INSTRUCTIONS"] diff --git a/apps/worker/app/services/document_agent/profile_agent.py b/apps/worker/app/services/document_agent/profile_agent.py new file mode 100644 index 000000000..6a73cf74d --- /dev/null +++ b/apps/worker/app/services/document_agent/profile_agent.py @@ -0,0 +1,40 @@ +"""Public entrypoint for document page anatomy profiling.""" + +from __future__ import annotations + +import os +from typing import Any + +from app.services.document_agent.coordinator import ProfileCoordinator +from app.services.document_agent.manifest import PageAnatomyMap + + +class ProfileAgent: + def __init__( + self, + *, + model: str | None = None, + settings: dict[str, Any] | None = None, + ) -> None: + self._model = model + self._settings = settings or {} + + def run( + self, + file_path: str, + job_id: str, + *, + output_dir: str | None = None, + db: Any | None = None, + ) -> PageAnatomyMap: + if not os.path.exists(file_path): + raise FileNotFoundError(file_path) + coordinator = ProfileCoordinator( + pdf_path=file_path, + job_id=job_id, + output_dir=output_dir, + db=db, + model=self._model, + settings=self._settings, + ) + return coordinator.run() diff --git a/apps/worker/app/services/document_agent/registry.py b/apps/worker/app/services/document_agent/registry.py new file mode 100644 index 000000000..80a17531f --- /dev/null +++ b/apps/worker/app/services/document_agent/registry.py @@ -0,0 +1,163 @@ +"""Agent tool registry with blackboard-based preconditions.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.state import AgentBlackboard + +ToolHandler = Callable[[ToolContext, dict[str, Any]], ToolResult] +Precondition = Callable[[AgentBlackboard], tuple[bool, str]] + + +def _always(_blackboard: AgentBlackboard) -> tuple[bool, str]: + return True, "" + + +@dataclass(frozen=True) +class ToolSpec: + name: str + description: str + parameters: dict[str, Any] + preconditions: tuple[Precondition, ...] + handler: ToolHandler + + def to_openai_schema(self) -> dict[str, Any]: + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.parameters, + }, + } + + +class ToolRegistry: + def __init__(self) -> None: + self._tools: dict[str, ToolSpec] = {} + + def register(self, spec: ToolSpec) -> None: + self._tools[spec.name] = spec + + def get(self, name: str) -> ToolSpec | None: + return self._tools.get(name) + + def openai_specs(self, blackboard: AgentBlackboard) -> list[dict[str, Any]]: + return [ + tool.to_openai_schema() + for tool in self._tools.values() + if self._preconditions_met(tool, blackboard)[0] + ] + + def allowed_names(self, blackboard: AgentBlackboard) -> list[str]: + return [ + name + for name, tool in self._tools.items() + if self._preconditions_met(tool, blackboard)[0] + ] + + def _preconditions_met( + self, + tool: ToolSpec, + blackboard: AgentBlackboard, + ) -> tuple[bool, str]: + for check in tool.preconditions: + ok, reason = check(blackboard) + if not ok: + return False, reason + return True, "" + + def dispatch( + self, + name: str, + ctx: ToolContext, + args: dict[str, Any], + ) -> ToolResult: + tool = self.get(name) + if tool is None: + return ToolResult(status="error", error=f"unknown tool: {name}") + ok, reason = self._preconditions_met(tool, ctx.blackboard) + if not ok: + return ToolResult( + status="precondition_unmet", + payload={ + "allowed_tools": self.allowed_names(ctx.blackboard), + "tool": name, + "reason": reason, + }, + error=reason, + ) + return tool.handler(ctx, args) + + +REGISTRY = ToolRegistry() + + +def register_tool( + *, + name: str, + description: str, + parameters: dict[str, Any] | None = None, + preconditions: tuple[Precondition, ...] | None = None, +) -> Callable[[ToolHandler], ToolHandler]: + def _decorator(handler: ToolHandler) -> ToolHandler: + REGISTRY.register( + ToolSpec( + name=name, + description=description, + parameters=parameters + or {"type": "object", "properties": {}, "required": []}, + preconditions=preconditions or (_always,), + handler=handler, + ) + ) + return handler + + return _decorator + + +def has_page_features(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.page_features), "page_features missing; run bootstrap probe first" + + +def has_page_labels(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.page_labels), "page_labels missing; run bootstrap classify first" + + +def has_doc_stats(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.doc_stats), "doc_stats missing; run bootstrap aggregate first" + + +def has_document_profile(blackboard: AgentBlackboard) -> tuple[bool, str]: + return blackboard.document_profile is not None, "document_profile missing; run planner first" + + +def not_is_scanned(blackboard: AgentBlackboard) -> tuple[bool, str]: + profile = blackboard.document_profile + return ( + profile is not None and not profile.is_scanned, + "document is scanned or profile is missing; text grep is unavailable", + ) + + +def has_toc_anchors(blackboard: AgentBlackboard) -> tuple[bool, str]: + return bool(blackboard.toc_anchor_pages), "toc anchors missing; call find_toc_anchors first" + + +def has_toc_result(blackboard: AgentBlackboard) -> tuple[bool, str]: + return blackboard.toc_result is not None, "toc_result missing; call extract_toc first" + + +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 new file mode 100644 index 000000000..bcdb1a1c5 --- /dev/null +++ b/apps/worker/app/services/document_agent/state.py @@ -0,0 +1,47 @@ +"""State carried by the document profile agent.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +from app.services.document_agent.manifest import ( + AgentVerdict, + DocumentProfile, + H1BoundaryResult, + PageFeature, + PageLabel, + ShardPlan, + TocAnchorPage, + TocResult, +) + + +class DocumentAgentState(str, Enum): + INIT = "init" + RUNNING = "running" + READY = "ready" + FAILED = "failed" + + +@dataclass +class AgentBlackboard: + page_count: int = 0 + document_profile: DocumentProfile | None = None + page_features: list[PageFeature] = field(default_factory=list) + page_labels: list[PageLabel] = field(default_factory=list) + doc_stats: dict[str, Any] = field(default_factory=dict) + extrema_pages: list[int] = field(default_factory=list) + toc_anchor_pages: list[TocAnchorPage] = field(default_factory=list) + toc_result: TocResult | None = None + toc_hierarchies: list[dict[str, Any]] | None = None + h1_result: H1BoundaryResult | None = None + shard_plan: ShardPlan | None = None + validation_report: dict[str, Any] | None = None + verdict: AgentVerdict | None = None + step_history: list[dict[str, Any]] = field(default_factory=list) + page_full_text_cache: dict[int, str] = field(default_factory=dict) + global_signals: dict[str, Any] = field(default_factory=dict) + errors: list[str] = field(default_factory=list) + diff --git a/apps/worker/app/services/document_agent/tools/__init__.py b/apps/worker/app/services/document_agent/tools/__init__.py index 8f54e85fa..00260327d 100644 --- a/apps/worker/app/services/document_agent/tools/__init__.py +++ b/apps/worker/app/services/document_agent/tools/__init__.py @@ -1,15 +1,14 @@ -"""Document split-agent Phase 1 tools.""" +"""Import tool modules so decorators register handlers.""" -from app.services.document_agent.tools.classify_special_pages import ( - classify_special_pages, -) -from app.services.document_agent.tools.probe_sample_pages import sample_pages -from app.services.document_agent.tools.probe_vlm_inspect import vlm_inspect_pages -from app.services.document_agent.tools.propose_shard_plan import propose_shard_plan +from app.services.document_agent.registry import REGISTRY -__all__ = [ - "classify_special_pages", - "propose_shard_plan", - "sample_pages", - "vlm_inspect_pages", -] +from . import extract_toc_with_boundaries as extract_toc_with_boundaries # noqa: F401 +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 propose_shard_plan as propose_shard_plan # noqa: F401 +from . import validate_anatomy_map as validate_anatomy_map # noqa: F401 +from . import verdict as verdict # noqa: F401 + +__all__ = ["REGISTRY"] diff --git a/apps/worker/app/services/document_agent/tools/classify_page_kinds.py b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py new file mode 100644 index 000000000..fb60765d7 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/classify_page_kinds.py @@ -0,0 +1,85 @@ +"""Rule-based page kind classification.""" + +from __future__ import annotations + +import time +from collections import Counter, defaultdict +from typing import Any + +from app.services.document_agent.manifest import PageFeature, PageLabel, ToolContext, ToolResult + + +def _label_feature(feature: PageFeature) -> PageLabel: + page = feature.page + if ( + feature.raw_text_length < 80 + and feature.image_coverage < 0.02 + and feature.drawings_count < 5 + ): + return PageLabel( + page=page, + kind="low_content", + confidence=0.78, + evidence={"signal": "low_text_image_drawings"}, + ) + if feature.orientation == "landscape": + return PageLabel( + page=page, + kind="landscape", + confidence=0.78, + evidence={"width": feature.width, "height": feature.height}, + ) + if feature.image_coverage >= 0.35 and feature.raw_text_length < 250: + return PageLabel( + page=page, + kind="image_heavy", + confidence=0.84, + evidence={"image_coverage": feature.image_coverage}, + ) + if feature.table_count > 0 or feature.drawings_count >= 80: + return PageLabel( + page=page, + kind="table_heavy", + confidence=0.72, + evidence={ + "table_count": feature.table_count, + "drawings_count": feature.drawings_count, + }, + ) + return PageLabel(page=page, kind="normal", confidence=0.65, evidence={}) + + +def classify_page_kinds(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + labels = [_label_feature(feature) for feature in ctx.blackboard.page_features] + ctx.blackboard.page_labels = labels + counts = Counter(label.kind for label in labels) + ctx.blackboard.global_signals["page_kind_counts"] = counts + features_by_page = {feature.page: feature for feature in ctx.blackboard.page_features} + samples: dict[str, list[dict[str, Any]]] = defaultdict(list) + for label in labels: + if len(samples[label.kind]) >= 8: + continue + feature = features_by_page.get(label.page) + samples[label.kind].append( + { + "page": label.page, + "confidence": label.confidence, + "evidence": label.evidence, + "raw_text_length": feature.raw_text_length if feature else None, + "image_coverage": feature.image_coverage if feature else None, + "table_count": feature.table_count if feature else None, + "drawings_count": feature.drawings_count if feature else None, + "text_preview": (feature.text_lines_preview[:4] if feature else []), + } + ) + return ToolResult( + status="ok", + payload={"page_kind_counts": dict(counts)}, + latency_ms=int((time.monotonic() - start) * 1000), + input_summary={"page_count": ctx.blackboard.page_count}, + output_summary={ + "page_kind_counts": dict(counts), + "sample_pages_by_kind": dict(samples), + }, + ) diff --git a/apps/worker/app/services/document_agent/tools/classify_special_pages.py b/apps/worker/app/services/document_agent/tools/classify_special_pages.py deleted file mode 100644 index 0b8efa6c8..000000000 --- a/apps/worker/app/services/document_agent/tools/classify_special_pages.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Classify sampled pages into Phase 1 special page kinds.""" - -from __future__ import annotations - -import json -from typing import Any - -from app.services.document_agent.manifest import SpecialKind -from app.services.document_agent.tools.llm_json import extract_json_object -from loguru import logger - -ALLOWED_KINDS: set[str] = { - "toc", - "blank", - "sparse", - "table_heavy", - "image_heavy", - "landscape", - "single_image", - "normal", -} - -PROMPT = """You are classifying sampled PDF pages for a document split planner. -Return ONLY valid json. - -Allowed special_kind values: -- toc: table of contents / contents pages -- blank: blank page -- sparse: very little useful content -- table_heavy: mostly tables or dense tabular rules -- image_heavy: many charts/photos/figures -- landscape: landscape page that should not be split through a landscape block -- single_image: one large screenshot/scanned image dominates the page -- normal: ordinary text page - -Use structural features first, and use VLM observations when present. Be conservative: -only mark toc/table/image/landscape/single_image when there is clear evidence. - -JSON schema: -{ - "pages": [ - {"page": 1, "special_kind": "normal", "confidence": 0.75, "note": "short reason"} - ], - "global_notes": "short summary" -} -""" - - -def _heuristic_kind(page: dict[str, Any]) -> tuple[SpecialKind, float, str]: - text = str(page.get("text_preview") or "") - text_len = int(page.get("text_length") or 0) - image_coverage = float(page.get("image_coverage") or 0.0) - table_count = int(page.get("table_count") or 0) - drawings_count = int(page.get("drawings_count") or 0) - orientation = str(page.get("orientation") or "") - is_blank_like = bool(page.get("is_blank_like")) - - toc_markers = ["目录", "contents", "table of contents"] - if any(marker.lower() in text.lower() for marker in toc_markers): - return "toc", 0.82, "text preview contains TOC marker" - if is_blank_like: - return "blank", 0.9, "very low text/image/drawing signal" - if image_coverage >= 0.72 and text_len < 250: - return "single_image", 0.82, "one or more images dominate the page" - if table_count > 0 or drawings_count >= 80: - return "table_heavy", 0.72, "table detector or dense ruled drawings fired" - if image_coverage >= 0.35: - return "image_heavy", 0.72, "high image coverage" - if orientation == "landscape": - return "landscape", 0.75, "page is landscape" - if text_len < 80: - return "sparse", 0.68, "short text and no stronger special signal" - return "normal", 0.65, "no special signal" - - -def heuristic_classify_special_pages( - sampled_pages: list[dict[str, Any]], -) -> dict[str, Any]: - pages = [] - for page in sampled_pages: - kind, confidence, note = _heuristic_kind(page) - pages.append( - { - "page": int(page.get("page_number") or 0), - "special_kind": kind, - "confidence": confidence, - "note": note, - } - ) - return {"pages": pages, "global_notes": "heuristic classification"} - - -def _normalize_llm_pages( - data: dict[str, Any], - sampled_pages: list[dict[str, Any]], -) -> dict[str, Any]: - by_page = {int(page.get("page_number") or 0): page for page in sampled_pages} - pages = [] - for item in data.get("pages", []) or []: - if not isinstance(item, dict): - continue - page_number = int(item.get("page") or item.get("page_number") or 0) - if page_number not in by_page: - continue - kind = str(item.get("special_kind") or item.get("kind") or "normal") - if kind not in ALLOWED_KINDS: - kind = "normal" - confidence = max(0.0, min(float(item.get("confidence") or 0.0), 1.0)) - if confidence <= 0: - confidence = 0.5 - pages.append( - { - "page": page_number, - "special_kind": kind, - "confidence": confidence, - "note": str(item.get("note") or "")[:300], - } - ) - - seen = {item["page"] for item in pages} - for fallback in heuristic_classify_special_pages(sampled_pages)["pages"]: - if fallback["page"] not in seen: - pages.append(fallback) - pages.sort(key=lambda item: item["page"]) - return { - "pages": pages, - "global_notes": str(data.get("global_notes") or data.get("notes") or "")[:1000], - } - - -def classify_special_pages( - sampled_pages: list[dict[str, Any]], - *, - vlm_observations: list[dict[str, Any]] | None = None, - model: str | None = None, - use_llm: bool = True, -) -> dict[str, Any]: - """Classify special pages with LLM, falling back to deterministic heuristics.""" - if not use_llm: - return heuristic_classify_special_pages(sampled_pages) - - try: - from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL - client = get_openai_client(model=effective_model) - payload = { - "sampled_pages": sampled_pages, - "vlm_observations": vlm_observations or [], - } - response = client.chat_completion( - messages=[ - {"role": "system", "content": PROMPT}, - { - "role": "user", - "content": "Classify these pages as json:\n" - + json.dumps(payload, ensure_ascii=False), - }, - ], - model=effective_model, - temperature=0.0, - max_tokens=1800, - response_format={"type": "json_object"}, - ) - return _normalize_llm_pages(extract_json_object(response), sampled_pages) - except Exception as exc: - logger.warning( - f"[document_agent.classify_special_pages] LLM classification failed, " - f"using heuristics: {exc}" - ) - return heuristic_classify_special_pages(sampled_pages) diff --git a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py new file mode 100644 index 000000000..4a424a5cc --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py @@ -0,0 +1,424 @@ +"""VLM-driven TOC anchor, boundary, and entry extraction.""" + +from __future__ import annotations + +import gc +import json +import os +import time +from pathlib import Path +from typing import Any, cast + +from shared.utils.token_estimate import estimate_tokens + +from app.services.document_agent.manifest import ( + TocAnchorPage, + TocResult, + ToolContext, + ToolResult, +) +from app.services.document_agent.registry import register_tool +from app.services.document_agent.tools.vlm_toc_extractor import ( + vlm_entries_to_toc_hierarchies, +) +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) +from loguru import logger + +# -- Constants ----------------------------------------------------------------- + +BOUNDARY_STEP_PAGES = 5 +MAX_BOUNDARY_ROUNDS = 6 +MAX_TOC_PAGES = BOUNDARY_STEP_PAGES * MAX_BOUNDARY_ROUNDS # 30 + + +# -- PyMuPDF workers (must be top-level for multiprocessing pickle) ------------ + + +@worker +def _render_single_page_worker( + queue, pdf_path: str, page_num: int, output_path: str, dpi: int +) -> None: + import pymupdf # type: ignore[import] + + try: + doc = pymupdf.open(pdf_path) + idx = page_num - 1 + if 0 <= idx < doc.page_count: + page = doc[idx] + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + pix = page.get_pixmap(matrix=mat) + pix.save(output_path) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "png_path": output_path}) + + +# -- VLM helpers --------------------------------------------------------------- + + +def _vlm_confirm_anchors( + anchor_pages: list[TocAnchorPage], + model: str, + budget: Any | None = None, +) -> tuple[list[TocAnchorPage], bool]: + """Phase 1: send all anchor PNGs to VLM, ask which are real TOC starts.""" + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + if not anchor_pages: + return [], False + + import base64 + + # Build multi-image message + content_parts: list[dict[str, Any]] = [ + { + "type": "text", + "text": ( + "You are a document structure analysis expert. " + "Below are screenshot(s) of candidate pages extracted from a PDF. " + "These pages contained keywords such as 'Table of Contents' / 'Contents' " + "during a text scan.\n\n" + "For each page, determine whether it is truly the **start page** of a " + "Table of Contents (TOC).\n\n" + "Criteria for a real TOC page:\n" + "- Contains a list of section titles paired with page numbers\n" + "- Titles are connected to page numbers via dots, ellipses, or spaces\n" + "- Titles have a systematic numbering scheme (e.g. 1. / 1.1 / Chapter 1)\n\n" + "NOT a TOC page:\n" + "- Body text that casually mentions 'contents'\n" + "- A page with only a 'Contents' heading but body text below\n\n" + "Return a strict JSON array (no markdown fences):\n" + '[{"page": , "is_toc_start": true/false, "reason": "brief reason"}]' + ), + } + ] + + for anchor in anchor_pages: + with open(anchor.png_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + content_parts.append( + { + "type": "text", + "text": f"\n--- Page {anchor.page} ---", + } + ) + content_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + + messages = cast(Any, [{"role": "user", "content": content_parts}]) + est = estimate_tokens(str(content_parts[0]["text"])) + len(anchor_pages) * 800 + if budget and not budget.try_reserve("visual", est): + logger.warning("[extract.toc] insufficient visual budget for anchor confirmation") + return [], True + + try: + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=messages, + model=model, + temperature=0.1, + max_tokens=500, + response_format={"type": "json_object"}, + ) + if budget: + budget.commit("visual", actual=usage.get("total_tokens", est), est=est) + data = json.loads(raw) + if isinstance(data, dict): + items = data.get("pages") or data.get("results") or data.get("data") or [] + if not items and len(data) == 1: + items = list(data.values())[0] + elif isinstance(data, list): + items = data + else: + items = [] + + confirmed_pages: set[int] = set() + for item in items: + if isinstance(item, dict) and item.get("is_toc_start"): + confirmed_pages.add(int(item["page"])) + + confirmed = [a for a in anchor_pages if a.page in confirmed_pages] + rejected = [a.page for a in anchor_pages if a.page not in confirmed_pages] + logger.info( + "[extract.toc] VLM confirmed {} TOC starts, rejected pages: {}", + len(confirmed), + rejected, + ) + return confirmed, False + except Exception as exc: + if budget: + budget.refund("visual", est=est) + logger.warning( + "[extract.toc] VLM anchor confirmation failed: {}, " + "falling back to no confirmed anchors (safe degradation)", + exc, + ) + return [], True + + +# -- Main tool ----------------------------------------------------------------- + + +@register_tool( + name="extract.toc_with_boundaries", + description=( + "VLM-confirms TOC anchor pages, then batch-classifies and extracts " + "TOC entries from rendered page windows using VLM." + ), +) +def extract_toc_with_boundaries( + ctx: ToolContext, _args: dict[str, Any] +) -> ToolResult: + start = time.monotonic() + anchors = ctx.blackboard.toc_anchor_pages + warnings: list[str] = [] + debug_info: dict[str, Any] = {} + + if not anchors: + logger.info("[extract.toc] no anchor pages, skipping") + ctx.blackboard.toc_result = TocResult( + method="none", + notes="No TOC anchor pages found by find.toc_anchor_pages", + ) + return ToolResult( + status="ok", + payload={"toc_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + ) + + model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") + if not model: + logger.warning("[extract.toc] no VLM model configured; skipping TOC extraction") + ctx.blackboard.toc_result = TocResult( + method="none", + notes="No VLM model configured for TOC extraction", + ) + return ToolResult( + status="ok", + payload={"toc_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["No VLM model configured; skipping TOC extraction."], + ) + dpi = int(ctx.settings.get("toc_png_dpi", "144")) + page_count = ctx.blackboard.page_count + output_dir = str( + Path(ctx.output_dir or os.path.expanduser("~/.knowhere/_debug_profile")) + / "toc_pages" + ) + os.makedirs(output_dir, exist_ok=True) + + # -- Phase 1: VLM confirm anchors ----------------------------------------- + confirmed, confirm_failed = _vlm_confirm_anchors(anchors, model, budget=ctx.budget) + if confirm_failed: + warnings.append("vlm_anchor_confirmation_failed") + debug_info["phase1_confirmed"] = [a.page for a in confirmed] + debug_info["phase1_rejected"] = [ + a.page for a in anchors if a not in confirmed + ] + + if not confirmed: + ctx.blackboard.toc_result = TocResult( + method="none", + notes="VLM rejected all TOC anchor candidates", + ) + return ToolResult( + status="ok", + payload={"toc_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["VLM rejected all anchor pages"], + debug=debug_info, + ) + + # -- Phase 2+3 (unified): batch classify + extract --------------------------- + # Instead of separate boundary detection (Phase 2) then per-page extraction + # (Phase 3), we send batches of BOUNDARY_STEP_PAGES images to VLM in one + # call. The VLM classifies each page (TOC vs non-TOC) AND extracts entries + # from TOC pages simultaneously. If the last page in a batch is still TOC, + # we expand the window and use prior entries as continuation context. + from app.services.document_agent.tools.vlm_toc_extractor import ( + vlm_extract_toc_batch, + ) + + all_entries: list[dict[str, Any]] = [] + all_toc_pages: list[int] = [] + toc_hierarchies: list[dict[str, Any]] = [] + batch_meta: list[dict[str, Any]] = [] + batch_trace: list[dict[str, Any]] = [] + + for anchor in confirmed: + anchor_page = anchor.page + region_entries: list[dict[str, Any]] = [] + region_toc_pages: list[int] = [] + region_scan_end = anchor_page + + for round_idx in range(MAX_BOUNDARY_ROUNDS): + batch_start = anchor_page + round_idx * BOUNDARY_STEP_PAGES + batch_end = min( + batch_start + BOUNDARY_STEP_PAGES - 1, page_count + ) + if batch_start > page_count: + break + + batch_pages = list(range(batch_start, batch_end + 1)) + logger.info( + "[extract.toc] batch round {}: pages {}-{} for anchor {}", + round_idx, batch_start, batch_end, anchor_page, + ) + + # Render all pages in this batch + page_pngs: list[tuple[int, str]] = [] + for page_num in batch_pages: + png_path = os.path.join(output_dir, f"toc_page_{page_num}.png") + run_in_child_process( + _render_single_page_worker, + ctx.pdf_path, + page_num, + png_path, + dpi, + timeout=60, + ) + page_pngs.append((page_num, png_path)) + + # Send batch to VLM — classify + extract in one call + batch_result = vlm_extract_toc_batch( + page_pngs=page_pngs, + model=model, + previous_entries=region_entries if region_entries else None, + ) + batch_meta.append(batch_result.meta) + + # Collect results + region_entries.extend(batch_result.all_entries) + region_toc_pages.extend(batch_result.toc_pages) + region_scan_end = batch_end + + batch_trace.append({ + "anchor": anchor_page, + "round": round_idx, + "batch_pages": batch_pages, + "toc_pages": batch_result.toc_pages, + "non_toc_pages": batch_result.non_toc_pages, + "entries_found": len(batch_result.all_entries), + }) + + # Determine if we need to continue expanding + # If the last page in the batch is NOT TOC, boundary found + last_page_is_toc = ( + batch_result.page_results + and batch_result.page_results[-1].is_toc + ) + if not last_page_is_toc: + logger.info( + "[extract.toc] boundary found: last page {} is not TOC", + batch_end, + ) + break + + # Last page is still TOC — continue expanding + if batch_end >= page_count: + break + logger.info( + "[extract.toc] last page {} still TOC, expanding window", + batch_end, + ) + + all_entries.extend(region_entries) + all_toc_pages.extend(region_toc_pages) + + if region_entries: + region_hierarchies = vlm_entries_to_toc_hierarchies( + region_entries, + toc_page_nums=region_toc_pages, + scan_end_page=region_scan_end, + page_count=page_count, + ) + toc_hierarchies.extend(region_hierarchies) + else: + logger.warning( + "[extract.toc] anchor {} produced no TOC entries", + anchor_page, + ) + + if not all_entries: + raise RuntimeError( + "VLM TOC extractor returned no entries for confirmed TOC pages" + ) + + debug_info["batch_trace"] = batch_trace + debug_info["batch_meta"] = batch_meta + debug_info["vlm_entry_count"] = len(all_entries) + + all_toc_pages_sorted = sorted(set(all_toc_pages)) + toc_region_count = len(toc_hierarchies) + + ctx.blackboard.toc_result = TocResult( + toc_pages=all_toc_pages_sorted, + method="vlm_batch", + notes=( + f"VLM confirmed {len(confirmed)} TOC starts, " + f"batch classify+extract found {toc_region_count} regions, " + f"toc_pages={all_toc_pages_sorted}" + ), + ) + ctx.blackboard.toc_hierarchies = toc_hierarchies if toc_hierarchies else None + ctx.blackboard.global_signals["vlm_toc_entries"] = { + "model": model, + "toc_pages": all_toc_pages_sorted, + "total_entries": len(all_entries), + "entries": all_entries, + "batch_meta": batch_meta, + } + + # Persist toc_hierarchies to disk for inspection / downstream reuse + if toc_hierarchies and ctx.output_dir: + toc_json_path = os.path.join(ctx.output_dir, "toc_hierarchies.json") + try: + with open(toc_json_path, "w", encoding="utf-8") as f: + json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) + logger.info("[extract.toc] wrote toc_hierarchies to {}", toc_json_path) + except Exception as exc: + logger.warning("[extract.toc] failed to write toc_hierarchies: {}", exc) + + # Build toc_ranges from confirmed TOC pages for summary + toc_ranges_out: list[list[int]] = [] + if toc_hierarchies: + for hier in toc_hierarchies: + toc_ranges_out.append(hier.get("toc_range", [])) + + toc_summary: dict[str, Any] = { + "toc_ranges": toc_ranges_out, + "toc_page_count": len(all_toc_pages_sorted), + "toc_entry_count": len(all_entries), + "toc_region_count": toc_region_count, + "toc_source": "vlm_batch", + } + if toc_hierarchies: + for i, hier in enumerate(toc_hierarchies): + tree = hier.get("toc_tree", {}) + toc_summary[f"region_{i}_level1_count"] = len(tree) + toc_summary[f"region_{i}_level1_titles"] = list(tree.keys())[:10] + + return ToolResult( + status="ok", + payload={ + "toc_count": len(toc_hierarchies) if toc_hierarchies else 0, + "toc_page_count": len(all_toc_pages_sorted), + "toc_region_count": toc_region_count, + }, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary=toc_summary, + warnings=warnings, + debug=debug_info, + ) + diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py new file mode 100644 index 000000000..db03c90eb --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -0,0 +1,237 @@ +"""Scan for TOC anchor pages and render their PNGs for VLM inspection.""" + +from __future__ import annotations + +import gc +import os +import time +from pathlib import Path +from typing import Any + +from app.services.document_agent.manifest import TocAnchorPage, ToolContext, ToolResult +from app.services.document_agent.registry import has_page_labels, register_tool +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) +from loguru import logger + +# CJK and English TOC keywords used for anchor detection. +TOC_KEYWORDS = {"目录", "目次", "contents", "tableofcontents", "table of contents"} + +# If a TOC keyword fingerprint appears on more than this fraction of total +# pages, it is treated as a recurring navigation element (header/footer link) +# rather than real TOC content. +RECURRING_ELEMENT_THRESHOLD = 0.30 + +# Hard cap on the number of candidate anchor pages sent to VLM. A real +# document never has more than ~30 TOC start pages. +MAX_ANCHOR_CANDIDATES = 30 + + +def _normalize_for_toc(text: str) -> str: + """Collapse whitespace for keyword matching.""" + return text.replace(" ", "").replace("\u3000", "").lower() + + +@worker +def _render_pages_worker( + queue, pdf_path: str, pages: list[int], output_dir: str, dpi: int +) -> None: + import pymupdf # type: ignore[import] + + results: list[dict[str, Any]] = [] + try: + doc = pymupdf.open(pdf_path) + for page_num in pages: + idx = page_num - 1 + if 0 <= idx < doc.page_count: + page = doc[idx] + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + pix = page.get_pixmap(matrix=mat) + png_name = f"toc_anchor_page_{page_num}.png" + png_path = os.path.join(output_dir, png_name) + pix.save(png_path) + results.append({"page": page_num, "png_path": png_path}) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "results": results}) + + +def _filter_recurring_elements( + matches: list[tuple[int, str, int]], + total_pages: int, +) -> set[int]: + """Remove pages whose TOC keyword pattern is a recurring navigation element. + + Each *match* is ``(page, raw_line, line_index)``. We build a **composite + fingerprint** per page by joining all its matches as + ``"raw_line@line_idx\\n..."``. If the same composite fingerprint appears + on more than ``RECURRING_ELEMENT_THRESHOLD`` of all pages, those pages are + header/footer false-positives. + + Example (SpaceX S-1): + - page 17 → ``"Table of Contents@0\\nTABLE OF CONTENTS@1"`` (unique → keeps) + - page 18 → ``"Table of Contents@1"`` (376 pages → recurring → filtered) + - page 401 → ``"Table of Contents@0"`` (~31 pages → VLM decides) + """ + # Collect all matches per page + page_matches: dict[int, list[tuple[str, int]]] = {} + for page, raw_line, line_idx in matches: + page_matches.setdefault(page, []).append((raw_line, line_idx)) + + # Build composite fingerprint per page (sorted by line_idx for stability) + page_fingerprints: dict[int, str] = {} + for page, hits in page_matches.items(): + hits_sorted = sorted(hits, key=lambda h: h[1]) + page_fingerprints[page] = "\n".join( + f"{raw}@{idx}" for raw, idx in hits_sorted + ) + + # Group pages by composite fingerprint + fp_groups: dict[str, list[int]] = {} + for page, fp in page_fingerprints.items(): + fp_groups.setdefault(fp, []).append(page) + + threshold = max(int(total_pages * RECURRING_ELEMENT_THRESHOLD), 1) + + surviving: set[int] = set() + for fp, pages in fp_groups.items(): + if len(pages) > threshold: + logger.info( + "[find.toc_anchor_pages] recurring pattern filtered: " + "{!r} appears on {}/{} pages", + fp[:60], + len(pages), + total_pages, + ) + else: + surviving.update(pages) + + return surviving + + +@register_tool( + name="find.toc_anchor_pages", + description=( + "Scan page text previews for TOC keywords, filter recurring " + "navigation elements, then render candidate PNGs for VLM confirmation." + ), + preconditions=(has_page_labels,), +) +def find_toc_anchor_pages(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + total_pages = ctx.blackboard.page_count + + # Scan text previews for TOC keywords and record per-line matches. + # Each entry: (page, raw_line_text, line_index) + # We use raw (original) text for fingerprinting so that casing + # differences (e.g. "Table of Contents" vs "TABLE OF CONTENTS") + # naturally produce distinct composite fingerprints. + keyword_matches: list[tuple[int, str, int]] = [] + raw_hit_pages: set[int] = set() + + for feature in ctx.blackboard.page_features: + page_matched = False + for line_idx, raw_line in enumerate(feature.text_lines_preview): + norm_line = _normalize_for_toc(raw_line) + for keyword in TOC_KEYWORDS: + if keyword in norm_line: + keyword_matches.append((feature.page, raw_line.strip(), line_idx)) + raw_hit_pages.add(feature.page) + page_matched = True + break # one match per line is enough + + # Fallback: check if a TOC keyword spans across adjacent lines. + # PyMuPDF sometimes splits large headings across lines, e.g. + # "目" + "录" or "Table of" + "Contents". Join the first few + # preview lines (where a page title would appear) and re-check + # with the same keywords and normalisation. + if not page_matched and feature.text_lines_preview: + head = feature.text_lines_preview[:10] + joined_head = _normalize_for_toc("".join(head)) + for keyword in TOC_KEYWORDS: + if keyword in joined_head: + keyword_matches.append((feature.page, keyword, 0)) + raw_hit_pages.add(feature.page) + logger.debug( + "[find.toc_anchor_pages] cross-line keyword '{}' " + "detected on page {} (head lines joined)", + keyword, + feature.page, + ) + break + + # Apply recurring element fingerprint filter + if keyword_matches: + anchor_pages = _filter_recurring_elements(keyword_matches, total_pages) + else: + anchor_pages = set() + + logger.info( + "[find.toc_anchor_pages] keyword scan: {} raw hits → {} after " + "fingerprint filter", + len(raw_hit_pages), + len(anchor_pages), + ) + + # Hard cap: a document never has more than ~30 real TOC start candidates. + if len(anchor_pages) > MAX_ANCHOR_CANDIDATES: + logger.warning( + "[find.toc_anchor_pages] {} candidates exceed cap of {}, truncating", + len(anchor_pages), + MAX_ANCHOR_CANDIDATES, + ) + anchor_pages = set(sorted(anchor_pages)[:MAX_ANCHOR_CANDIDATES]) + + if not anchor_pages: + logger.info("[find.toc_anchor_pages] no TOC keyword pages found") + ctx.blackboard.toc_anchor_pages = [] + return ToolResult( + status="ok", + payload={"anchor_count": 0}, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={"anchor_count": 0, "pages": []}, + ) + + # Render candidate pages as PNGs for downstream VLM confirmation + sorted_pages = sorted(anchor_pages) + output_dir = str( + Path(ctx.output_dir or os.path.expanduser("~/.knowhere/_debug_profile")) + / "toc_pages" + ) + os.makedirs(output_dir, exist_ok=True) + + dpi = int(ctx.settings.get("toc_png_dpi", "144")) + result = run_in_child_process( + _render_pages_worker, ctx.pdf_path, sorted_pages, output_dir, dpi, timeout=120 + ) + + anchors: list[TocAnchorPage] = [] + for item in result.get("results") or []: + page = int(item["page"]) + anchors.append( + TocAnchorPage(page=page, png_path=item["png_path"], source="text_scan") + ) + + ctx.blackboard.toc_anchor_pages = anchors + logger.info( + "[find.toc_anchor_pages] found {} anchor pages: {}", + len(anchors), + [a.page for a in anchors], + ) + + return ToolResult( + status="ok", + payload={"anchor_count": len(anchors)}, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={ + "anchor_count": len(anchors), + "pages": [a.to_dict() for a in anchors], + }, + ) + diff --git a/apps/worker/app/services/document_agent/tools/grep_text.py b/apps/worker/app/services/document_agent/tools/grep_text.py new file mode 100644 index 000000000..6c59f5268 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/grep_text.py @@ -0,0 +1,79 @@ +"""Generic full-document text grep for native PDFs.""" + +from __future__ import annotations + +import re +import time +from typing import Any + +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.pdf_text import read_page_texts +from app.services.document_agent.registry import has_page_features, not_is_scanned, register_tool + + +def _load_page_texts(ctx: ToolContext) -> dict[int, str]: + if ctx.blackboard.page_full_text_cache: + return dict(ctx.blackboard.page_full_text_cache) + pages = list(range(1, ctx.blackboard.page_count + 1)) + texts = read_page_texts(ctx.pdf_path, pages, timeout=300) + ctx.blackboard.page_full_text_cache = texts + return texts + + +@register_tool( + name="grep.text", + description="Search full PDF text for a substring or regex. Available only for native PDFs.", + parameters={ + "type": "object", + "properties": { + "query": {"type": "string"}, + "regex": {"type": "boolean", "default": False}, + "case_sensitive": {"type": "boolean", "default": False}, + "max_results": {"type": "integer", "default": 30}, + "context_chars": {"type": "integer", "default": 80}, + }, + "required": ["query"], + }, + preconditions=(has_page_features, not_is_scanned), +) +def grep_text(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + query = str(args.get("query") or "").strip() + if not query: + return ToolResult( + status="error", + error="grep.text requires query", + latency_ms=int((time.monotonic() - start) * 1000), + ) + use_regex = bool(args.get("regex", False)) + case_sensitive = bool(args.get("case_sensitive", False)) + max_results = max(1, min(int(args.get("max_results") or 30), 100)) + context_chars = max(20, min(int(args.get("context_chars") or 80), 300)) + flags = 0 if case_sensitive else re.IGNORECASE + pattern = re.compile(query if use_regex else re.escape(query), flags) + results: list[dict[str, Any]] = [] + for page, text in sorted(_load_page_texts(ctx).items()): + for match in pattern.finditer(text): + start_idx = max(match.start() - context_chars, 0) + end_idx = min(match.end() + context_chars, len(text)) + results.append( + { + "page": page, + "char_offset": match.start(), + "snippet": text[start_idx:end_idx].replace("\n", " "), + } + ) + if len(results) >= max_results: + break + if len(results) >= max_results: + break + summary = {"query": query, "hit_count": len(results), "results": results} + ctx.blackboard.global_signals.setdefault("grep_history", []).append( + {"query": query, "hit_count": len(results), "sample_pages": [item["page"] for item in results[:10]]} + ) + return ToolResult( + status="ok", + payload=summary, + latency_ms=int((time.monotonic() - start) * 1000), + output_summary={"query": query, "hit_count": len(results)}, + ) diff --git a/apps/worker/app/services/document_agent/tools/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py new file mode 100644 index 000000000..3b31d97e0 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -0,0 +1,108 @@ +"""Generic VLM inspection tool for selected PDF pages.""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from typing import Any, cast + +from app.services.document_agent.manifest import ToolContext, ToolResult +from app.services.document_agent.registry import has_page_features, register_tool +from app.services.document_agent.visual import render_pages +from shared.utils.token_estimate import estimate_tokens + + +@register_tool( + name="inspect.pages", + description="Render arbitrary PDF pages and ask the VLM a custom profiling question.", + parameters={ + "type": "object", + "properties": { + "pages": {"type": "array", "items": {"type": "integer"}}, + "question": {"type": "string"}, + }, + "required": ["pages", "question"], + }, + preconditions=(has_page_features,), +) +def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + pages = sorted( + { + int(page) + for page in (args.get("pages") or []) + if 1 <= int(page) <= ctx.blackboard.page_count + } + )[:10] + if not pages: + return ToolResult( + status="error", + error="inspect.pages requires at least one valid page", + latency_ms=int((time.monotonic() - start) * 1000), + ) + question = str(args.get("question") or "Describe the document structure visible on these pages.") + pngs = render_pages(ctx, pages, folder_name="inspect_pages", prefix="inspect") + model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") + prompt = ( + "You are inspecting PDF page screenshots for a document profiling agent. " + "Answer strict JSON with keys: observations, implications, recommended_next_action. " + "observations must be an array of {page, summary, visual_kind}. " + f"Question: {question}" + ) + est = estimate_tokens(prompt) + len(pngs) * 800 + if not model: + payload = {"pages": pages, "pngs": pngs, "note": "No VLM model configured."} + ctx.blackboard.global_signals.setdefault("visual_inspections", []).append(payload) + return ToolResult( + status="ok", + payload=payload, + latency_ms=int((time.monotonic() - start) * 1000), + warnings=["No VLM model configured; returned rendered page paths only."], + ) + if not ctx.budget.try_reserve("visual", est): + return ToolResult( + status="error", + error="insufficient visual budget", + latency_ms=int((time.monotonic() - start) * 1000), + ) + + content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + for item in pngs: + with open(str(item["png_path"]), "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + content_parts.append({"type": "text", "text": f"\n--- Page {item['page']} ---"}) + content_parts.append( + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}} + ) + try: + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + client = get_openai_client(model=model) + raw, usage = client.chat_completion_with_usage( + messages=cast(Any, [{"role": "user", "content": content_parts}]), + model=model, + temperature=0.0, + max_tokens=1200, + response_format={"type": "json_object"}, + ) + ctx.budget.commit("visual", actual=usage.get("total_tokens", est), est=est) + try: + payload: dict[str, Any] = json.loads(raw) + except json.JSONDecodeError: + payload = {"raw": raw} + if isinstance(payload, dict): + payload.setdefault("pages", pages) + else: + payload = {"result": payload, "pages": pages} + ctx.blackboard.global_signals.setdefault("visual_inspections", []).append(payload) + return ToolResult( + status="ok", + payload=payload, + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=usage.get("total_tokens", 0), + ) + except Exception: + ctx.budget.refund("visual", est=est) + raise diff --git a/apps/worker/app/services/document_agent/tools/llm_json.py b/apps/worker/app/services/document_agent/tools/llm_json.py deleted file mode 100644 index 9531e3690..000000000 --- a/apps/worker/app/services/document_agent/tools/llm_json.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Small JSON helpers for split-agent LLM tools.""" - -from __future__ import annotations - -import json -import re -from typing import Any - - -def extract_json_object(text: str) -> dict[str, Any]: - """Parse JSON from a model response that may include light prose/fences.""" - raw = (text or "").strip() - if not raw: - raise ValueError("empty JSON response") - if raw.startswith("```"): - raw = re.sub(r"^```(?:json)?\s*", "", raw, flags=re.IGNORECASE) - raw = re.sub(r"\s*```$", "", raw) - try: - data = json.loads(raw) - except json.JSONDecodeError: - start = raw.find("{") - end = raw.rfind("}") - if start < 0 or end <= start: - raise - data = json.loads(raw[start : end + 1]) - if not isinstance(data, dict): - raise ValueError("expected JSON object") - return data 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 new file mode 100644 index 000000000..a60ad2695 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/match_h1_pages.py @@ -0,0 +1,195 @@ +"""Match TOC level-1 headings to body pages via PyMuPDF text search.""" + +from __future__ import annotations + +import re +import time +import unicodedata +from typing import Any + +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 loguru import logger + + +# ── Text normalization for matching ────────────────────────────────────── + +_LEADING_NUMBER_RE = re.compile( + r"""^ + (?: + [#]+\s* + | 第\s*[零一二三四五六七八九十百千\d]+\s*[章节篇部分] + | [零一二三四五六七八九十百千]+\s*[、。，,] + | [（(]\s*[零一二三四五六七八九十百千\d]+\s*[）)] + | \d+(?:\.\d+)*\.?\s* + | [IVXLCDM]+\.?\s+ + | [A-Za-z]\.\s+ + | Chapter\s+\w+\s* + ) + """, + re.VERBOSE | re.IGNORECASE, +) + +_PAGE_SUFFIX_RE = re.compile(r"[\s\.\-·…]+\d+\s*$") + + +def _normalize(text: str) -> str: + """Normalize text for fuzzy heading matching.""" + text = unicodedata.normalize("NFKC", text or "") + text = re.sub(r"\s+", " ", text).strip() + return text + + +def _clean_toc_title(title: str) -> str: + """Remove leading numbering/hashes and trailing page numbers from a TOC title.""" + cleaned = _PAGE_SUFFIX_RE.sub("", title or "").strip() + cleaned = _LEADING_NUMBER_RE.sub("", cleaned).strip() + return cleaned + + +def _extract_level1_titles(toc_hierarchies: list[dict[str, Any]]) -> list[str]: + """Extract level-1 titles from toc_hierarchies. + + Each hierarchy dict contains ``toc_tree`` – a nested dict where top-level + keys are level-1 headings (values are sub-heading dicts). + """ + titles: list[str] = [] + for hier in toc_hierarchies: + toc_tree = hier.get("toc_tree") or {} + for raw_title in toc_tree.keys(): + cleaned = _clean_toc_title(raw_title) + if cleaned and len(cleaned) >= 2: + titles.append(cleaned) + return titles + + +@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) + + # Strict substring matching: for each level-1 title, find the first body page + h1_candidates: list[H1Candidate] = [] + matched_titles: list[str] = [] + unmatched_titles: list[str] = [] + + for title in level1_titles: + normalized_title = _normalize(title) + found = False + for page in search_pages: + text = page_texts.get(page, "") + normalized_text = _normalize(text) + if normalized_title in normalized_text: + # Find the matched line for evidence + matched_line = "" + for line in text.splitlines(): + if normalized_title in _normalize(line): + matched_line = line.strip()[:100] + break + + h1_candidates.append( + H1Candidate( + title=title, + page=page, + confidence=0.88, + matched_line=matched_line, + source="toc_exact_top", + evidence={ + "normalized_needle": normalized_title, + "page_text_length": len(text), + }, + ) + ) + matched_titles.append(title) + found = True + break # Only first match per title + + if not found: + unmatched_titles.append(title) + + # Deduplicate: if multiple titles map to the same page, keep the first + seen_pages: set[int] = set() + deduped: list[H1Candidate] = [] + for candidate in h1_candidates: + if candidate.page not in seen_pages: + seen_pages.add(candidate.page) + deduped.append(candidate) + h1_candidates = deduped + + 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 new file mode 100644 index 000000000..7ae7f0195 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/persist_anatomy_map.py @@ -0,0 +1,63 @@ +"""Persist anatomy map artifacts and optional database records.""" + +from __future__ import annotations + +import json +import os +import time +from pathlib import Path +from typing import Any + +from app.services.document_agent.manifest import PageAnatomyMap, ToolContext, ToolResult + + +def _artifact_dir(ctx: ToolContext) -> Path: + if ctx.output_dir: + return Path(ctx.output_dir) + base = Path(os.path.expanduser("~/.knowhere/_debug_profile")) + return base / Path(ctx.pdf_path).stem + + +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") + return PageAnatomyMap( + job_id=ctx.job_id, + file_path=ctx.pdf_path, + page_count=ctx.blackboard.page_count, + page_features=ctx.blackboard.page_features, + page_labels=ctx.blackboard.page_labels, + toc_result=ctx.blackboard.toc_result, + h1_result=ctx.blackboard.h1_result, + shard_plan=ctx.blackboard.shard_plan, + document_profile=ctx.blackboard.document_profile, + toc_hierarchies=ctx.blackboard.toc_hierarchies, + global_signals=ctx.blackboard.global_signals, + trace_summary={ + "budget": ctx.budget.snapshot(), + "validation": ctx.blackboard.validation_report, + }, + ) + + +def persist_anatomy_map(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + anatomy = build_anatomy_map(ctx) + output_dir = _artifact_dir(ctx) + output_dir.mkdir(parents=True, exist_ok=True) + artifact_path = output_dir / "anatomy_map.json" + artifact_path.write_text( + json.dumps(anatomy.to_dict(), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + if ctx.trace: + ctx.trace.set_anatomy_map(anatomy, str(artifact_path)) + return ToolResult( + status="ok", + payload={"artifact_path": str(artifact_path)}, + latency_ms=int((time.monotonic() - start) * 1000), + ) diff --git a/apps/worker/app/services/document_agent/tools/probe_page_features.py b/apps/worker/app/services/document_agent/tools/probe_page_features.py new file mode 100644 index 000000000..678557fad --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/probe_page_features.py @@ -0,0 +1,140 @@ +"""Full-page structural probing.""" + +from __future__ import annotations + +import gc +import time +from typing import Any + +from app.services.document_agent.manifest import PageFeature, ToolContext, ToolResult +from app.services.document_agent.pdf_text import top_lines +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) +from loguru import logger + + +def _rect_area(rect: Any) -> float: + width = max(float(getattr(rect, "width", 0.0) or 0.0), 0.0) + height = max(float(getattr(rect, "height", 0.0) or 0.0), 0.0) + return width * height + + +def _image_coverage(page: Any, page_area: float) -> tuple[float, int]: + if page_area <= 0: + return 0.0, 0 + area = 0.0 + images = page.get_images(full=True) or [] + seen: set[tuple[float, float, float, float]] = set() + for image in images: + xref = image[0] + try: + rects = page.get_image_rects(xref) or [] + except Exception: + rects = [] + for rect in rects: + key = ( + round(float(getattr(rect, "x0", 0.0) or 0.0), 2), + round(float(getattr(rect, "y0", 0.0) or 0.0), 2), + round(float(getattr(rect, "x1", 0.0) or 0.0), 2), + round(float(getattr(rect, "y1", 0.0) or 0.0), 2), + ) + if key in seen: + continue + seen.add(key) + area += _rect_area(rect) + return min(area / page_area, 1.0), len(images) + + +def _table_count(page: Any) -> int: + try: + finder = page.find_tables() + return len(getattr(finder, "tables", []) or []) + except Exception: + return 0 + + +def _probe_one(page: Any, page_number: int) -> dict[str, Any]: + rect = page.rect + area = max(_rect_area(rect), 1.0) + text = page.get_text() or "" + raw_text_length = len(text.strip()) + image_coverage, image_count = _image_coverage(page, area) + try: + drawings_count = len(page.get_drawings() or []) + except Exception: + drawings_count = 0 + orientation = "landscape" if float(rect.width) > float(rect.height) else "portrait" + return { + "page": page_number, + "raw_text_length": raw_text_length, + "text_density": round(raw_text_length / area * 10000, 4), + "image_coverage": round(image_coverage, 4), + "image_count": image_count, + "table_count": _table_count(page), + "drawings_count": drawings_count, + "orientation": orientation, + "width": round(float(rect.width), 2), + "height": round(float(rect.height), 2), + "is_blank_like": raw_text_length < 20 and image_coverage < 0.02 and drawings_count < 5, + "text_lines_preview": top_lines(text, max_lines=30), + } + + +@worker +def _probe_worker(queue, pdf_path: str) -> None: + import pymupdf # type: ignore[import] + + features: list[dict[str, Any]] = [] + page_count = 0 + try: + doc = pymupdf.open(pdf_path) + page_count = int(doc.page_count) + for idx in range(page_count): + features.append(_probe_one(doc[idx], idx + 1)) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "page_count": page_count, "features": features}) + + +def probe_page_features(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + try: + result = run_in_child_process(_probe_worker, ctx.pdf_path, timeout=300) + features = [ + PageFeature( + page=int(item["page"]), + raw_text_length=int(item.get("raw_text_length") or 0), + text_density=float(item.get("text_density") or 0.0), + image_coverage=float(item.get("image_coverage") or 0.0), + image_count=int(item.get("image_count") or 0), + table_count=int(item.get("table_count") or 0), + drawings_count=int(item.get("drawings_count") or 0), + orientation=str(item.get("orientation") or "portrait"), # type: ignore[arg-type] + width=float(item.get("width") or 0.0), + height=float(item.get("height") or 0.0), + is_blank_like=bool(item.get("is_blank_like")), + text_lines_preview=list(item.get("text_lines_preview") or []), + ) + for item in (result.get("features") or []) + ] + ctx.blackboard.page_features = sorted(features, key=lambda f: f.page) + ctx.blackboard.page_count = int(result.get("page_count") or len(features)) + ctx.blackboard.global_signals["total_pages"] = ctx.blackboard.page_count + logger.info("[document_agent] probed {} pages", ctx.blackboard.page_count) + return ToolResult( + status="ok", + payload={"page_count": ctx.blackboard.page_count}, + latency_ms=int((time.monotonic() - start) * 1000), + ) + except Exception as exc: + return ToolResult( + status="error", + error=str(exc), + latency_ms=int((time.monotonic() - start) * 1000), + ) diff --git a/apps/worker/app/services/document_agent/tools/probe_sample_pages.py b/apps/worker/app/services/document_agent/tools/probe_sample_pages.py deleted file mode 100644 index b5b7e1122..000000000 --- a/apps/worker/app/services/document_agent/tools/probe_sample_pages.py +++ /dev/null @@ -1,213 +0,0 @@ -"""Feature-page sampling for the Phase 1 split agent.""" - -from __future__ import annotations - -import gc -import statistics -from typing import Any, Literal - -from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker - -SampleStrategy = Literal["stratified", "uniform", "key_pages"] - - -def _choose_sample_indices( - page_count: int, - *, - strategy: SampleStrategy = "stratified", - max_samples: int = 25, -) -> list[int]: - if page_count <= 0: - return [] - if max_samples <= 0: - return [] - if page_count <= max_samples: - return list(range(page_count)) - - if strategy == "key_pages": - candidates = [0, 1, 2, 3, 4, page_count - 5, page_count - 4, page_count - 3, page_count - 2, page_count - 1] - return sorted({idx for idx in candidates if 0 <= idx < page_count})[:max_samples] - - if strategy == "uniform": - if max_samples == 1: - return [0] - return sorted( - { - round(i * (page_count - 1) / (max_samples - 1)) - for i in range(max_samples) - } - ) - - edge_each = min(5, max_samples // 3) - edge_indices = list(range(edge_each)) + list(range(page_count - edge_each, page_count)) - remaining = max_samples - len(set(edge_indices)) - middle_start = edge_each - middle_end = page_count - edge_each - 1 - middle_indices: list[int] = [] - if remaining > 0 and middle_start <= middle_end: - if remaining == 1: - middle_indices = [(middle_start + middle_end) // 2] - else: - middle_indices = [ - round(middle_start + i * (middle_end - middle_start) / (remaining - 1)) - for i in range(remaining) - ] - return sorted({idx for idx in edge_indices + middle_indices if 0 <= idx < page_count}) - - -def _rect_area(rect: Any) -> float: - return max(float(getattr(rect, "width", 0.0) or 0.0), 0.0) * max( - float(getattr(rect, "height", 0.0) or 0.0), - 0.0, - ) - - -def _measure_image_coverage(page: Any, page_area: float) -> tuple[float, int]: - if page_area <= 0: - return 0.0, 0 - image_area = 0.0 - images = page.get_images(full=True) or [] - seen_rects: set[tuple[float, float, float, float]] = set() - for image in images: - if not image: - continue - xref = image[0] - try: - rects = page.get_image_rects(xref) or [] - except Exception: - rects = [] - for rect in rects: - key = ( - round(float(getattr(rect, "x0", 0.0) or 0.0), 2), - round(float(getattr(rect, "y0", 0.0) or 0.0), 2), - round(float(getattr(rect, "x1", 0.0) or 0.0), 2), - round(float(getattr(rect, "y1", 0.0) or 0.0), 2), - ) - if key in seen_rects: - continue - seen_rects.add(key) - image_area += _rect_area(rect) - return min(image_area / page_area, 1.0), len(images) - - -def _font_stats(page: Any) -> dict[str, float | int]: - sizes: list[float] = [] - try: - text_dict = page.get_text("dict") or {} - except Exception: - text_dict = {} - for block in text_dict.get("blocks", []) or []: - for line in block.get("lines", []) or []: - for span in line.get("spans", []) or []: - size = float(span.get("size") or 0.0) - if size > 0: - sizes.append(size) - if not sizes: - return {"min": 0.0, "max": 0.0, "mean": 0.0, "median": 0.0, "count": 0} - return { - "min": min(sizes), - "max": max(sizes), - "mean": statistics.fmean(sizes), - "median": statistics.median(sizes), - "count": len(sizes), - } - - -def _table_count(page: Any) -> int: - try: - finder = page.find_tables() - return len(getattr(finder, "tables", []) or []) - except Exception: - return 0 - - -def _extract_page_features(page: Any, page_index: int) -> dict[str, Any]: - rect = page.rect - page_area = max(_rect_area(rect), 1.0) - text = page.get_text() or "" - text_len = len(text.strip()) - image_coverage, image_count = _measure_image_coverage(page, page_area) - try: - drawings_count = len(page.get_drawings() or []) - except Exception: - drawings_count = 0 - orientation = "landscape" if float(rect.width) > float(rect.height) else "portrait" - text_density = text_len / page_area * 10000 - table_count = _table_count(page) - is_blank_like = text_len < 20 and image_coverage < 0.02 and drawings_count < 5 - - return { - "page_index": page_index, - "page_number": page_index + 1, - "width": float(rect.width), - "height": float(rect.height), - "orientation": orientation, - "text_length": text_len, - "text_density": round(text_density, 4), - "image_count": image_count, - "image_coverage": round(image_coverage, 4), - "table_count": table_count, - "drawings_count": drawings_count, - "font_size_stats": _font_stats(page), - "is_blank_like": is_blank_like, - "text_preview": " ".join(text.split())[:500], - } - - -@worker -def _sample_pages_worker( - queue, - pdf_path: str, - strategy: str, - max_samples: int, -) -> None: - import pymupdf - - doc = pymupdf.open(pdf_path) - try: - page_count = int(doc.page_count) - safe_strategy: SampleStrategy = "stratified" - if strategy == "uniform": - safe_strategy = "uniform" - elif strategy == "key_pages": - safe_strategy = "key_pages" - indices = _choose_sample_indices( - page_count, - strategy=safe_strategy, - max_samples=max_samples, - ) - sampled_pages = [_extract_page_features(doc[idx], idx) for idx in indices] - finally: - doc.close() - gc.collect() - - queue.put( - { - "ok": True, - "page_count": page_count, - "sample_indices": indices, - "sampled_pages": sampled_pages, - } - ) - - -def sample_pages( - pdf_path: str, - *, - strategy: SampleStrategy = "stratified", - max_samples: int = 25, - timeout: int = 120, -) -> dict[str, Any]: - """Sample structural page features in an isolated PyMuPDF child process.""" - result = run_in_child_process( - _sample_pages_worker, - pdf_path, - strategy, - max_samples, - timeout=timeout, - ) - return { - "page_count": int(result.get("page_count") or 0), - "sample_indices": list(result.get("sample_indices") or []), - "sampled_pages": list(result.get("sampled_pages") or []), - } diff --git a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py deleted file mode 100644 index b073236a5..000000000 --- a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Selective VLM inspection for ambiguous sampled PDF pages.""" - -from __future__ import annotations - -import base64 -import os -import tempfile -from typing import Any - -from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker -from loguru import logger -from openai.types.chat import ( - ChatCompletionContentPartImageParam, - ChatCompletionContentPartParam, - ChatCompletionContentPartTextParam, - ChatCompletionMessageParam, -) - -DEFAULT_QUESTION = ( - "Inspect these PDF page screenshots. For each page, decide whether it is a " - "table-heavy page, image-heavy page, table of contents, blank/sparse page, " - "landscape page, single-image page, or normal page. Return compact JSON with " - "items: [{page, judgement, confidence, note}]." -) - - -@worker -def _render_vlm_pages_worker( - queue, - pdf_path: str, - page_indices: list[int], - dpi: int, - out_dir: str, -) -> None: - import pymupdf - - doc = pymupdf.open(pdf_path) - rendered: list[dict[str, Any]] = [] - try: - mat = pymupdf.Matrix(dpi / 72, dpi / 72) - for idx in page_indices: - if idx < 0 or idx >= doc.page_count: - continue - page = doc[idx] - pix = page.get_pixmap(matrix=mat, alpha=False) - out_path = os.path.join(out_dir, f"vlm_probe_p{idx + 1}.png") - pix.save(out_path) - rendered.append({"page_index": idx, "page_number": idx + 1, "path": out_path}) - pix = None - page = None - finally: - doc.close() - queue.put({"ok": True, "rendered": rendered}) - - -def _png_to_data_url(path: str) -> str | None: - try: - with open(path, "rb") as file: - data = base64.b64encode(file.read()).decode("utf-8") - return f"data:image/png;base64,{data}" - except Exception as exc: - logger.warning(f"[document_agent.vlm_inspect] failed to encode {path}: {exc}") - return None - - -def _call_vlm( - *, - image_items: list[dict[str, Any]], - question: str, - model: str | None = None, - max_tokens: int = 900, -) -> tuple[str, dict[str, int]]: - from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - effective_model = model or settings.IMAGE_MODEL or "qwen3.5-flash" - client = get_openai_client(model=effective_model) - content: list[ChatCompletionContentPartParam] = [ - ChatCompletionContentPartTextParam( - type="text", - text=question, - ) - ] - - for item in image_items: - url = item.get("data_url") - if not url: - continue - content.append( - ChatCompletionContentPartTextParam( - type="text", - text=f"Page {item['page_number']}:", - ) - ) - content.append( - ChatCompletionContentPartImageParam( - type="image_url", - image_url={"url": url}, - ) - ) - - messages: list[ChatCompletionMessageParam] = [{"role": "user", "content": content}] - return client.chat_completion_with_usage( - messages=messages, - model=effective_model, - temperature=0.0, - max_tokens=max_tokens, - ) - - -def vlm_inspect_pages( - pdf_path: str, - *, - page_indices: list[int], - question: str = DEFAULT_QUESTION, - dpi: int = 120, - model: str | None = None, - max_tokens: int = 900, - timeout: int = 60, -) -> dict[str, Any]: - """Render selected 0-based pages and ask the configured VLM to inspect them.""" - if not page_indices: - return {"observations": [], "raw_response": "", "usage": {}} - - with tempfile.TemporaryDirectory(prefix="doc_agent_vlm_") as tmp_dir: - result = run_in_child_process( - _render_vlm_pages_worker, - pdf_path, - sorted(set(page_indices)), - dpi, - tmp_dir, - timeout=timeout, - ) - image_items = [] - for item in result.get("rendered", []) or []: - data_url = _png_to_data_url(item["path"]) - if data_url is not None: - image_items.append({**item, "data_url": data_url}) - - if not image_items: - return {"observations": [], "raw_response": "", "usage": {}} - - response, usage = _call_vlm( - image_items=image_items, - question=question, - model=model, - max_tokens=max_tokens, - ) - return { - "observations": [ - { - "page": item["page_number"], - "page_index": item["page_index"], - "vlm_judgement": response, - "confidence": None, - } - for item in image_items - ], - "raw_response": response, - "usage": usage, - } 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 1db72592b..82be670f6 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 @@ -1,361 +1,365 @@ -"""Shard planning for the Phase 1 split agent.""" +"""LLM-guided long-PDF shard planning from document profile evidence.""" from __future__ import annotations import json -from collections import Counter -from hashlib import sha256 +import os +import time from typing import Any from app.services.document_agent.manifest import ( - GlobalSignals, - ShardManifest, - ShardSignal, - SpecialPage, + Shard, + ShardPlan, + ToolContext, + ToolResult, ) -from app.services.document_agent.tools.llm_json import extract_json_object -from loguru import logger +from app.services.document_agent.registry import has_doc_stats, has_h1_result, has_toc_result, register_tool +from app.services.document_agent.validators import single_shard_plan, validate_shard_plan +from shared.utils.token_estimate import estimate_tokens -PROMPT = """You are planning PDF shards for a downstream parser. -Return ONLY valid json. -Goal: -- Cover every page from 1 to page_count exactly once. -- Prefer shards as close to max_pages_per_shard pages as possible without exceeding it. -- Do not cut through obvious table-heavy ranges, continuous image/landscape blocks, or likely TOC pages. -- Align cuts near safer normal/sparse pages when possible. - -JSON schema: -{ - "cuts": [ - {"start": 1, "end": 199, "predominant_kind": "text_dense", "rationale": "short reason"} - ], - "global_notes": "short summary" -} - -Allowed predominant_kind values: text_dense, table_heavy, image_heavy, mixed, landscape_block, toc, sparse. -""" - -ALLOWED_PREDOMINANT = { - "text_dense", - "table_heavy", - "image_heavy", - "mixed", - "landscape_block", - "toc", - "sparse", -} +def _thresholds(ctx: ToolContext) -> tuple[int, int, int]: + threshold = int( + ctx.settings.get("shard_threshold") + or os.environ.get("PARSE_AGENT_SHARD_THRESHOLD", "200") + ) + min_pages = int( + ctx.settings.get("min_pages_per_shard") + or os.environ.get("PARSE_AGENT_MIN_PAGES_PER_SHARD", "20") + ) + max_pages = int( + ctx.settings.get("max_pages_per_shard") + or os.environ.get("PARSE_AGENT_MAX_PAGES_PER_SHARD", "200") + ) + return threshold, min_pages, max_pages -def _special_by_page(classifications: dict[str, Any]) -> dict[int, dict[str, Any]]: - by_page = {} - for item in classifications.get("pages", []) or []: - if not isinstance(item, dict): +def _cuts_to_shards(cuts: list[tuple[int, str, str, float]], page_count: int) -> list[Shard]: + shards: list[Shard] = [] + previous = 0 + for cut_page, anchor_type, evidence, confidence in cuts: + if cut_page <= previous: continue - page = int(item.get("page") or 0) - if page > 0: - by_page[page] = item - return by_page - - -def _kind_for_range(start: int, end: int, by_page: dict[int, dict[str, Any]]) -> str: - kinds = [ - str(item.get("special_kind") or item.get("kind") or "normal") - for page, item in by_page.items() - if start <= page <= end - ] - if not kinds: - return "text_dense" - counts = Counter(kinds) - special_total = sum(count for kind, count in counts.items() if kind != "normal") - if special_total == 0: - return "text_dense" - top_kind, top_count = counts.most_common(1)[0] - if top_kind in {"table_heavy"}: - return "table_heavy" - if top_kind in {"image_heavy", "single_image"}: - return "image_heavy" - if top_kind == "landscape": - return "landscape_block" - if top_kind == "toc": - return "toc" - if top_kind in {"blank", "sparse"} and top_count >= max(1, (end - start + 1) // 2): - return "sparse" - return "mixed" - - -def _fallback_cuts( - page_count: int, - *, - max_pages_per_shard: int, - classifications: dict[str, Any], -) -> list[dict[str, Any]]: - by_page = _special_by_page(classifications) - cuts = [] - start = 1 - while start <= page_count: - target_end = min(start + max_pages_per_shard - 1, page_count) - end = target_end - if target_end < page_count: - # Prefer a safe boundary close to the shard limit without exceeding it. - target_kind = str(by_page.get(target_end, {}).get("special_kind") or "normal") - if target_kind not in {"blank", "sparse", "normal"}: - window_start = max(start, target_end - 5) - candidates = [] - for page in range(window_start, target_end): - kind = str(by_page.get(page, {}).get("special_kind") or "normal") - priority = {"blank": 0, "sparse": 1, "normal": 2}.get(kind) - if priority is not None: - candidates.append((abs(page - target_end), priority, page)) - if candidates: - end = min(candidates)[2] - cuts.append( - { - "start": start, - "end": end, - "predominant_kind": _kind_for_range(start, end, by_page), - "rationale": "deterministic fallback cut near max shard size", - } + shards.append( + Shard( + shard_index=len(shards), + page_start=previous + 1, + page_end=cut_page, + page_offset=previous, + anchor_type=anchor_type, # type: ignore[arg-type] + anchor_evidence=evidence, + confidence=confidence, + ) ) - start = end + 1 - return cuts + previous = cut_page + if previous < page_count: + shards.append( + Shard( + shard_index=len(shards), + page_start=previous + 1, + page_end=page_count, + page_offset=previous, + anchor_type="forced_max_size", + anchor_evidence="final shard", + confidence=1.0, + ) + ) + return shards -def _normalize_cuts( - cuts: list[dict[str, Any]], +def _build_prompt( *, page_count: int, - max_pages_per_shard: int, - classifications: dict[str, Any], -) -> list[dict[str, Any]]: - normalized = [] - expected = 1 - for raw in cuts: - if not isinstance(raw, dict): - continue - start = int(raw.get("start") or raw.get("page_start") or 0) - end = int(raw.get("end") or raw.get("page_end") or 0) - if start != expected or end < start or end > page_count: - raise ValueError("LLM shard cuts are not contiguous") - kind = str(raw.get("predominant_kind") or "mixed") - if kind not in ALLOWED_PREDOMINANT: - kind = "mixed" - normalized.append( - { - "start": start, - "end": end, - "predominant_kind": kind, - "rationale": str(raw.get("rationale") or "")[:500], - } - ) - expected = end + 1 - if expected != page_count + 1: - raise ValueError("LLM shard cuts do not cover all pages") - if not normalized: - raise ValueError("empty LLM shard cuts") - # Avoid accepting pathological single giant cuts except naturally small docs. - if page_count > max_pages_per_shard * 2 and any( - cut["end"] - cut["start"] + 1 > max_pages_per_shard * 2 - for cut in normalized - ): - raise ValueError("LLM shard cut exceeds hard tolerance") - return normalized + min_pages: int, + max_pages: int, + doc_stats: dict[str, Any], + page_kind_counts: dict[str, int], + toc_pages: list[int], + h1_pages: list[dict[str, Any]], + profile: dict[str, Any] | None, + visual_evidence: list[dict[str, Any]], + grep_history: list[dict[str, Any]], +) -> str: + payload = { + "page_count": page_count, + "min_pages_per_shard": min_pages, + "max_pages_per_shard": max_pages, + "page_kind_counts": page_kind_counts, + "doc_stats": doc_stats, + "toc_pages": toc_pages, + "h1_pages": h1_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" + "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 " + "shard sizes, and over-fragmentation.\n" + "- Prefer fewer, semantically coherent shards over many tiny shards.\n" + "- Keep each cut rationale under 120 characters.\n" + "- Every resulting shard length must be between min_pages_per_shard and " + "max_pages_per_shard, except the final shard may be shorter only when no better " + "valid split exists. Check each segment length exactly before returning.\n" + "- If no split is useful, return enabled=false and cuts=[] even for a long document.\n" + "Output schema:\n" + "{\n" + ' "enabled": boolean,\n' + ' "cuts": [\n' + " {\"cut_after_page\": number, \"anchor_type\": \"h1_boundary\" | " + "\"blank_separator\" | \"forced_max_size\", " + "\"confidence\": number, \"rationale\": string}\n" + " ],\n" + ' "reason": "llm_boundary_decision" | "not_needed" | "too_large",\n' + ' "rationale": string\n' + "}\n" + "Payload:\n" + + json.dumps(payload, ensure_ascii=False) + ) -def _build_global_signals( - *, - sampled_pages: list[dict[str, Any]], - classifications: dict[str, Any], -) -> GlobalSignals: - pages = classifications.get("pages", []) or [] - toc_pages = [ - int(page.get("page") or 0) - for page in pages - if str(page.get("special_kind") or page.get("kind")) == "toc" - ] - sample_size = len(sampled_pages) - if sample_size <= 0: - return GlobalSignals( - has_toc=bool(toc_pages), - toc_pages=toc_pages, - landscape_ratio=0.0, - table_page_ratio=0.0, - image_page_ratio=0.0, - sample_size=0, - notes=str(classifications.get("global_notes") or ""), - ) - landscape_count = sum(1 for page in sampled_pages if page.get("orientation") == "landscape") - table_count = sum( - 1 - for page in pages - if str(page.get("special_kind") or page.get("kind")) == "table_heavy" - ) - image_count = sum( - 1 - for page in pages - if str(page.get("special_kind") or page.get("kind")) in {"image_heavy", "single_image"} - ) - return GlobalSignals( - has_toc=bool(toc_pages), - toc_pages=toc_pages, - landscape_ratio=landscape_count / sample_size, - table_page_ratio=table_count / sample_size, - image_page_ratio=image_count / sample_size, - sample_size=sample_size, - notes=str(classifications.get("global_notes") or ""), - ) +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. + sanitized = (text or "").strip() + if len(sanitized) > max_length: + sanitized = sanitized[:max_length].rstrip() + "…" + return sanitized -def build_manifest_from_cuts( - *, - file_uri: str, - job_id: str, +def _validate_cut_lengths( + cuts: list[tuple[int, str, str, float]], page_count: int, - sampled_pages: list[dict[str, Any]], - classifications: dict[str, Any], - cuts: list[dict[str, Any]], - decision_log_ref: str = "local-debug", -) -> ShardManifest: - by_page = _special_by_page(classifications) - shards = [] - for cut in cuts: - start = int(cut["start"]) - end = int(cut["end"]) - special_pages = [] - for page in range(start, end + 1): - item = by_page.get(page) - if not item: - continue - kind = str(item.get("special_kind") or item.get("kind") or "normal") - if kind == "normal": - continue - special_pages.append( - SpecialPage( - page=page, - kind=kind, # type: ignore[arg-type] - confidence=float(item.get("confidence") or 0.0), - note=str(item.get("note") or ""), - ) + min_pages: int, + max_pages: int, +) -> None: + previous = 0 + for cut_page, *_ in cuts: + if cut_page - previous < min_pages: + raise ValueError( + f"LLM cut plan creates shard length {cut_page - previous} < min_pages={min_pages}" ) - shards.append( - ShardSignal( - page_start=start, - page_end=end, - page_offset=start - 1, - predominant_kind=cut["predominant_kind"], - special_pages=special_pages, - cut_rationale=str(cut.get("rationale") or ""), + if cut_page - previous > max_pages: + raise ValueError( + f"LLM cut plan creates shard length {cut_page - previous} > max_pages={max_pages}" ) + previous = cut_page + if page_count - previous < min_pages and cuts: + raise ValueError( + f"LLM cut plan creates final shard length {page_count - previous} < min_pages={min_pages}" + ) + if page_count - previous > max_pages: + raise ValueError( + f"LLM cut plan creates final shard length {page_count - previous} > max_pages={max_pages}" ) - - manifest = ShardManifest( - job_id=job_id, - file_uri=file_uri, - file_sha=_hash_file(file_uri), - page_count=page_count, - shard_count=len(shards), - shards=shards, - global_signals=_build_global_signals( - sampled_pages=sampled_pages, - classifications=classifications, - ), - decision_log_ref=decision_log_ref, - ) - manifest.validate() - return manifest -def _hash_file(file_uri: str) -> str: - digest = sha256() - try: - with open(file_uri, "rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - except OSError: - return sha256(file_uri.encode("utf-8")).hexdigest() +def _parse_llm_plan( + raw: str, + page_count: int, + min_pages: int, + max_pages: int, +) -> tuple[bool, list[tuple[int, str, str, float]], str, str]: + data = json.loads(raw) + if not isinstance(data, dict): + raise ValueError("LLM shard plan is not an object") + enabled = bool(data.get("enabled")) + reason = str(data.get("reason") or ("llm_boundary_decision" if enabled else "not_needed")) + rationale = _sanitize_rationale(str(data.get("rationale") or "")) + cuts: list[tuple[int, str, str, float]] = [] + for item in data.get("cuts") or []: + if not isinstance(item, dict): + continue + raw_cut_page = item.get("cut_after_page") + if raw_cut_page is None: + continue + cut_page = int(raw_cut_page) + 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"}: + 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)) + cuts = sorted({cut[0]: cut for cut in cuts}.values(), key=lambda cut: cut[0]) + if enabled: + _validate_cut_lengths(cuts, page_count, min_pages, max_pages) + return enabled, cuts, reason, rationale -def propose_shard_plan( +def _deterministic_guardrail_plan( *, - file_uri: str, - job_id: str, page_count: int, - sampled_pages: list[dict[str, Any]], - classifications: dict[str, Any], - max_pages_per_shard: int = 199, - model: str | None = None, - use_llm: bool = True, -) -> dict[str, Any]: - """Produce a validated shard proposal and manifest.""" - max_pages_per_shard = max(1, int(max_pages_per_shard)) + min_pages: int, + max_pages: int, + h1_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 + ] + if eligible: + chosen = max(eligible) + cut_page = chosen - 1 + cuts.append((cut_page, "h1_boundary", f"guardrail H1 start page {chosen}", 0.35)) + previous = cut_page + else: + cuts.append((target, "forced_max_size", "guardrail max shard size", 0.25)) + previous = target + # Merge final shard into previous if it's smaller than min_pages + if cuts and (page_count - cuts[-1][0]) < min_pages: + cuts.pop() + return cuts, "too_large" - cuts: list[dict[str, Any]] + +@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), +) +def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + page_count = ctx.blackboard.page_count + threshold, min_pages, max_pages = _thresholds(ctx) + if page_count <= threshold: + plan = single_shard_plan(page_count) + ctx.blackboard.shard_plan = plan + return ToolResult( + status="ok", + payload={"enabled": False, "shard_count": len(plan.shards)}, + 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( + 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) + warnings: list[str] = [] raw_response = "" - if use_llm and page_count > max_pages_per_shard: + rationale = "" + llm_attempted = False + if model and ctx.budget.try_reserve("plan", prompt_tokens_est): try: - from shared.core.config import settings + llm_attempted = True from shared.services.ai.openai_compatible_client_sync import get_openai_client - effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL - client = get_openai_client(model=effective_model) - payload = { - "page_count": page_count, - "max_pages_per_shard": max_pages_per_shard, - "sampled_pages": sampled_pages, - "page_classifications": classifications, - } - raw_response = client.chat_completion( - messages=[ - {"role": "system", "content": PROMPT}, - { - "role": "user", - "content": "Propose a shard plan as json:\n" - + json.dumps(payload, ensure_ascii=False), - }, - ], - model=effective_model, + 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=2200, + max_tokens=1600, response_format={"type": "json_object"}, ) - data = extract_json_object(raw_response) - cuts = _normalize_cuts( - list(data.get("cuts") or []), - page_count=page_count, - max_pages_per_shard=max_pages_per_shard, - classifications=classifications, - ) - if data.get("global_notes") and not classifications.get("global_notes"): - classifications = {**classifications, "global_notes": data.get("global_notes")} + 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: - logger.warning( - f"[document_agent.propose_shard_plan] LLM planning failed, " - f"using fallback: {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 = _fallback_cuts( - page_count, - max_pages_per_shard=max_pages_per_shard, - classifications=classifications, + 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 after malformed LLM shard decision." else: - cuts = _fallback_cuts( - page_count, - max_pages_per_shard=max_pages_per_shard, - classifications=classifications, - ) + 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, + }, + ) - manifest = build_manifest_from_cuts( - file_uri=file_uri, - job_id=job_id, - page_count=page_count, - sampled_pages=sampled_pages, - classifications=classifications, - cuts=cuts, + shards = _cuts_to_shards(cuts, page_count) + enabled = len(shards) > 1 + if not enabled: + reason = "not_needed" + plan = ShardPlan( + enabled=enabled, + reason=reason, # type: ignore[arg-type] + shards=shards, + validation=validate_shard_plan( + ShardPlan(enabled=enabled, reason=reason, shards=shards), # type: ignore[arg-type] + page_count=page_count, + min_pages=min_pages, + max_pages=max_pages, + ), + ) + ctx.blackboard.shard_plan = plan + return ToolResult( + status="ok", + payload={ + "enabled": plan.enabled, + "reason": plan.reason, + "shard_count": len(plan.shards), + "valid": plan.validation.valid, + }, + latency_ms=int((time.monotonic() - start) * 1000), + tokens_used=ctx.budget.snapshot()["plan"]["used"], + input_summary={ + "page_count": page_count, + "h1_count": len(h1_pages), + "model": model, + }, + output_summary={ + "enabled": plan.enabled, + "reason": plan.reason, + "rationale": rationale, + "shards": [shard.to_dict() for shard in plan.shards], + }, + warnings=warnings, + debug={ + "prompt_excerpt": prompt[:4000], + "raw_response_excerpt": raw_response[:4000], + "llm_attempted": llm_attempted, + }, ) - return { - "cuts": cuts, - "manifest": manifest, - "manifest_dict": manifest.to_dict(), - "raw_response": raw_response, - } 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 new file mode 100644 index 000000000..c9e8273c4 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/validate_anatomy_map.py @@ -0,0 +1,70 @@ +"""Validate the current anatomy blackboard.""" + +from __future__ import annotations + +import os +import time +from typing import Any + +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, +) +from app.services.document_agent.validators import validate_anatomy_map + + +def _thresholds(ctx: ToolContext) -> tuple[int, int]: + min_pages = int( + ctx.settings.get("min_pages_per_shard") + or os.environ.get("PARSE_AGENT_MIN_PAGES_PER_SHARD", "20") + ) + max_pages = int( + ctx.settings.get("max_pages_per_shard") + or os.environ.get("PARSE_AGENT_MAX_PAGES_PER_SHARD", "200") + ) + return min_pages, max_pages + + +@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), +) +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( + status="error", + error="blackboard is missing anatomy outputs", + latency_ms=int((time.monotonic() - start) * 1000), + ) + anatomy = PageAnatomyMap( + job_id=ctx.job_id, + file_path=ctx.pdf_path, + page_count=ctx.blackboard.page_count, + page_features=ctx.blackboard.page_features, + page_labels=ctx.blackboard.page_labels, + toc_result=ctx.blackboard.toc_result, + h1_result=ctx.blackboard.h1_result, + shard_plan=ctx.blackboard.shard_plan, + document_profile=ctx.blackboard.document_profile, + global_signals=ctx.blackboard.global_signals, + trace_summary={}, + ) + min_pages, max_pages = _thresholds(ctx) + report = validate_anatomy_map(anatomy, min_pages=min_pages, max_pages=max_pages) + ctx.blackboard.validation_report = report.to_dict() + if ctx.blackboard.shard_plan: + ctx.blackboard.shard_plan.validation = report + return ToolResult( + status="ok" if report.valid else "invalid", + payload=report.to_dict(), + latency_ms=int((time.monotonic() - start) * 1000), + ) diff --git a/apps/worker/app/services/document_agent/tools/verdict.py b/apps/worker/app/services/document_agent/tools/verdict.py new file mode 100644 index 000000000..bada1ba22 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/verdict.py @@ -0,0 +1,39 @@ +"""Agent verdict tool.""" + +from __future__ import annotations + +import time +from typing import Any + +from app.services.document_agent.manifest import AgentVerdict, ToolContext, ToolResult +from app.services.document_agent.registry import has_shard_plan, register_tool + + +@register_tool( + name="verdict", + description="Finish the document profile run with success or abort.", + parameters={ + "type": "object", + "properties": { + "status": {"type": "string", "enum": ["success", "abort"]}, + "rationale": {"type": "string"}, + }, + "required": ["status", "rationale"], + }, + preconditions=(has_shard_plan,), +) +def verdict(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + start = time.monotonic() + status = str(args.get("status") or "abort") + if status not in {"success", "abort"}: + status = "abort" + ctx.blackboard.verdict = AgentVerdict( + status=status, # type: ignore[arg-type] + rationale=str(args.get("rationale") or ""), + ) + return ToolResult( + status="ok", + payload=ctx.blackboard.verdict.to_dict(), + latency_ms=int((time.monotonic() - start) * 1000), + ) + diff --git a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py new file mode 100644 index 000000000..7c1d62479 --- /dev/null +++ b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py @@ -0,0 +1,361 @@ +"""VLM-native TOC entry extraction and hierarchy conversion.""" + +from __future__ import annotations + +import base64 +import json +import time +from dataclasses import dataclass +from typing import Any, cast + + +# --------------------------------------------------------------------------- +# Batch-mode prompt: send a window of candidate pages in one VLM call. +# The VLM first classifies each page, then extracts entries only from TOC pages. +# --------------------------------------------------------------------------- + +VLM_TOC_BATCH_PROMPT = """\ +You will receive {page_count} consecutive page screenshots from a document. +Some of these pages may be Table of Contents (TOC) pages, while others may be +regular body text, section dividers, blank pages, or other non-TOC content. + +**Your task has two parts:** + +### Part 1: Classify each page +For each page, decide whether it is a TOC page or not. + +A page IS a TOC page when it shows a STRUCTURED LISTING of document sections, +recognizable by MOST of these visual patterns: +- Multiple entry lines, each pairing a section/chapter TITLE with a PAGE NUMBER +- Leader characters (dots "......", dashes "------", or whitespace) connecting + titles on the left to page numbers aligned on the right +- Systematic numbering in the titles (1. / 1.1 / Chapter 1 / 一、 / 第一章, etc.) +- An explicit heading such as "Table of Contents", "Contents", "目录", or "目次" + (may appear only on the first page of a multi-page TOC) + +A page is NOT a TOC page when: +- It contains narrative paragraphs or body text, even if the text has numbered + headings (e.g. "1.0.1 为建立并落实..." followed by explanatory sentences) +- It is a section divider / title page with only a single heading and no listing +- It is blank or nearly blank +- It shows data tables, charts, or images rather than a contents listing +- It has numbered definitions or terms with explanations (e.g. "2.0.3 风险 risk") + — these are glossary/body content, NOT a TOC + +The KEY distinction: TOC entries are SHORT titles pointing to page numbers. +Body text has EXPLANATORY content after the heading. If a numbered item is +followed by sentences of explanation, it is body text, not a TOC entry. + +### Part 2: Extract entries from TOC pages only +For each page you classify as TOC, extract every entry with: +- title: the section/chapter name, verbatim, without trailing dots or leaders. + Combine wrapped lines into one string. Include numbering prefixes. +- page_number: integer for plain numbers, string for non-numeric (iv, F-1), + null when no page reference is visible. +- level: hierarchy depth from visual cues (1=top-level, 2=indented sub-entry, 3+=deeper). + Category headers or group labels without page numbers → level 1. + +Do NOT include the TOC heading itself ("Table of Contents", "目录", etc.) or +column labels ("Page", "页码"). + +Return strict JSON (no markdown fences): +{{ + "pages": [ + {{ + "page": , + "is_toc": true/false, + "entries": [{{"title": "...", "page_number": ..., "level": ...}}, ...] + }}, + ... + ] +}} + +For non-TOC pages, set "entries" to an empty array []. +""" + +VLM_TOC_BATCH_CONTINUATION = """\ + +--- Continuation Context --- +Previous batch(es) already confirmed TOC pages and extracted these entries: + +{previous_summary} + +Last active section: Level {last_l1_level}: "{last_l1_title}" + +Use this to maintain hierarchy consistency for any TOC pages in this batch. +""" + + +@dataclass +class BatchPageResult: + """Result for a single page within a batch VLM call.""" + + page: int + is_toc: bool + entries: list[dict[str, Any]] + + +@dataclass +class BatchTocResult: + """Result from a batch VLM TOC extraction call.""" + + page_results: list[BatchPageResult] + toc_pages: list[int] # pages classified as TOC + non_toc_pages: list[int] # pages classified as non-TOC + all_entries: list[dict[str, Any]] # entries from TOC pages only + meta: dict[str, Any] + + +def _build_batch_continuation(previous_entries: list[dict[str, Any]]) -> str: + """Build continuation context for batch mode.""" + if not previous_entries: + return "" + + tail = previous_entries[-8:] + summary_lines = [] + for entry in tail: + level = entry.get("level", "?") + title = entry.get("title", "?") + page_number = entry.get("page_number") + suffix = f" -> p.{page_number}" if page_number is not None else "" + summary_lines.append(f" L{level}: {title}{suffix}") + + if len(previous_entries) > 8: + summary_lines.insert( + 0, f" ... ({len(previous_entries) - 8} earlier entries omitted)" + ) + + previous_summary = "\n".join(summary_lines) + last_l1 = None + for entry in reversed(previous_entries): + if entry.get("level") == 1: + last_l1 = entry + break + + if last_l1 is None: + return ( + "\n\n--- Continuation Context ---\n" + f"Previous batch extracted entries:\n{previous_summary}\n" + ) + + return VLM_TOC_BATCH_CONTINUATION.format( + previous_summary=previous_summary, + last_l1_level=last_l1.get("level", 1), + last_l1_title=last_l1.get("title", "?"), + ) + + +def vlm_extract_toc_batch( + *, + page_pngs: list[tuple[int, str]], + model: str, + previous_entries: list[dict[str, Any]] | None = None, +) -> BatchTocResult: + """Extract TOC entries from a batch of page images in a single VLM call. + + Args: + page_pngs: list of (page_number, png_path) pairs, in page order. + model: VLM model name. + previous_entries: entries from prior batches, for continuation context. + + Returns: + BatchTocResult with per-page classification and extracted entries. + """ + from loguru import logger + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + if not page_pngs: + return BatchTocResult( + page_results=[], toc_pages=[], non_toc_pages=[], + all_entries=[], meta={}, + ) + + prompt_text = VLM_TOC_BATCH_PROMPT.format(page_count=len(page_pngs)) + prompt_text += _build_batch_continuation(previous_entries or []) + + content_parts: list[dict[str, Any]] = [ + {"type": "text", "text": prompt_text}, + ] + for page_num, png_path in page_pngs: + with open(png_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode() + content_parts.append( + {"type": "text", "text": f"\n--- Page {page_num} ---"} + ) + content_parts.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{img_b64}"}, + } + ) + + start = time.monotonic() + 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.1, + max_tokens=8192, + response_format={"type": "json_object"}, + ) + elapsed_ms = int((time.monotonic() - start) * 1000) + + data = json.loads(raw) + raw_pages: list[dict[str, Any]] = [] + if isinstance(data, dict): + raw_pages = data.get("pages", []) + elif isinstance(data, list): + raw_pages = data + + # Build lookup from VLM response + page_lookup: dict[int, dict[str, Any]] = {} + for item in raw_pages: + if isinstance(item, dict) and "page" in item: + page_lookup[int(item["page"])] = item + + # Process results for each page in the original order + page_results: list[BatchPageResult] = [] + toc_pages: list[int] = [] + non_toc_pages: list[int] = [] + all_entries: list[dict[str, Any]] = [] + + for page_num, _png_path in page_pngs: + vlm_page = page_lookup.get(page_num, {}) + is_toc = bool(vlm_page.get("is_toc", False)) + raw_entries = vlm_page.get("entries", []) + + entries: list[dict[str, Any]] = [] + if is_toc: + for entry_item in raw_entries: + if not isinstance(entry_item, dict): + continue + title = str(entry_item.get("title") or "").strip() + if not title: + continue + try: + level = int(entry_item.get("level") or 1) + except (TypeError, ValueError): + level = 1 + entries.append( + { + "title": title, + "page_number": entry_item.get("page_number"), + "level": level, + } + ) + toc_pages.append(page_num) + else: + non_toc_pages.append(page_num) + + page_results.append( + BatchPageResult(page=page_num, is_toc=is_toc, entries=entries) + ) + all_entries.extend(entries) + + logger.info( + "[vlm_toc_batch] {} pages: toc={} non_toc={} entries={} elapsed={}ms", + len(page_pngs), + toc_pages, + non_toc_pages, + len(all_entries), + elapsed_ms, + ) + + return BatchTocResult( + page_results=page_results, + toc_pages=toc_pages, + non_toc_pages=non_toc_pages, + all_entries=all_entries, + meta={ + "pages_sent": [p for p, _ in page_pngs], + "model": model, + "elapsed_ms": elapsed_ms, + "usage": dict(usage), + "raw_response_length": len(raw), + "has_continuation_context": bool(previous_entries), + }, + ) + + + +def build_toc_tree(entries: list[dict[str, Any]]) -> dict[str, Any]: + root: dict[str, Any] = {} + stack: list[tuple[dict[str, Any], int]] = [(root, 0)] + positive_levels = [ + int(entry["level"]) + for entry in entries + if isinstance(entry.get("level"), int) and entry["level"] > 0 + ] + level_for_minus_one = max(positive_levels) + 1 if positive_levels else 1 + + for entry in entries: + heading = str(entry.get("title") or "").strip() + if not heading: + continue + original_level = entry.get("level", 1) + level = level_for_minus_one if original_level == -1 else int(original_level or 1) + while len(stack) > 1 and stack[-1][1] >= level: + stack.pop() + parent = stack[-1][0] + parent[heading] = {} + stack.append((parent[heading], level)) + return root + + +def build_toc_with_level_md(entries: list[dict[str, Any]]) -> str: + """Build a compact Markdown table of TOC entries (heading + level only). + + Accepts both raw VLM entries (key='title') and stored toc_with_level + dicts (key='heading'). Call this dynamically at consumption time rather + than persisting the MD string. + """ + if not entries: + return "" + lines = ["| heading | level |", "|---------|-------|"] + for entry in entries: + heading = str( + entry.get("heading") or entry.get("title") or "" + ).strip().replace("|", "\\|") + level = entry.get("level", 1) + lines.append(f"| {heading} | {level} |") + return "\n".join(lines) + + +def vlm_entries_to_toc_hierarchies( + entries: list[dict[str, Any]], + *, + toc_page_nums: list[int], + scan_end_page: int | None = None, + page_count: int | None = None, +) -> list[dict[str, Any]]: + if not entries or not toc_page_nums: + return [] + + toc_with_level = [] + for entry in entries: + toc_with_level.append( + { + "heading": str(entry.get("title") or "").strip(), + "level": entry.get("level", 1), + "page_number": entry.get("page_number"), + } + ) + + start_page = min(toc_page_nums) + end_page = max(toc_page_nums) + if scan_end_page is None: + scan_end_page = start_page + if page_count is not None: + scan_end_page = min(scan_end_page, page_count) + + return [ + { + "toc_range": [start_page, end_page], + "toc_range_unit": "page", + "scan_range": [start_page, scan_end_page], + "source": "vlm", + "toc_with_level": toc_with_level, + "toc_tree": build_toc_tree(entries), + } + ] + diff --git a/apps/worker/app/services/document_agent/trace.py b/apps/worker/app/services/document_agent/trace.py new file mode 100644 index 000000000..b455b63e8 --- /dev/null +++ b/apps/worker/app/services/document_agent/trace.py @@ -0,0 +1,173 @@ +"""Best-effort parse-agent trace buffering and database persistence.""" + +from __future__ import annotations + +import time +from datetime import datetime +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from loguru import logger + +from app.services.document_agent.manifest import PageAnatomyMap, ToolResult + + +class ParseRunRecorder: + def __init__(self, *, job_id: str, db: Any | None = None) -> None: + self.run_id = f"prof_{uuid4().hex[:12]}" + self.job_id = job_id + self._db = db + self._started = time.monotonic() + self._steps: list[dict[str, Any]] = [] + self._anatomy: PageAnatomyMap | None = None + self._artifact_path: str | None = None + + def record_step( + self, + *, + round_index: int, + actor: str, + action_type: str, + result: ToolResult, + tool_name: str | None = None, + tool_args: dict[str, Any] | None = None, + ) -> None: + self._steps.append( + { + "round_index": round_index, + "actor": actor, + "action_type": action_type, + "tool_name": tool_name, + "tool_args": tool_args or {}, + "observation": { + "status": result.status, + "payload_keys": sorted(result.payload.keys()), + "payload": result.payload, + "input_summary": result.input_summary, + "output_summary": result.output_summary, + "warnings": list(result.warnings), + "debug": result.debug, + "error": result.error, + }, + "tokens_used": result.tokens_used, + "latency_ms": result.latency_ms, + "created_at": datetime.utcnow(), + } + ) + + def set_anatomy_map(self, anatomy: PageAnatomyMap, artifact_path: str) -> None: + self._anatomy = anatomy + self._artifact_path = artifact_path + self.write_trace_json(str(Path(artifact_path).with_name("trace.json"))) + + def write_trace_artifact( + self, + output_dir: str | None, + *, + final_status: str, + summary: dict[str, Any] | None = None, + ) -> None: + if output_dir is None: + return + self.write_trace_json( + str(Path(output_dir) / "trace.json"), + final_status=final_status, + summary=summary, + ) + + def write_trace_json( + self, + trace_path: str, + *, + final_status: str | None = None, + summary: dict[str, Any] | None = None, + ) -> None: + try: + import json + + serializable_steps = [] + for step in self._steps: + item = dict(step) + created_at = item.get("created_at") + if created_at is not None and hasattr(created_at, "isoformat"): + item["created_at"] = created_at.isoformat() + serializable_steps.append(item) + Path(trace_path).write_text( + json.dumps( + { + "run_id": self.run_id, + "job_id": self.job_id, + "final_status": final_status, + "summary": summary, + "artifact_path": self._artifact_path, + "steps": serializable_steps, + }, + ensure_ascii=False, + indent=2, + default=str, + ), + encoding="utf-8", + ) + except Exception as exc: + logger.debug(f"parse agent trace json write failed: {exc}") + + def summary(self) -> dict[str, Any]: + return { + "run_id": self.run_id, + "step_count": len(self._steps), + "artifact_path": self._artifact_path, + "latency_ms": int((time.monotonic() - self._started) * 1000), + } + + def flush(self, *, final_status: str, summary: dict[str, Any] | None = None) -> None: + if self._db is None: + return + try: + from shared.models.database.parse_agent import ParseRun, ParseStep + from shared.models.database.document_page_plan import DocumentPagePlan + + run = ParseRun( + run_id=self.run_id, + job_id=self.job_id, + kind="profile", + final_status=final_status, + rounds_count=max((step["round_index"] for step in self._steps), default=0) + 1, + total_tokens=sum(int(step.get("tokens_used") or 0) for step in self._steps), + total_latency_ms=int((time.monotonic() - self._started) * 1000), + summary=summary or self.summary(), + ) + self._db.add(run) + for index, step in enumerate(self._steps): + self._db.add( + ParseStep( + step_id=f"prst_{uuid4().hex[:12]}", + run_id=self.run_id, + round_index=int(step["round_index"]), + actor=str(step["actor"]), + action_type=str(step["action_type"]), + tool_name=step.get("tool_name"), + tool_args=step.get("tool_args"), + observation=step.get("observation"), + tokens_used=int(step.get("tokens_used") or 0), + latency_ms=int(step.get("latency_ms") or 0), + created_at=step.get("created_at"), + ) + ) + if self._anatomy is not None: + self._db.add( + DocumentPagePlan( + page_plan_id=f"dpp_{uuid4().hex[:12]}", + job_id=self.job_id, + page_count=self._anatomy.page_count, + shard_plan=self._anatomy.shard_plan.to_dict(), + global_signals=self._anatomy.global_signals, + ) + ) + self._db.flush() + except Exception as exc: + logger.debug(f"parse agent trace flush failed: {exc}") + try: + self._db.rollback() + except Exception: + pass diff --git a/apps/worker/app/services/document_agent/validators.py b/apps/worker/app/services/document_agent/validators.py new file mode 100644 index 000000000..2f06371ba --- /dev/null +++ b/apps/worker/app/services/document_agent/validators.py @@ -0,0 +1,102 @@ +"""Validation and repair for page anatomy outputs.""" + +from __future__ import annotations + +from app.services.document_agent.manifest import ( + PageAnatomyMap, + Shard, + ShardPlan, + ValidationReport, +) + + +def validate_shard_plan( + plan: ShardPlan, + *, + page_count: int, + min_pages: int, + max_pages: int, +) -> ValidationReport: + errors: list[str] = [] + warnings: list[str] = [] + if not plan.shards: + errors.append("shard_plan has no shards") + return ValidationReport(valid=False, errors=errors, warnings=warnings) + expected_start = 1 + for shard in sorted(plan.shards, key=lambda item: item.shard_index): + if shard.page_start != expected_start: + errors.append( + f"shard {shard.shard_index} starts at {shard.page_start}, expected {expected_start}" + ) + if shard.page_end < shard.page_start: + errors.append(f"shard {shard.shard_index} has invalid range") + if shard.page_offset != shard.page_start - 1: + errors.append(f"shard {shard.shard_index} page_offset mismatch") + length = shard.page_end - shard.page_start + 1 + if plan.enabled and length > max_pages: + errors.append(f"shard {shard.shard_index} exceeds max_pages={max_pages}") + if plan.enabled and length < min_pages: + errors.append(f"shard {shard.shard_index} shorter than min_pages={min_pages}") + expected_start = shard.page_end + 1 + if expected_start != page_count + 1: + errors.append("shard_plan does not cover full document") + return ValidationReport(valid=not errors, errors=errors, warnings=warnings) + + +def single_shard_plan(page_count: int) -> ShardPlan: + return ShardPlan( + enabled=False, + reason="not_needed", + shards=[ + Shard( + shard_index=0, + page_start=1, + page_end=max(page_count, 1), + page_offset=0, + anchor_type="forced_max_size", + anchor_evidence="document within shard threshold", + confidence=1.0, + ) + ], + ) + + +def validate_anatomy_map( + anatomy: PageAnatomyMap, + *, + min_pages: int, + max_pages: int, +) -> ValidationReport: + errors: list[str] = [] + warnings: list[str] = [] + page_count = anatomy.page_count + feature_pages = {feature.page for feature in anatomy.page_features} + label_pages = {label.page for label in anatomy.page_labels} + expected_pages = set(range(1, page_count + 1)) + if feature_pages != expected_pages: + errors.append("page_features do not cover every page") + 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") + shard_report = validate_shard_plan( + anatomy.shard_plan, + page_count=page_count, + min_pages=min_pages, + max_pages=max_pages, + ) + errors.extend(shard_report.errors) + warnings.extend(shard_report.warnings) + if anatomy.shard_plan.enabled: + forced_count = sum( + 1 + for shard in anatomy.shard_plan.shards + if shard.anchor_type == "forced_max_size" + ) + if forced_count == len(anatomy.shard_plan.shards): + warnings.append("all shards are based on forced max-size boundaries") + return ValidationReport(valid=not errors, errors=errors, warnings=warnings) diff --git a/apps/worker/app/services/document_agent/visual.py b/apps/worker/app/services/document_agent/visual.py new file mode 100644 index 000000000..5a6e6da3c --- /dev/null +++ b/apps/worker/app/services/document_agent/visual.py @@ -0,0 +1,86 @@ +"""Shared page rendering helpers for document-agent visual reasoning.""" + +from __future__ import annotations + +import gc +import os +from pathlib import Path +from typing import Any + +from app.services.document_agent.manifest import ToolContext +from app.services.document_parser.formats.pdf.pymupdf_subprocess import ( + run_in_child_process, + worker, +) + + +@worker +def _render_pages_worker( + queue, + pdf_path: str, + pages: list[int], + output_dir: str, + dpi: int, + prefix: str, +) -> None: + import pymupdf # type: ignore[import] + + results: list[dict[str, Any]] = [] + try: + doc = pymupdf.open(pdf_path) + for page_num in pages: + idx = page_num - 1 + if 0 <= idx < doc.page_count: + page = doc[idx] + mat = pymupdf.Matrix(dpi / 72.0, dpi / 72.0) + pix = page.get_pixmap(matrix=mat) + png_name = f"{prefix}_page_{page_num}.png" + png_path = os.path.join(output_dir, png_name) + pix.save(png_path) + results.append({"page": page_num, "png_path": png_path}) + finally: + try: + doc.close() + except Exception: + pass + gc.collect() + queue.put({"ok": True, "results": results}) + + +def visual_output_dir(ctx: ToolContext, folder_name: str = "agent_visuals") -> str: + output_dir = str( + Path(ctx.output_dir or os.path.expanduser("~/.knowhere/_debug_profile")) + / folder_name + ) + os.makedirs(output_dir, exist_ok=True) + return output_dir + + +def render_pages( + ctx: ToolContext, + pages: list[int], + *, + folder_name: str = "agent_visuals", + prefix: str = "visual", + dpi: int | None = None, + timeout: int = 120, +) -> list[dict[str, Any]]: + if not pages: + return [] + page_count = max(int(ctx.blackboard.page_count or 0), 0) + bounded_pages = sorted({page for page in pages if 1 <= page <= page_count}) + if not bounded_pages: + return [] + output_dir = visual_output_dir(ctx, folder_name=folder_name) + effective_dpi = dpi or int(ctx.settings.get("agent_png_dpi", "144")) + result = run_in_child_process( + _render_pages_worker, + ctx.pdf_path, + bounded_pages, + output_dir, + effective_dpi, + prefix, + timeout=timeout, + ) + return list(result.get("results") or []) + diff --git a/apps/worker/app/services/document_parser/assets/inline_asset.py b/apps/worker/app/services/document_parser/assets/inline_asset.py index 9e792ee1d..1d67eae09 100644 --- a/apps/worker/app/services/document_parser/assets/inline_asset.py +++ b/apps/worker/app/services/document_parser/assets/inline_asset.py @@ -10,7 +10,6 @@ def build_image_asset_row( summary: str, know_id: str, addtime: str, - page_nums: str = "", ) -> ParsedRow: return ParsedRow( content=content, @@ -22,7 +21,6 @@ def build_image_asset_row( tokens="", connectto="", addtime=addtime, - page_nums=page_nums, ) @@ -34,7 +32,6 @@ def build_table_asset_row( keywords: str, know_id: str, addtime: str, - page_nums: str = "", ) -> ParsedRow: return ParsedRow( content=content, @@ -46,5 +43,5 @@ def build_table_asset_row( tokens="", connectto="", addtime=addtime, - page_nums=page_nums, ) + diff --git a/apps/worker/app/services/document_parser/formats/atlas/parser.py b/apps/worker/app/services/document_parser/formats/atlas/parser.py index a3210d67e..ddac6a6d0 100644 --- a/apps/worker/app/services/document_parser/formats/atlas/parser.py +++ b/apps/worker/app/services/document_parser/formats/atlas/parser.py @@ -444,7 +444,6 @@ def _vlm_task(page_num, img_name): know_id=know_id, addtime=time_stamp, tokens=tokens, - page_nums=str(page_num), ) ) diff --git a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py index b71ace46a..ed5f51abf 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py @@ -39,7 +39,6 @@ class MarkdownImageAssetRequest: last_context: str image_summary: str | None timestamp: str - current_page_number: int seen_images: dict[str, dict[str, str]] summary_image: bool row_index: int @@ -65,7 +64,6 @@ def build_markdown_image_asset( source_path=source_path, cache_entry=request.seen_images[image_binary_hash], timestamp=request.timestamp, - current_page_number=request.current_page_number, ) relative_image_path = f"images/{request.image_name}{image_suffix}" @@ -91,7 +89,6 @@ def build_markdown_image_asset( summary=image_summary_field, know_id=image_know_id, timestamp=request.timestamp, - current_page_number=request.current_page_number, ) cache_entry = { "relative_img_path": relative_image_path, @@ -166,7 +163,6 @@ def _build_duplicate_image_asset( source_path: Path, cache_entry: dict[str, str], timestamp: str, - current_page_number: int, ) -> MarkdownImageAsset: row_values = _build_image_row_values( content=cache_entry["img_content"], @@ -174,7 +170,6 @@ def _build_duplicate_image_asset( summary=cache_entry["img_summary_field"], know_id=cache_entry["temp_uid"], timestamp=timestamp, - current_page_number=current_page_number, ) try: source_path.unlink() @@ -205,7 +200,6 @@ def _build_image_row_values( summary: str, know_id: str, timestamp: str, - current_page_number: int, ) -> ParserRowValues: image_row = build_image_asset_row( content=content, @@ -213,7 +207,6 @@ def _build_image_row_values( summary=summary, know_id=know_id, addtime=timestamp, - page_nums=str(current_page_number) if current_page_number > 0 else "", ) return cast(ParserRowValues, image_row.to_list()) diff --git a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py index b7c6da508..95e5088e8 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/parse_state.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py @@ -1,6 +1,6 @@ from __future__ import annotations -import re + from collections.abc import Callable from dataclasses import dataclass, field from typing import Any @@ -35,8 +35,7 @@ class MarkdownParseState: inner_paths: list[str] = field(default_factory=list) error_line_numbers: list[int] = field(default_factory=list) table_lines: list[str] = field(default_factory=list) - current_page_number: int = 0 - chunk_pages: set[int] = field(default_factory=set) + base_level: int | None = None path: str = "" path_counter: dict[str, int] = field(default_factory=dict) @@ -50,45 +49,38 @@ def __post_init__(self) -> None: self.path = self.relative_root def record_page_marker(self, line: str) -> bool: + """Detect and skip HTML comment lines (page markers, slide markers). + + Page number tracking has been removed; PAGE MEMORY will provide + accurate page numbers in a future release. + """ if "" not in line: return False if "page" not in line and "Slide number" not in line: return False - - page_match = re.search(r"page\s+(\d+)", line) - if page_match: - self.current_page_number = int(page_match.group(1)) - else: - self.current_page_number += 1 - self.chunk_pages.add(self.current_page_number) return True def flush_current_content(self) -> None: - page_numbers = self._format_chunk_pages() self.rows = self.row_updater( self.rows, self.content_items, self.path, self.llm_parameters, self.timestamp, - page_numbers, + "", 1500, True, ) self.content_items = [] - self.chunk_pages = set() - if self.current_page_number > 0: - self.chunk_pages.add(self.current_page_number) def flush_placeholder_chunk(self) -> None: - page_numbers = self._format_chunk_pages() self.rows = self.row_updater( self.rows, [], self.path, self.llm_parameters, self.timestamp, - page_numbers, + "", 1500, True, ) @@ -135,8 +127,6 @@ def append_content_item(self, item: str) -> None: def append_plain_text(self, text: str) -> None: self.content_items.append(text.strip()) - if self.current_page_number > 0: - self.chunk_pages.add(self.current_page_number) def append_row(self, row: ParserRowValues) -> None: self.rows.append(row) @@ -181,5 +171,4 @@ def to_dataframe(self) -> pd.DataFrame: ) return process_dup_paths_df(rows_builder.to_dataframe()) - def _format_chunk_pages(self) -> str: - return ",".join(str(page) for page in sorted(self.chunk_pages)) + diff --git a/apps/worker/app/services/document_parser/formats/markdown/parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py index a33be14e9..614bab23a 100755 --- a/apps/worker/app/services/document_parser/formats/markdown/parser.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py @@ -213,51 +213,98 @@ def parse_md( md_lines=None, base_llm_paras=None, relative_root=None, + toc_hierarchies=None, + lines_with_heading=None, ): - if md_lines is None and file_path is not None: - from app.services.common.file_loading import is_remote, load_file_bytes - - if is_remote(file_path): - file_bytes = load_file_bytes(file_path) - md_content = file_bytes.decode("utf-8") - md_lines = md_content.splitlines() - else: - with open(file_path, "r", encoding="utf-8") as file: - md_lines = file.readlines() - - md_lines = [line.strip() for line in md_lines if line.strip() != ""] + if lines_with_heading is not None: + # ── Phase A bypass ── + # Caller (e.g. oversized PDF shard-first path) already ran per-shard + # heading prediction and passed in the merged lines_with_heading. + # Skip TOC detection and heading prediction entirely. + logger.info( + f"📌 Using pre-identified headings ({len(lines_with_heading)} lines), " + f"skipping TOC detection and heading prediction" + ) + else: + # ── Phase A: TOC detection + heading prediction ── + if md_lines is None and file_path is not None: + from app.services.common.file_loading import is_remote, load_file_bytes + + if is_remote(file_path): + file_bytes = load_file_bytes(file_path) + md_content = file_bytes.decode("utf-8") + md_lines = md_content.splitlines() + else: + with open(file_path, "r", encoding="utf-8") as file: + md_lines = file.readlines() - # Preprocess: merge multi-line HTML tables into single lines - md_lines = merge_html_tables(md_lines) + md_lines = [line.strip() for line in md_lines if line.strip() != ""] - # Detect TOC using async LLM-based detection - toc_model_name = ( - base_llm_paras.get("model_name", settings.NORMOL_MODEL) - if base_llm_paras - else settings.NORMOL_MODEL - ) - hierarchy_model_name = ( - (base_llm_paras.get("hierarchy_model_name") or toc_model_name) - if base_llm_paras - else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) - ) + # Preprocess: merge multi-line HTML tables into single lines + md_lines = merge_html_tables(md_lines) - with stage_timer( - "md.detect_toc", line_count=len(md_lines), model_name=toc_model_name - ): - toc_hierarchies, md_lines = detect_tocs_in_texts( - md_lines, - model_name=toc_model_name, - hierarchy_model_name=hierarchy_model_name, + # Detect TOC using async LLM-based detection + toc_model_name = ( + base_llm_paras.get("model_name", settings.NORMOL_MODEL) + if base_llm_paras + else settings.NORMOL_MODEL ) + hierarchy_model_name = ( + (base_llm_paras.get("hierarchy_model_name") or toc_model_name) + if base_llm_paras + else (settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL) + ) + + if toc_hierarchies is not None: + # Pre-detected TOC from upstream (e.g. DOC_AGENT VLM-based extraction). + # Skip row-based detection entirely — TOC pages have already been + # physically stripped from the PDF, so no TOC rows exist in md_lines. + logger.info( + f"📌 Using pre-detected TOC hierarchies " + f"({len(toc_hierarchies)} regions), " + f"skipping detect_tocs_in_texts" + ) + else: + with stage_timer( + "md.detect_toc", line_count=len(md_lines), model_name=toc_model_name + ): + toc_hierarchies, md_lines = detect_tocs_in_texts( + md_lines, + model_name=toc_model_name, + hierarchy_model_name=hierarchy_model_name, + ) - # Save toc_hierarchies.json to output_dir (will be included in final zip package) - if toc_hierarchies: - toc_json_path = os.path.join(output_dir, "toc_hierarchies.json") - with open(toc_json_path, "w", encoding="utf-8") as f: - json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) - logger.info(f"Saved TOC hierarchies to {toc_json_path}") + # Save toc_hierarchies.json to output_dir (will be included in final zip package) + if toc_hierarchies: + toc_json_path = os.path.join(output_dir, "toc_hierarchies.json") + with open(toc_json_path, "w", encoding="utf-8") as f: + json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) + logger.info(f"Saved TOC hierarchies to {toc_json_path}") + + # Find layout.json path + layout_json_path = os.path.join(output_dir, "layout.json") + if not os.path.exists(layout_json_path): + layout_json_path = None + logger.debug("layout.json not found, META features will not be added") + + # estimate hierarchies with toc_hierarchies context + with stage_timer( + "md.predict_headings", + line_count=len(md_lines), + smart_parse=base_llm_paras["smart_title_parse"], + model_name=hierarchy_model_name, + ): + lines_with_heading = eval_md_headings( + md_lines, + source_type, + toc_hierarchies=toc_hierarchies, + smart_parse=base_llm_paras["smart_title_parse"], + model_name=hierarchy_model_name, + output_dir=output_dir, + layout_json_path=layout_json_path, + ) + # ── Phase B: MarkdownParseState traversal ── # Clean old artifacts to prevent accumulation across debug runs. # In production each job uses a fresh workspace so rmtree never triggers. tb_dir = os.path.join(output_dir, "tables") @@ -282,29 +329,6 @@ def parse_md( row_updater=update_df_list, ) - # Find layout.json path - layout_json_path = os.path.join(output_dir, "layout.json") - if not os.path.exists(layout_json_path): - layout_json_path = None - logger.debug("layout.json not found, META features will not be added") - - # estimate hierarchies with toc_hierarchies context - with stage_timer( - "md.predict_headings", - line_count=len(md_lines), - smart_parse=base_llm_paras["smart_title_parse"], - model_name=hierarchy_model_name, - ): - lines_with_heading = eval_md_headings( - md_lines, - source_type, - toc_hierarchies=toc_hierarchies, - smart_parse=base_llm_paras["smart_title_parse"], - model_name=hierarchy_model_name, - output_dir=output_dir, - layout_json_path=layout_json_path, - ) - logger.debug("Parsing md data... total_lines={}", len(lines_with_heading)) for i, line in enumerate(lines_with_heading): if parser_state.record_page_marker(line): @@ -347,7 +371,6 @@ def parse_md( last_context=last_context, image_summary=img_summary, timestamp=parser_state.timestamp, - current_page_number=parser_state.current_page_number, seen_images=parser_state.seen_images, summary_image=bool(base_llm_paras["summary_image"]), row_index=len(parser_state.rows), @@ -407,7 +430,6 @@ def parse_md( table_dir=tb_dir, table_count=parser_state.table_count, timestamp=parser_state.timestamp, - current_page_number=parser_state.current_page_number, summary_table=bool(base_llm_paras["summary_table"]), row_index=len(parser_state.rows), ) diff --git a/apps/worker/app/services/document_parser/formats/markdown/table_asset.py b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py index 19f795f4f..a489e036d 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/table_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py @@ -32,7 +32,6 @@ class MarkdownTableAssetRequest: table_dir: str table_count: int timestamp: str - current_page_number: int summary_table: bool row_index: int @@ -62,9 +61,6 @@ def build_markdown_table_asset( keywords="", know_id=gen_str_codes((request.table_html + str(request.table_count))), addtime=request.timestamp, - page_nums=str(request.current_page_number) - if request.current_page_number > 0 - else "", ) deferred_task = None 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 6a8dfaf4a..c5baf58f3 100755 --- a/apps/worker/app/services/document_parser/formats/pdf/parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -1,211 +1,12 @@ # pyright: reportArgumentType=false -import json import os -import re from app.services.document_parser.formats.markdown.parser import parse_md from app.services.document_parser.providers.mineru.pdf_service import parse_via_full -from app.services.document_parser.formats.pdf.pymupdf_subprocess import worker from app.services.document_parser.support.stage_profiler import stage_timer from loguru import logger - -def _inject_page_markers(output_dir: str) -> None: - """Inject markers into full.md using layout.json page info. - - Reads layout.json to find the first text content of each page, - then searches for that text in full.md and inserts a page marker above it. - - If layout.json is not available (e.g. fast path without MinerU), - this function does nothing gracefully. - """ - layout_path = os.path.join(output_dir, "layout.json") - md_path = os.path.join(output_dir, "full.md") - - if not os.path.exists(layout_path) or not os.path.exists(md_path): - logger.debug("layout.json or full.md not found, skipping page marker injection") - return - - try: - with open(layout_path, "r", encoding="utf-8") as f: - layout_data = json.load(f) - except (json.JSONDecodeError, IOError) as e: - logger.warning(f"Failed to read layout.json: {e}") - return - - pdf_info = layout_data.get("pdf_info", []) - if not pdf_info: - return - - with open(md_path, "r", encoding="utf-8") as f: - md_lines = f.readlines() - - # Build anchor map: {normalized_text: page_number (1-based)} - # Use the first text span of each page's first para_block as anchor - anchors = [] # list of (anchor_text, page_num) - for page in pdf_info: - page_idx = page.get("page_idx", 0) - page_num = page_idx + 1 # 1-based page number - - # Find the first non-empty text content in this page - anchor_text = None - for block in page.get("para_blocks", []): - for line in block.get("lines", []): - for span in line.get("spans", []): - content = span.get("content", "").strip() - if content and len(content) >= 3: # skip very short anchors - anchor_text = content - break - if anchor_text: - break - if anchor_text: - break - - if anchor_text: - anchors.append((anchor_text, page_num)) - - if not anchors: - logger.debug( - "No anchor texts found in layout.json, skipping page marker injection" - ) - return - - # Match anchors against md_lines and insert markers - # Process from end to start so line indices don't shift - insertions = [] # list of (line_index, page_num) - used_lines = set() - - for anchor_text, page_num in anchors: - # Normalize anchor for matching - anchor_norm = re.sub(r"\s+", " ", anchor_text).strip() - if len(anchor_norm) < 3: - continue - - # Search for anchor in md_lines (use first 50 chars for substring match) - search_key = anchor_norm[:50] - for i, line in enumerate(md_lines): - if i in used_lines: - continue - line_norm = re.sub(r"^#+\s*", "", line.strip()) - line_norm = re.sub(r"\s+", " ", line_norm).strip() - if search_key in line_norm: - insertions.append((i, page_num)) - used_lines.add(i) - break - - if not insertions: - logger.debug("No page marker matches found, skipping injection") - return - - # Sort by line index descending to insert from bottom to top - insertions.sort(key=lambda x: x[0], reverse=True) - for line_idx, page_num in insertions: - md_lines.insert(line_idx, f"\n") - - # Write back - with open(md_path, "w", encoding="utf-8") as f: - f.writelines(md_lines) - - logger.info(f"Injected {len(insertions)} page markers into full.md") - - -def _inject_page_markers_pymupdf(pdf_path: str, output_dir: str) -> None: - """Inject markers into full.md for pymupdf4llm fast path. - - Must run inside the same process that holds the PyMuPDF import. - """ - import pymupdf - - md_path = os.path.join(output_dir, "full.md") - if not os.path.exists(md_path): - return - - try: - doc = pymupdf.open(pdf_path) - except Exception: - return - - with open(md_path, "r", encoding="utf-8") as f: - md_lines = f.readlines() - - anchors = [] - for page_idx in range(len(doc)): - page = doc[page_idx] - page_num = page_idx + 1 - blocks = page.get_text("blocks") - for block in blocks: - if block[6] == 0: - text = block[4].strip().split("\n")[0].strip() - if text and len(text) >= 3: - anchors.append((text, page_num)) - break - - doc.close() - - if not anchors: - return - - insertions = [] - used_lines = set() - - for anchor_text, page_num in anchors: - anchor_norm = re.sub(r"\s+", " ", anchor_text).strip() - search_key = anchor_norm[:50] - for i, line in enumerate(md_lines): - if i in used_lines: - continue - line_norm = re.sub(r"^#+\s*", "", line.strip()) - line_norm = re.sub(r"\s+", " ", line_norm).strip() - if search_key in line_norm: - insertions.append((i, page_num)) - used_lines.add(i) - break - - if not insertions: - return - - insertions.sort(key=lambda x: x[0], reverse=True) - for line_idx, page_num in insertions: - md_lines.insert(line_idx, f"\n") - - with open(md_path, "w", encoding="utf-8") as f: - f.writelines(md_lines) - - -# ─── Child-process workers (top-level for pickling) ───────────────── - - -@worker -def _fast_path_worker(queue, pdf_path, output_dir, image_dir): - """Child process: pymupdf4llm extraction + page marker injection.""" - import pymupdf - import pymupdf4llm - - doc = pymupdf.open(pdf_path) - try: - md_text = pymupdf4llm.to_markdown( - doc, - write_images=True, - image_path=image_dir, - image_format="png", - ) - finally: - doc.close() - - full_md_path = os.path.join(output_dir, "full.md") - with open(full_md_path, "w", encoding="utf-8") as f: - f.write(md_text) - - _inject_page_markers_pymupdf(pdf_path, output_dir) - - img_count = len([n for n in os.listdir(image_dir) if n.endswith(".png")]) - queue.put( - { - "ok": True, - "md_chars": len(md_text), - "image_count": img_count, - } - ) +from shared.core.config import settings def parse_pdfs( @@ -220,7 +21,7 @@ def parse_pdfs( route = profile.route if profile else "standard" base_llm_paras.update({"doc_name": filename}) - # ── Atlas routing: bypass MinerU entirely, use PyMuPDF for per-page chunking ── + # ── Atlas routing: bypass MinerU entirely ── if profile and profile.doc_category == "atlas": logger.info(f"📐 Atlas detected, bypassing MinerU for {filename}") from app.services.document_parser.formats.atlas.parser import parse_atlas @@ -229,41 +30,22 @@ def parse_pdfs( pdf_path, output_dir, base_llm_paras, relative_root, profile=profile ) - # TODO: Re-enable fast path after thorough debugging. - # Conservative strategy: until the fast path (pymupdf4llm) is fully validated, - # all non-atlas PDFs are forced to MinerU (standard route) regardless of what - # DocProfiler recommends. The routing logic below is intentionally bypassed. - # - # Original fast-path block (keep for reference, do NOT delete): - # if route == "fast": - # logger.info(f"⚡ Fast path: extracting with pymupdf4llm for {filename}") - # - # os.makedirs(output_dir, exist_ok=True) - # image_dir = os.path.join(output_dir, "images") - # os.makedirs(image_dir, exist_ok=True) - # - # with stage_timer("pdf.extract.fast", filename=filename): - # result = run_in_child_process( - # _fast_path_worker, pdf_path, output_dir, image_dir, - # ) - # logger.info( - # f"⚡ Fast path done: {result['md_chars']} chars, " - # f"{result['image_count']} images" - # ) - # else: - # with stage_timer("pdf.extract.standard", filename=filename): - # parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) - # _inject_page_markers(output_dir) + # ── Oversized PDF: doc_agent → shard → parallel MinerU → merge → parse_md ── + if profile and profile.page_count > settings.MAX_PDF_PAGE_LIMIT: + logger.info( + f"📄 Oversized PDF: {profile.page_count} pages > " + f"{settings.MAX_PDF_PAGE_LIMIT} limit, entering shard pipeline" + ) + return _parse_oversized_pdf( + pdf_path, filename, output_dir, base_llm_paras, + profile=profile, relative_root=relative_root, s3_key=s3_key, + ) - logger.info( - f"🛡️ Conservative mode: forcing MinerU (standard) for {filename} [route={route}]" - ) + # ── Standard single-pass MinerU ── + logger.info(f"📄 Standard MinerU parse for {filename} [route={route}]") with stage_timer("pdf.extract.standard", filename=filename): parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) - # Inject page markers from MinerU layout.json - _inject_page_markers(output_dir) - logger.info("✅ PDF parsing step 1 complete: text extracted") with stage_timer("pdf.parse_md", filename=filename): @@ -274,3 +56,202 @@ def parse_pdfs( base_llm_paras=base_llm_paras, relative_root=relative_root, ) + + +def _parse_oversized_pdf( + pdf_path, filename, output_dir, base_llm_paras, + profile=None, relative_root=None, s3_key=None, +): + """Handle PDFs exceeding MinerU's page limit via shard-first hierarchy. + + Pipeline: + 1. DOC_AGENT → shard plan + TOC + 2. bin_pack → merged shards + 3. split_pdf (exclude TOC pages) + 4. MinerU per shard (parallel) + 5. **Per-shard heading prediction** (parallel) ← NEW + 6. Merge lines_with_heading + images + 7. parse_md Phase B (skip TOC detection + heading prediction) + """ + from concurrent.futures import ThreadPoolExecutor, as_completed + from dataclasses import dataclass + + from app.services.document_parser.formats.markdown.parser import ( + eval_md_headings, + merge_html_tables, + ) + from app.services.document_parser.formats.pdf.shard_merger import merge_images + from app.services.document_parser.formats.pdf.shard_splitter import ( + bin_pack_shards, + run_doc_agent, + split_pdf, + ) + + job_id = base_llm_paras.get("doc_name", filename) + + # 1. Run doc_agent to get full anatomy map (shard plan + TOC info) + with stage_timer("pdf.doc_agent", filename=filename): + anatomy = run_doc_agent(pdf_path, job_id=job_id, output_dir=output_dir) + + agent_shards = anatomy.shard_plan.shards + + # 2. Extract TOC info from anatomy for page exclusion and heading constraint + toc_pages: set[int] = set() + toc_hierarchies = None + if anatomy.toc_result and anatomy.toc_result.toc_pages: + toc_pages = set(anatomy.toc_result.toc_pages) + toc_hierarchies = anatomy.toc_hierarchies + logger.info( + f"📌 DOC_AGENT TOC detected: {len(toc_pages)} pages to exclude " + f"({sorted(toc_pages)}), " + f"{len(toc_hierarchies) if toc_hierarchies else 0} hierarchy regions" + ) + + # 3. Bin-pack agent shards to maximize MinerU page limit + merged_shards = bin_pack_shards(agent_shards, max_pages=settings.MAX_PDF_PAGE_LIMIT) + logger.info( + f"📦 Bin-packed {len(agent_shards)} agent shards → " + f"{len(merged_shards)} MinerU shards" + ) + for ms in merged_shards: + logger.info( + f" shard_{ms.shard_index}: pages {ms.page_start}-{ms.page_end} " + f"({ms.page_count} pages)" + ) + + # 4. Physically split PDF (exclude TOC pages if detected) + work_dir = os.path.join(output_dir, "_shards") + os.makedirs(work_dir, exist_ok=True) + with stage_timer("pdf.split", filename=filename): + shard_pdf_paths, _page_remap = split_pdf( + pdf_path, merged_shards, work_dir, + exclude_pages=toc_pages if toc_pages else None, + ) + + # 5. Parse each shard via MinerU (parallel) + shard_output_dirs: list[str | None] = [None] * len(shard_pdf_paths) + concurrency = settings.MINERU_SHARD_CONCURRENCY + + def _parse_single_shard(shard_idx, shard_pdf): + shard_out = os.path.join(work_dir, f"shard_{shard_idx}_output") + os.makedirs(shard_out, exist_ok=True) + shard_filename = ( + f"{os.path.splitext(filename)[0]}_shard{shard_idx}.pdf" + ) + logger.info( + f" 🔄 MinerU shard_{shard_idx}: parsing" + ) + parse_via_full(shard_pdf, shard_filename, shard_out, s3_key=None) + return shard_out + + with stage_timer( + "pdf.mineru_parallel", filename=filename, shard_count=len(shard_pdf_paths) + ): + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = { + executor.submit(_parse_single_shard, i, shard_pdf_path): i + for i, shard_pdf_path in enumerate(shard_pdf_paths) + } + for future in as_completed(futures): + idx = futures[future] + shard_output_dirs[idx] = future.result() + + # 6. Per-shard heading prediction (parallel) + @dataclass + class ShardHeadingResult: + shard_index: int + lines_with_heading: list[str] + heading_count: int + + smart_parse = base_llm_paras.get("smart_title_parse", True) + hierarchy_model_name = ( + base_llm_paras.get("hierarchy_model_name") + or base_llm_paras.get("model_name", settings.NORMOL_MODEL) + ) + + def _predict_shard_headings(shard_idx: int, shard_out_dir: str) -> ShardHeadingResult: + """Run full heading prediction pipeline on a single shard's full.md.""" + md_path = os.path.join(shard_out_dir, "full.md") + if not os.path.exists(md_path): + logger.warning(f"shard_{shard_idx}: full.md not found, returning empty") + return ShardHeadingResult(shard_index=shard_idx, lines_with_heading=[], heading_count=0) + + with open(md_path, "r", encoding="utf-8") as f: + md_lines = f.readlines() + md_lines = [line.strip() for line in md_lines if line.strip() != ""] + md_lines = merge_html_tables(md_lines) + + # TOC context: first TOC shared by all shards; subsequent TOCs assigned + # by page boundary. For simplicity, all TOCs are passed since pred_titles + # only matches headings actually present in this shard's content. + shard_toc = toc_hierarchies + + lines_with_heading = eval_md_headings( + md_lines, + source_type="md", + toc_hierarchies=shard_toc, + smart_parse=smart_parse, + model_name=hierarchy_model_name, + output_dir=shard_out_dir, + layout_json_path=( + os.path.join(shard_out_dir, "layout.json") + if os.path.exists(os.path.join(shard_out_dir, "layout.json")) + else None + ), + ) + + heading_count = sum(1 for line in lines_with_heading if line.startswith("#")) + logger.info( + f" ✅ shard_{shard_idx}: {heading_count} headings identified " + f"from {len(lines_with_heading)} lines" + ) + return ShardHeadingResult( + shard_index=shard_idx, + lines_with_heading=lines_with_heading, + heading_count=heading_count, + ) + + shard_heading_results: list[ShardHeadingResult | None] = [None] * len(shard_output_dirs) + + with stage_timer( + "pdf.shard_headings", filename=filename, shard_count=len(shard_output_dirs) + ): + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = { + executor.submit(_predict_shard_headings, i, shard_dir): i + for i, shard_dir in enumerate(shard_output_dirs) + if shard_dir is not None + } + for future in as_completed(futures): + idx = futures[future] + shard_heading_results[idx] = future.result() + + # 7. Merge: concatenate lines_with_heading (in shard order) + merge images + all_lines_with_heading: list[str] = [] + total_headings = 0 + for result in shard_heading_results: + if result is not None: + all_lines_with_heading.extend(result.lines_with_heading) + total_headings += result.heading_count + + logger.info( + f"📎 Merged {len(shard_heading_results)} shards: " + f"{len(all_lines_with_heading)} lines, {total_headings} headings" + ) + + with stage_timer("pdf.merge_images", filename=filename): + merge_images(shard_output_dirs, output_dir) + + logger.info("✅ Shard-first hierarchy complete, entering parse_md Phase B") + + # 8. parse_md Phase B only (skip TOC detection + heading prediction) + with stage_timer("pdf.parse_md", filename=filename): + return parse_md( + output_dir, + source_type="md", + base_llm_paras=base_llm_paras, + relative_root=relative_root, + lines_with_heading=all_lines_with_heading, + ) + + diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py new file mode 100644 index 000000000..55afa3ca0 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py @@ -0,0 +1,28 @@ +"""Merge MinerU image outputs from multiple shards into a single unified output.""" + +from __future__ import annotations + +import os +import shutil + +from loguru import logger + + +def merge_images(shard_dirs: list[str], target_dir: str) -> None: + """Copy all images from shard images/ dirs into target_dir/images/.""" + target_img_dir = os.path.join(target_dir, "images") + os.makedirs(target_img_dir, exist_ok=True) + total = 0 + for shard_dir in shard_dirs: + if shard_dir is None: + continue + img_dir = os.path.join(shard_dir, "images") + if not os.path.isdir(img_dir): + continue + for fname in os.listdir(img_dir): + src = os.path.join(img_dir, fname) + dst = os.path.join(target_img_dir, fname) + if os.path.isfile(src): + shutil.copy2(src, dst) + total += 1 + logger.info(f"Merged {total} images → {target_img_dir}") diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py new file mode 100644 index 000000000..ebf599145 --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/pdf/shard_splitter.py @@ -0,0 +1,166 @@ +"""PDF shard splitting: doc_agent integration + bin-packing + physical split.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import pymupdf +from loguru import logger + +if TYPE_CHECKING: + from app.services.document_agent.manifest import PageAnatomyMap, Shard + + +@dataclass +class MergedShard: + """A contiguous page range within MinerU's per-request page limit.""" + + shard_index: int + page_start: int # 1-based inclusive + page_end: int # 1-based inclusive + + @property + def page_count(self) -> int: + return self.page_end - self.page_start + 1 + + @property + def page_offset(self) -> int: + """Offset to add to MinerU's 0-based page_idx to get original page_idx.""" + return self.page_start - 1 + + +def run_doc_agent( + pdf_path: str, job_id: str, output_dir: str +) -> "PageAnatomyMap": + """Run doc_agent ProfileCoordinator and return the full anatomy map. + + Returns the complete PageAnatomyMap so callers can access TOC info + (toc_result.toc_pages, toc_hierarchies) in addition to the shard plan. + + Raises RuntimeError if the agent fails or produces no shards. + """ + from app.services.document_agent.coordinator import ProfileCoordinator + + agent_output_dir = os.path.join(output_dir, "_doc_agent") + os.makedirs(agent_output_dir, exist_ok=True) + + coordinator = ProfileCoordinator( + pdf_path=pdf_path, + job_id=job_id, + output_dir=agent_output_dir, + ) + anatomy = coordinator.run() + + if not anatomy.shard_plan.enabled or not anatomy.shard_plan.shards: + raise RuntimeError( + f"Doc agent did not produce a valid shard plan for {job_id}" + ) + + shards = anatomy.shard_plan.shards + logger.info( + f"📋 Doc agent: {len(shards)} shards via {anatomy.shard_plan.reason}" + ) + return anatomy + + +def bin_pack_shards( + agent_shards: list["Shard"], + max_pages: int, +) -> list[MergedShard]: + """Greedy left-to-right bin-packing: merge adjacent agent shards up to max_pages.""" + if not agent_shards: + return [] + + merged: list[MergedShard] = [] + cur_start = agent_shards[0].page_start + cur_end = agent_shards[0].page_end + + for shard in agent_shards[1:]: + if shard.page_end - cur_start + 1 <= max_pages: + cur_end = shard.page_end + else: + merged.append( + MergedShard(len(merged), page_start=cur_start, page_end=cur_end) + ) + cur_start = shard.page_start + cur_end = shard.page_end + + merged.append( + MergedShard(len(merged), page_start=cur_start, page_end=cur_end) + ) + return merged + + +def split_pdf( + pdf_path: str, + shards: list[MergedShard], + work_dir: str, + exclude_pages: set[int] | None = None, +) -> tuple[list[str], dict[int, int] | None]: + """Physically split PDF into sub-PDFs using PyMuPDF. + + Args: + pdf_path: Path to the source PDF. + shards: Merged shard ranges to extract. + work_dir: Directory for temporary shard PDFs. + exclude_pages: Optional set of 1-based page numbers to strip + (e.g. TOC pages detected by DOC_AGENT). + + Returns: + (shard_paths, page_remap) + - shard_paths: one temp PDF path per shard. + - page_remap: when pages are excluded, maps each shard's local + 0-based page index to the original 1-based page number. + ``None`` when no pages are excluded. + """ + doc = pymupdf.open(pdf_path) + paths: list[str] = [] + page_remap: dict[int, int] | None = None + + if exclude_pages: + page_remap = {} + logger.info( + f"📌 Excluding {len(exclude_pages)} pages from PDF: " + f"{sorted(exclude_pages)}" + ) + + try: + global_new_idx = 0 # running counter across all shards + for shard in shards: + sub_doc = pymupdf.open() + shard_included = 0 + for page_num in range(shard.page_start, shard.page_end + 1): + if exclude_pages and page_num in exclude_pages: + continue + sub_doc.insert_pdf( + doc, + from_page=page_num - 1, + to_page=page_num - 1, + ) + if page_remap is not None: + page_remap[global_new_idx] = page_num + global_new_idx += 1 + shard_included += 1 + + shard_path = os.path.join(work_dir, f"shard_{shard.shard_index}.pdf") + if shard_included > 0: + sub_doc.save(shard_path) + paths.append(shard_path) + else: + logger.warning( + f" ⚠️ shard_{shard.shard_index}: all pages excluded, skipping" + ) + sub_doc.close() + + excluded_in_shard = shard.page_count - shard_included + logger.info( + f" ✂️ shard_{shard.shard_index}: " + f"pages {shard.page_start}-{shard.page_end} " + f"({shard_included} included" + f"{f', {excluded_in_shard} excluded' if excluded_in_shard else ''})" + ) + finally: + doc.close() + return paths, page_remap diff --git a/apps/worker/app/services/document_parser/orchestration/parse_session.py b/apps/worker/app/services/document_parser/orchestration/parse_session.py index 7e1141a84..b90323553 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -109,20 +109,37 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: f"ℹ️ VLM rejected atlas for {parse_input.filename}, routing as generic" ) - pdf_page_limit = settings.MAX_PDF_PAGE_LIMIT - if profile.file_type == "pdf" and profile.page_count > pdf_page_limit: - raise ValidationException( - user_message=( - f"Document too large: {profile.page_count} pages exceeds the {pdf_page_limit}-page limit. " - "Please split the document and upload in smaller batches." - ), - violations=[ - { + if profile.file_type == "pdf" and profile.page_count > settings.MAX_PDF_PAGE_LIMIT: + if profile.page_count > settings.OVERSIZED_PDF_SOFT_LIMIT: + raise ValidationException( + user_message=( + f"This document has {profile.page_count} pages. Processing ultra-long " + f"documents (over {settings.OVERSIZED_PDF_SOFT_LIMIT} pages) requires " + "dedicated resources. Please contact our support team for assistance." + ), + violations=[{ "field": "page_count", - "description": f"PDF has {profile.page_count} pages, limit is {pdf_page_limit}", - } - ], - ) + "description": ( + f"PDF has {profile.page_count} pages, " + f"soft limit is {settings.OVERSIZED_PDF_SOFT_LIMIT}" + ), + }], + ) + if not settings.OVERSIZED_PDF_SHARD_ENABLED: + raise ValidationException( + user_message=( + f"Document has {profile.page_count} pages, exceeding the " + f"{settings.MAX_PDF_PAGE_LIMIT}-page limit. Please split the " + "document into smaller parts and upload them separately." + ), + violations=[{ + "field": "page_count", + "description": ( + f"PDF has {profile.page_count} pages, " + f"limit is {settings.MAX_PDF_PAGE_LIMIT}" + ), + }], + ) if profile.doc_category == "atlas": filename, internal_output_filename, relative_root, full_output_dir = ( diff --git a/apps/worker/app/services/document_parser/structure/heading_candidates.py b/apps/worker/app/services/document_parser/structure/heading_candidates.py index f9d060349..0a7596aed 100644 --- a/apps/worker/app/services/document_parser/structure/heading_candidates.py +++ b/apps/worker/app/services/document_parser/structure/heading_candidates.py @@ -15,11 +15,13 @@ HEADING_COLUMNS = Index(["id", "heading", "level", "reason"]) -def get_max_lvl(code_str: str): - match = re.search(r"\[([^]]+)]", code_str) - if not match: - return "Sure" +def get_max_lvl(code_str: str) -> int: + """Extract the maximum hierarchy depth from a POS trigger code string. + ``code_str`` is always ``str(pos_code)`` where *pos_code* is a list of + integers, so the ``[…]`` bracket match is guaranteed. + """ + match = re.search(r"\[([^]]+)]", code_str) nums = [int(item.strip()) for item in match.group(1).split(",")] max_value = int(max(nums)) return max_value if max_value > 1 else -2 @@ -86,7 +88,14 @@ def judge_by_conditions( return pos_triggered_code -def remove_by_conditions(text, include_punc: bool = False): +def remove_by_conditions(text, *, include_punc: bool = False): + """Evaluate negative (non-heading) conditions against *text*. + + ``include_punc`` controls the end-of-line punctuation rule. It is + intentionally **off** during initial scanning so that lines remain heading candidates. + The punctuation check is enabled only during the ``judge_negs`` second pass (after merges + may have altered heading text). + """ neg_conditions = [ r"^\d{3,}", r"(?i)(^https?://\S+|^www\.\S+|^P\.S|^\b\d{0,2}\s*(?:a\.m|p\.m)\b)", @@ -101,7 +110,7 @@ def remove_by_conditions(text, include_punc: bool = False): r"|left|right|begin|end|overline|underline|hat|vec|tilde)\b" r")" ), - r"^0\.\d+[\u4e00-\u9fa5A-Za-z\S]*", + r"^0\.\d+\S*", r"^\d*\.\d+$", r"[。！；].+", ( @@ -119,11 +128,15 @@ def remove_by_conditions(text, include_punc: bool = False): for regex in neg_conditions: neg_triggered_code.append(1 if re.search(regex, text) else 0) + # End-of-line punctuation — only checked during judge_negs second pass. if include_punc: neg_triggered_code.append(1 if re.search(r"[.,;，。；]$", text) else 0) else: neg_triggered_code.append(0) + MAX_HEADING_TOKENS = 10 + neg_triggered_code.append(1 if count_cn_en(text) > MAX_HEADING_TOKENS else 0) + return neg_triggered_code @@ -136,10 +149,13 @@ def md_heading_match(line, as_is: bool = True): return (line, level) if as_is else (line.lstrip("#").strip(), level) +# Pre-compute zero-filled code arrays so non-heading lines get correct-width reason strings. +_ZERO_POS_CODE = judge_by_conditions("") +_ZERO_NEG_CODE = remove_by_conditions("") + + def filter_markdown_headings( md_lines: list[str], - num_pos: int = 17, - num_neg: int = 7, layout_json_path: str | None = None, ) -> pd.DataFrame: meta_ctx = None @@ -159,9 +175,7 @@ def filter_markdown_headings( if _is_non_heading_markdown_line(line): est_level = -1 - zero_pos_code = [0] * num_pos - zero_neg_code = [0] * num_neg - reason = f"POS {zero_pos_code} NEG {zero_neg_code}" + reason = f"POS {_ZERO_POS_CODE} NEG {_ZERO_NEG_CODE}" if meta_ctx: reason += " META [0, 0, 0]" line = "Figure/Image" @@ -242,8 +256,8 @@ def postprocess_headings(df: pd.DataFrame, task: str, max_depth: int = -1) -> pd if task == "merge_continuous": return _merge_continuous_non_headings(df) - if task == "merge_short" or task == "collapse": - return _collapse_heading_groups(df, task) + if task == "merge_short": + return _merge_short_heading_groups(df) return df @@ -377,7 +391,7 @@ def _merge_continuous_non_headings(df: pd.DataFrame) -> pd.DataFrame: return pd.DataFrame(denoised_rows, columns=HEADING_COLUMNS) -def _collapse_heading_groups(df: pd.DataFrame, task: str) -> pd.DataFrame: +def _merge_short_heading_groups(df: pd.DataFrame) -> pd.DataFrame: group_to_indices = defaultdict(list) for index, row in df.iterrows(): level = row["level"] @@ -387,24 +401,22 @@ def _collapse_heading_groups(df: pd.DataFrame, task: str) -> pd.DataFrame: checked_pairs = set() for _, indices in group_to_indices.items(): - _collapse_recursive(df, task, indices, merge_threshold=3, checked_pairs=checked_pairs) - - if task == "merge_short": - drop_between = df.index[ - df["reason"].astype(str).str.startswith("Merged into", na=False) - ].tolist() - if drop_between: - logger.debug( - f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows" - ) - df.drop(drop_between, inplace=True) - df.reset_index(drop=True, inplace=True) + _merge_short_recursive(df, indices, merge_threshold=3, checked_pairs=checked_pairs) + + drop_between = df.index[ + df["reason"].astype(str).str.startswith("Merged into", na=False) + ].tolist() + if drop_between: + logger.debug( + f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows" + ) + df.drop(drop_between, inplace=True) + df.reset_index(drop=True, inplace=True) return df -def _collapse_recursive( +def _merge_short_recursive( df: pd.DataFrame, - task: str, indices: list[int], merge_threshold: int = 3, checked_pairs: set[tuple[int, int]] | None = None, @@ -423,18 +435,11 @@ def _collapse_recursive( between = df.loc[index + 1 : next_index - 1] current_text = df.at[index, "heading"].strip() - next_text = df.at[next_index, "heading"].strip() - if task == "merge_short" and len(between) > 0: + if len(between) > 0: _merge_short_between_headings( df, between, index, next_index, current_text, merge_threshold ) - elif task == "collapse" and len(between) == 0: - logger.debug( - f"⚠️ Empty between i={current_text[:15]}, j={next_text[:15]} => set i.level=-1, j.level=Not Sure" - ) - df.at[index, "level"] = -2 - df.at[next_index, "level"] = -2 sub_between = between[between["level"] != -1] code_to_sub = defaultdict(list) @@ -445,7 +450,7 @@ def _collapse_recursive( code_to_sub[(level, reason)].append(row_index) for _, sub_indices in code_to_sub.items(): - _collapse_recursive(df, task, sub_indices, merge_threshold, checked_pairs) + _merge_short_recursive(df, sub_indices, merge_threshold, checked_pairs) def _merge_short_between_headings( diff --git a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py index 6f314416b..bea5e8540 100644 --- a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py +++ b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py @@ -2,8 +2,6 @@ from __future__ import annotations import os -import re -from collections import Counter from collections.abc import Callable from typing import Any @@ -13,8 +11,6 @@ from app.services.document_parser.support.text_helpers import count_cn_en, truncate_text_by_tokens from loguru import logger -from shared.core.exceptions.domain_exceptions import WorkerHandlingException - PLACEHOLDER_REASON = "__PLACEHOLDER__" HierarchyJudge = Callable[..., list[dict[str, Any]]] @@ -22,332 +18,17 @@ SaveIntermediateCsv = Callable[[pd.DataFrame, str | None, str], None] -def build_level_mapping( - df: pd.DataFrame, origin_lvls: list[int], mode: str = "max" -) -> tuple[pd.DataFrame, dict[str, dict[str, Any]]]: - mapped_df = df.copy() - mapped_df["origin_level"] = origin_lvls - - mapping = mapped_df.groupby("reason")["level"].apply(list).to_dict() - - processed_mapping: dict[str, dict[str, Any]] = {} - for reason, lvls in mapping.items(): - positive_lvls = [lvl for lvl in lvls if lvl > -1] - counts = Counter(lvls) - - if not positive_lvls: - mapped_lvl = -1 - elif mode == "max": - mapped_lvl = max(positive_lvls) - elif mode == "freq": - mapped_lvl = counts.most_common(1)[0][0] - else: - raise WorkerHandlingException( - internal_message=f"wrong input mode: {mode}. Must be 'max' or 'freq'" - ) - - processed_mapping[reason] = { - "lvls": lvls, - "positive_lvls": positive_lvls, - "freqs": dict(counts), - "mapped_lvl": mapped_lvl, - } - return mapped_df, processed_mapping - - -def execute_level_mapping( - df: pd.DataFrame, mapping: dict[str, dict[str, Any]] -) -> pd.DataFrame: - def map_row(row: pd.Series) -> int: - reason = row["reason"] - if reason in mapping: - return int(mapping[reason]["mapped_lvl"]) - return int(row["level"]) - - mapped_df = df.copy() - origin_est_lvls = mapped_df["level"].tolist() - mapped_df["level"] = mapped_df.apply(map_row, axis=1) - mapped_df["origin_level"] = origin_est_lvls - return mapped_df - - -def extract_non_neg_code(reason_str: str) -> str: - """Extract the non-NEG code from a heading reason string.""" - if not reason_str or not isinstance(reason_str, str): - return "" - neg_match = re.search(r"\s*NEG\s*\[[^\]]*\]", reason_str) - if neg_match: - before_neg = reason_str[: neg_match.start()] - after_neg = reason_str[neg_match.end() :] - return (before_neg + after_neg).strip() - return reason_str.strip() - - -def build_non_neg_mapping(lvl_mapping: dict[str, dict[str, Any]]) -> dict[str, int]: - non_neg_levels: dict[str, list[int]] = {} - for reason, info in lvl_mapping.items(): - non_neg_code = extract_non_neg_code(reason) - mapped_lvl = int(info.get("mapped_lvl", -1)) - if non_neg_code: - non_neg_levels.setdefault(non_neg_code, []).append(mapped_lvl) - - non_neg_mapping: dict[str, int] = {} - for non_neg_code, levels in non_neg_levels.items(): - positive_levels = [lvl for lvl in levels if lvl > -1] - if positive_levels: - level_counts = Counter(positive_levels) - non_neg_mapping[non_neg_code] = level_counts.most_common(1)[0][0] - else: - non_neg_mapping[non_neg_code] = -1 - - return non_neg_mapping - - -def handle_unseen_codes( - df: pd.DataFrame, - level_dfs: list[pd.DataFrame], - lvl_mapping: dict[str, dict[str, Any]], - output_dir: str | None = None, - window_half_size: int = 10, - strategy: str = "double_mapping", -) -> dict[str, dict[str, Any]]: - """Extend first-chunk reason mapping to reason codes only seen in later chunks.""" - - def extract_reason_signature(reason: str) -> str: - return reason.strip() if reason else "" - - def has_neg_signal(reason_str: str) -> bool: - if not reason_str or not isinstance(reason_str, str): - return False - neg_match = re.search(r"NEG\s*\[([^\]]*)\]", reason_str) - if not neg_match: - return False - neg_content = neg_match.group(1) - try: - nums = [int(x.strip()) for x in neg_content.split(",") if x.strip()] - return any(x >= 1 for x in nums) - except Exception: - return False - - def build_context_window( - target_idx: int, known_codes_set: set[str], total_rows: int, half_size: int = 10 - ) -> dict[str, Any]: - min_start = max(0, target_idx - half_size) - min_end = min(total_rows - 1, target_idx + half_size) - - start_idx = min_start - end_idx = min_end - - found_known_above = False - found_known_below = False - known_positions: list[int] = [] - - for index in range(start_idx, target_idx): - reason = df.iloc[index].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_above = True - known_positions.append(index) - - for index in range(target_idx + 1, end_idx + 1): - reason = df.iloc[index].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_below = True - known_positions.append(index) - - if not found_known_above and min_start > 0: - search_idx = min_start - 1 - while search_idx >= 0: - reason = df.iloc[search_idx].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_above = True - known_positions.append(search_idx) - start_idx = search_idx - break - search_idx -= 1 - - if not found_known_below and min_end < total_rows - 1: - search_idx = min_end + 1 - while search_idx < total_rows: - reason = df.iloc[search_idx].get("reason", "") - sig = extract_reason_signature(reason) - if sig in known_codes_set: - found_known_below = True - known_positions.append(search_idx) - end_idx = search_idx - break - search_idx += 1 - - return { - "start": start_idx, - "end": end_idx, - "found_known": found_known_above or found_known_below, - "known_positions": known_positions, - } - - non_neg_mapping = build_non_neg_mapping(lvl_mapping) - known_codes = set(lvl_mapping.keys()) - - all_codes_in_full: dict[str, dict[str, Any]] = {} - for seg_idx, seg_df in enumerate(level_dfs): - for _, row in seg_df.iterrows(): - reason = row.get("reason", "") - sig = extract_reason_signature(reason) - if not sig or sig == PLACEHOLDER_REASON: - continue - if sig not in all_codes_in_full: - all_codes_in_full[sig] = { - "first_seg": seg_idx, - "first_id": row.get("id", 0), - "reason": reason, - } - - unseen_codes: dict[str, dict[str, Any]] = {} - unseen_neg_filtered: dict[str, dict[str, Any]] = {} - for sig, info in all_codes_in_full.items(): - if sig in known_codes: - continue - if has_neg_signal(info["reason"]): - unseen_neg_filtered[sig] = info - else: - unseen_codes[sig] = info - - logger.info( - f"Unseen codes total: {len(unseen_codes) + len(unseen_neg_filtered)}, " - f"NEG filtered: {len(unseen_neg_filtered)}, to process: {len(unseen_codes)}" - ) - - for sig in unseen_neg_filtered: - lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NEG_FILTERED"} - - if unseen_codes: - if strategy == "double_mapping": - fallback_success = 0 - fallback_failed = 0 - failed_codes = [] - for sig in unseen_codes: - non_neg_code = extract_non_neg_code(sig) - if non_neg_code in non_neg_mapping: - mapped_level = non_neg_mapping[non_neg_code] - lvl_mapping[sig] = { - "mapped_lvl": mapped_level, - "note": f"NON_NEG_FALLBACK from '{non_neg_code}'", - } - fallback_success += 1 - else: - lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NO_MATCH_FALLBACK"} - fallback_failed += 1 - failed_codes.append( - f"'{non_neg_code}' (from '{sig[:60]}...')" - if len(sig) > 60 - else f"'{non_neg_code}' (from '{sig}')" - ) - - logger.info( - f"Double mapping result: success={fallback_success}, failed={fallback_failed}" - ) - if failed_codes: - logger.warning( - f"Failed codes (non_neg not in mapping): {failed_codes[:5]}" - f"{'...' if len(failed_codes) > 5 else ''}" - ) - - elif strategy == "window_llm" and output_dir: - total_rows = len(df) - windows: list[dict[str, Any]] = [] - for sig, info in unseen_codes.items(): - first_id = info["first_id"] - first_seg = info["first_seg"] - df_indices = df.index[df["id"] == first_id].tolist() - if df_indices: - first_df_idx = df_indices[0] - window_info = build_context_window( - first_df_idx, known_codes, total_rows, window_half_size - ) - windows.append( - { - "code": sig, - "first_id": first_id, - "first_seg": first_seg, - "start": window_info["start"], - "end": window_info["end"], - "found_known": window_info["found_known"], - } - ) - - sorted_windows = sorted(windows, key=lambda window: window["start"]) - merged_windows: list[dict[str, Any]] = [] - current_window: dict[str, Any] | None = None - - for window in sorted_windows: - if current_window is None: - current_window = { - "start": window["start"], - "end": window["end"], - "codes": [window["code"]], - "segments": [window["first_seg"]], - } - elif window["start"] <= current_window["end"]: - current_window["end"] = max(current_window["end"], window["end"]) - current_window["codes"].append(window["code"]) - current_window["segments"].append(window["first_seg"]) - else: - merged_windows.append(current_window) - current_window = { - "start": window["start"], - "end": window["end"], - "codes": [window["code"]], - "segments": [window["first_seg"]], - } - - if current_window: - merged_windows.append(current_window) - - windows_dir = os.path.join(output_dir, "merged_windows") - os.makedirs(windows_dir, exist_ok=True) - - unseen_codes_set = set(unseen_codes.keys()) - unseen_neg_set = set(unseen_neg_filtered.keys()) - - for index, merged_window in enumerate(merged_windows): - window_df = df.iloc[ - merged_window["start"] : merged_window["end"] + 1 - ].copy() - - def get_code_status(row: pd.Series) -> str: - reason = row.get("reason", "") - sig = extract_reason_signature(reason) - if not sig: - return "" - if sig in unseen_codes_set: - return "UNSEEN_TARGET" - if sig in unseen_neg_set: - return "NEG_TO_NEGATIVE_ONE" - if sig in known_codes: - return "KNOWN" - return "" - - window_df["code_status"] = window_df.apply(get_code_status, axis=1) - window_path = os.path.join( - windows_dir, - f"window_{index + 1:02d}_rows_" - f"{merged_window['start']}-{merged_window['end']}.csv", - ) - window_df.to_csv(window_path, index=False, encoding="utf-8-sig") - - logger.debug( - f"Window LLM: {len(merged_windows)} windows created in {windows_dir}" - ) - - return lvl_mapping - - def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: - """Collapse consecutive body rows into placeholder rows before LLM chunking.""" + """Collapse consecutive body rows into placeholder rows before LLM chunking. + + The output DataFrame has columns [id, heading, reason]. The ``level`` + column is intentionally NOT forwarded to the LLM — preliminary estimates + were found to mislead the model more often than they helped. The naive- + stage body-text detection (level == -1) is still used here to decide which + rows become placeholders vs candidates. + """ if df is None or len(df) == 0: - return pd.DataFrame(columns=["id", "heading", "level", "reason"]) + return pd.DataFrame(columns=["id", "heading", "reason"]) rows: list[dict[str, Any]] = [] index = 0 @@ -376,7 +57,6 @@ def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: { "id": f"{start_id}-{end_id}", "heading": f"[{run_length} BODY LINES]", - "level": "-", "reason": PLACEHOLDER_REASON, } ) @@ -387,17 +67,12 @@ def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame: { "id": int(row["id"]), "heading": str(row["heading"]), - "level": ( - int(lvl_int) - if lvl_int is not None and lvl_int != -2 - else "Not Sure" - ), "reason": str(row.get("reason", "") or ""), } ) index += 1 - return pd.DataFrame(rows, columns=["id", "heading", "level", "reason"]) + return pd.DataFrame(rows, columns=["id", "heading", "reason"]) def split_heading_table( @@ -413,7 +88,8 @@ def split_heading_table( current_rows: list[list[Any]] = [] current_len = 0 for _, row in working_df.iterrows(): - row_filtered = row.drop(labels=["reason"], errors="ignore") + # Drop internal-only columns before measuring token length + row_filtered = row.drop(labels=["reason", "level"], errors="ignore") row_len = sum(count_cn_en(str(value)) for value in row_filtered.values) if current_len + row_len > threshold and current_rows: @@ -478,112 +154,68 @@ def execute_llm_heading_hierarchy( f"threshold={prompt_limt} | max_start={max_len}" ) - basic_idx = 0 - for idx, chunk in enumerate(level_dfs): - if (chunk["reason"].astype(str) != PLACEHOLDER_REASON).any(): - basic_idx = idx - break - basic_df = level_dfs[basic_idx] - if basic_idx != 0: - logger.info( - f"smart parse => promoted chunk {basic_idx} as basic_df " - f"(chunks 0..{basic_idx - 1} contain only placeholders)" - ) - full_preds: pd.DataFrame | None = None try: with stage_timer( "heading.hierarchy_llm", chunk_count=len(level_dfs), - base_chunk_rows=len(basic_df), compact_enabled=compact_enabled, source_row_count=len(raw_preds), model_name=model_name, ): - logger.debug("smart parse => interpreting hierarchy patterns...") - df4llm = basic_df.drop(columns=["reason"]).copy() - df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) - logger.debug(f"DataFrame transformation completed, rows: {len(df4llm)}") - - layout_res = hierarchy_judge( - df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings" - ) - - layout_level_by_id: dict[Any, Any] = {} - if isinstance(layout_res, list): - for item in layout_res: - if isinstance(item, dict) and "id" in item and "level" in item: - layout_level_by_id[item["id"]] = item["level"] - - def level_for(row_id: Any) -> Any: - if row_id in layout_level_by_id: - return layout_level_by_id[row_id] - try: - return layout_level_by_id.get(int(row_id), -1) - except (TypeError, ValueError): - return -1 - - base_preds = ( - basic_df[["id", "heading", "reason"]].copy().reset_index(drop=True) - ) - base_preds.insert(2, "level", base_preds["id"].map(level_for)) - save_intermediate_csv( - base_preds, output_dir, f"preds_3_llm_base{csv_suffix}" - ) - + # ── Per-chunk independent LLM calls ── llm_levels: dict[int, Any] = {} - for _, row in base_preds.iterrows(): - row_id = row["id"] - if isinstance(row_id, bool): + + for chunk_idx, chunk_df in enumerate(level_dfs): + # Skip chunks that contain only placeholders + non_placeholder_mask = chunk_df["reason"].astype(str) != PLACEHOLDER_REASON + if not non_placeholder_mask.any(): + logger.debug( + f"smart parse => chunk {chunk_idx}: all placeholders, skipping" + ) continue - if isinstance(row_id, int): - llm_levels[row_id] = row["level"] - if len(level_dfs) > 1: - placeholder_mask_base = base_preds["reason"].eq(PLACEHOLDER_REASON) - figure_mask_base = base_preds["heading"].eq("Figure/Image") - exclude_mask_base = placeholder_mask_base | figure_mask_base - base_preds_for_mapping = base_preds[~exclude_mask_base].copy() - base_origin_for_mapping = basic_df.loc[ - ~exclude_mask_base.values, "level" - ].tolist() + df4llm = chunk_df.drop(columns=["reason", "level"], errors="ignore").copy() + df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm) - base_preds_for_mapping, lvl_mapping = build_level_mapping( - base_preds_for_mapping, base_origin_for_mapping, mode="freq" + logger.info( + f"smart parse => chunk {chunk_idx}/{len(level_dfs)}: " + f"sending {len(df4llm)} rows to LLM" ) - logger.debug( - f"mapping development finished: {len(lvl_mapping)} rules " - f"(placeholders and Figure/Image excluded)" + chunk_result = hierarchy_judge( + df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings" ) - logger.debug( - f"mapping dataframe to levels across {len(level_dfs)} chunks..." - ) - lvl_mapping = handle_unseen_codes( - preds_for_llm, level_dfs, lvl_mapping, output_dir + if isinstance(chunk_result, list): + for item in chunk_result: + if isinstance(item, dict) and "id" in item and "level" in item: + try: + llm_levels[int(item["id"])] = item["level"] + except (TypeError, ValueError): + pass + + # Save per-chunk intermediate CSV + chunk_preds = ( + chunk_df[["id", "heading", "reason"]].copy().reset_index(drop=True) ) - - for level_df in level_dfs: - placeholder_mask_chunk = level_df["reason"].eq(PLACEHOLDER_REASON) - figure_mask_chunk = level_df["heading"].eq("Figure/Image") - exclude_mask_chunk = placeholder_mask_chunk | figure_mask_chunk - non_excluded = level_df[~exclude_mask_chunk].copy() - if not non_excluded.empty: - non_excluded = execute_level_mapping(non_excluded, lvl_mapping) - for _, row in non_excluded.iterrows(): - row_id = row["id"] - if isinstance(row_id, bool): - continue - if isinstance(row_id, int): - llm_levels[row_id] = row["level"] - logger.info( - f"multi-chunk mapping produced {len(llm_levels)} id->level entries" + chunk_preds.insert( + 2, "level", + chunk_preds["id"].map( + lambda rid: llm_levels.get( + int(rid) if not isinstance(rid, str) or rid.isdigit() else -1, -1 + ) + ), ) - else: - logger.info( - "single chunk - skipping reason-code mapping, using LLM output directly" + save_intermediate_csv( + chunk_preds, output_dir, + f"preds_llm{csv_suffix}_{chunk_idx}" ) + logger.info( + f"smart parse => per-chunk LLM produced {len(llm_levels)} " + f"id->level entries across {len(level_dfs)} chunks" + ) + full_preds = raw_preds.copy()[["id", "heading", "level", "reason"]] def resolve_level(row_id: Any) -> int: @@ -596,11 +228,8 @@ def resolve_level(row_id: Any) -> int: return int(level) except (TypeError, ValueError): return -1 - + full_preds["level"] = full_preds["id"].map(resolve_level).astype(int) - save_intermediate_csv( - full_preds, output_dir, f"preds_4_llm_final{csv_suffix}" - ) except Exception as exc: logger.warning( @@ -608,3 +237,4 @@ def resolve_level(row_id: Any) -> int: ) full_preds = fallback_hierarchy(raw_preds.copy()) return full_preds + diff --git a/apps/worker/app/services/document_parser/structure/layout_parser.py b/apps/worker/app/services/document_parser/structure/layout_parser.py index 47f0650af..78c9f56ea 100755 --- a/apps/worker/app/services/document_parser/structure/layout_parser.py +++ b/apps/worker/app/services/document_parser/structure/layout_parser.py @@ -9,8 +9,6 @@ postprocess_headings, ) from app.services.document_parser.structure.heading_llm_executor import ( - build_level_mapping, - execute_level_mapping, execute_llm_heading_hierarchy, ) from app.services.document_parser.structure.heading_tree import ( @@ -121,20 +119,16 @@ def format_toc_context_for_llm(toc_context) -> str: formatted_blocks.append("- No TOC entries available") continue - for entry in toc_entries: - if not isinstance(entry, dict): - continue - - heading = str(entry.get("heading", "")).strip().replace("\n", " ") - if not heading: - continue + # Dynamically generate MD table from JSON list (more token-efficient) + from app.services.document_agent.tools.vlm_toc_extractor import ( + build_toc_with_level_md, + ) - level = entry.get("level") - line_id = entry.get("id") - if isinstance(level, int): - formatted_blocks.append(f"- level {level} | id {line_id} | {heading}") - else: - formatted_blocks.append(f"- id {line_id} | {heading}") + md_table = build_toc_with_level_md(toc_entries) + if md_table: + formatted_blocks.append(md_table) + else: + formatted_blocks.append("- No TOC entries available") return "\n".join(formatted_blocks) @@ -325,12 +319,32 @@ def _compute_zone_boundaries(toc_hierarchies, coordinate_mode="post_removal"): def _resolve_first_toc_boundary(toc_hierarchies=None, first_toc_ele_num=None): - """Resolve the earliest available first-TOC boundary across coordinate sources.""" + """Resolve the earliest available first-TOC boundary across coordinate sources. + + Page-based TOC boundaries (``toc_range_unit == "page"``, produced by + DOC_AGENT/VLM for PDF/PPT) are in *page numbers*, NOT line/element IDs. + They must NOT be used for pre-TOC row removal because ``raw_preds["id"]`` + are line indices that restart from 0 in each shard. For these documents + the DOC_AGENT has already handled shard splitting around TOC pages. + """ toc_range_start = None + toc_unit = None if toc_hierarchies: first_range = toc_hierarchies[0].get("toc_range") if first_range and len(first_range) == 2: toc_range_start = first_range[0] + toc_unit = toc_hierarchies[0].get("toc_range_unit") + + # Page-based coordinates cannot be compared against line/element IDs. + if toc_unit == "page": + if first_toc_ele_num is not None: + # DOCX fallback: element-based boundary is safe to use. + return first_toc_ele_num + logger.debug( + "📌 Skipping pre-TOC removal: TOC uses page-based coordinates " + "(DOC_AGENT already handled shard boundaries)" + ) + return None candidates = [ value for value in (toc_range_start, first_toc_ele_num) if value is not None @@ -575,53 +589,33 @@ def _process_single_zone(zone_idx, zone_start, zone_end, zone_toc): f"📌 Spliced {len(pre_toc_rows)} pre-TOC lines back into predictions" ) - # Save heading_preds as preds_5 - save_intermediate_csv(heading_preds, output_dir, "preds_5_final_output") + # Save final heading predictions + save_intermediate_csv(heading_preds, output_dir, "preds_final") return heading_preds def est_hierarchies_naive(raw_preds, proceed_smart=True, output_dir=None): - """Detect hierarchies by non-LLM + """Regex-only heading hierarchy estimation (LLM fallback). - Args: - raw_preds: raw data - proceed_smart: whether to proceed with smart parsing - output_dir: output directory, used to save intermediate results CSV - """ - logger.debug("🚀 non-llm parsing => recursive processing") - save_preds = raw_preds.copy() - - heading_preds = postprocess_headings(raw_preds, task="collapse") - save_preds.insert( - save_preds.columns.get_loc("level") + 1, - "lvl_cola", - heading_preds["level"].tolist(), - ) + When used as the primary pipeline step (before LLM), this function + determines the initial candidate/body split that ``compact_for_llm`` + relies on (level > -1 → candidate, -1 → body text). - heading_preds = postprocess_headings(heading_preds, task="judge_negs") - save_preds.insert( - save_preds.columns.get_loc("lvl_cola") + 1, - "lvl_neg", - heading_preds["level"].tolist(), - ) - save_preds["reason"] = heading_preds["reason"] + When used as a **fallback** after LLM failure, the returned levels + must be usable for tree construction. Single-level POS matches + (``get_max_lvl`` returns -2 for patterns like) + are normalized to level 1 so the output forms a valid hierarchy. - # mapping based on freq - if not proceed_smart: - heading_preds["level"] = heading_preds["level"].map( - lambda x: -1 if x == -2 else x - ) - heading_preds, lvl_mapping = build_level_mapping( - heading_preds, heading_preds["level"].tolist(), mode="freq" - ) - heading_preds = execute_level_mapping(heading_preds, lvl_mapping) - heading_preds.drop("origin_level", axis=1, inplace=True) - save_preds.insert( - save_preds.columns.get_loc("lvl_neg") + 1, - "lvl_map", - heading_preds["level"].tolist(), - ) + The ``proceed_smart`` and ``output_dir`` parameters are retained for API + compatibility but have no effect. + """ + logger.debug("🚀 non-llm parsing => judge_negs filtering") + heading_preds = postprocess_headings(raw_preds, task="judge_negs") + # legitimate heading candidates; default them to top-level. + heading_preds["level"] = heading_preds["level"].map( + lambda x: 1 if x == -2 else x + ) return heading_preds @@ -635,22 +629,19 @@ def est_hierarchies_llm( output_dir=None, csv_suffix="", ): - """LLM-based hierarchy detection — first chunk via LLM, remaining chunks via reason-code mapping. + """LLM-based hierarchy detection — all chunks evaluated independently. When ``KB_LAYOUT_LLM_COMPACT_INPUT`` is enabled (default), consecutive ``level == -1`` rows in ``raw_preds`` are folded into a single placeholder - row (``[N BODY LINES]``) before chunking. This shrinks the prompt, makes - most documents fit into a single chunk (skipping the lossy reason-code - mapping), and preserves the positional signal for the LLM. + row (``[N BODY LINES]``) before chunking. This shrinks the prompt and + preserves the positional signal for the LLM. Strategy: - 1. (Optional) Compact raw_preds so consecutive body rows become placeholders. - 2. Send only the first chunk to LLM for hierarchy prediction. - 3. Collect ``{id -> level}`` from the LLM response (int ids only). - 4. For multi-chunk docs, extend that mapping via reason-code mapping on - chunks 1..N (placeholders excluded). - 5. Expand the id->level mapping back onto the ORIGINAL ``raw_preds``; - any row not present in the mapping defaults to ``level = -1``. + 1. Compact raw_preds so consecutive body rows become placeholders. + 2. Split into chunks and send each independently to LLM. + 3. Collect ``{id -> level}`` from LLM responses (int ids only). + 4. Map levels back onto the ORIGINAL ``raw_preds``; + any row not in the map defaults to ``level = -1``. Args: raw_preds: raw data diff --git a/apps/worker/app/services/document_parser/tables/table_asset_writer.py b/apps/worker/app/services/document_parser/tables/table_asset_writer.py index 679ee028e..c8e5196f5 100644 --- a/apps/worker/app/services/document_parser/tables/table_asset_writer.py +++ b/apps/worker/app/services/document_parser/tables/table_asset_writer.py @@ -15,7 +15,6 @@ class TableAssetInput: keywords: str know_id: str addtime: str - page_nums: str = "" content: str | None = None tokens: str = "" length: int | None = None @@ -39,7 +38,6 @@ def write_table_asset(table_input: TableAssetInput) -> ParsedRow: tokens=table_input.tokens, connectto="", addtime=table_input.addtime, - page_nums=table_input.page_nums, length=table_input.length, ) diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index 3b3afbe61..b4eeda866 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -63,9 +63,27 @@ class StorageConfig(BaseModel): default=104857600, description="Maximum file size in bytes" ) MAX_PDF_PAGE_LIMIT: int = Field( - default=600, + default=200, ge=1, - description="Maximum allowed PDF page count before parsing is rejected", + description="Maximum PDF page count for single-pass MinerU parsing. " + "Documents exceeding this trigger the shard pipeline when enabled.", + ) + OVERSIZED_PDF_SHARD_ENABLED: bool = Field( + default=False, + description="Enable doc_agent shard pipeline for PDFs exceeding MAX_PDF_PAGE_LIMIT. " + "When False, oversized PDFs are rejected.", + ) + OVERSIZED_PDF_SOFT_LIMIT: int = Field( + default=1500, + ge=1, + description="Soft page limit for oversized PDF shard pipeline. " + "Documents exceeding this are rejected with a contact-support message.", + ) + MINERU_SHARD_CONCURRENCY: int = Field( + default=3, + ge=1, + le=10, + description="Maximum concurrent MinerU API calls for shard parsing.", ) SUPPORTED_EXTENSIONS: str = Field( default=".doc,.docx,.pdf,.txt,.xls,.xlsx,.pptx,.jpg,.jpeg,.png,.md", diff --git a/packages/shared-python/shared/models/database/__init__.py b/packages/shared-python/shared/models/database/__init__.py index a0aba9d3d..2d0e7d8a8 100644 --- a/packages/shared-python/shared/models/database/__init__.py +++ b/packages/shared-python/shared/models/database/__init__.py @@ -20,10 +20,12 @@ RetrievalRun, RetrievalStep, ) +from .document_page_plan import DocumentPagePlan from .demo_materialization import DemoMaterialization from .guest_device import GuestDevice from .job import Job from .job_result import JobChunk, JobResult +from .parse_agent import ParseRun, ParseStep # 3. Job-related log models from .job_state_audit_log import JobStateAuditLog @@ -54,12 +56,15 @@ "Document", "DocumentSection", "DocumentChunk", + "DocumentPagePlan", "DemoMaterialization", "GraphNode", "GraphEdge", "RetrievalHitStat", "RetrievalRun", "RetrievalStep", + "ParseRun", + "ParseStep", "StripePriceConfig", "PaymentRecord", "JobStateAuditLog", diff --git a/packages/shared-python/shared/models/database/document_page_plan.py b/packages/shared-python/shared/models/database/document_page_plan.py new file mode 100644 index 000000000..ee2fb228f --- /dev/null +++ b/packages/shared-python/shared/models/database/document_page_plan.py @@ -0,0 +1,32 @@ +"""Persisted page anatomy and future page-processing plans.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Dict, Optional + +from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String +from sqlalchemy.orm import Mapped, mapped_column + +from shared.core.database import Base +from shared.utils.utc_now import utc_now_naive + + +class DocumentPagePlan(Base): + __tablename__ = "document_page_plan" + + page_plan_id: Mapped[str] = mapped_column(String(36), primary_key=True) + job_id: Mapped[str] = mapped_column( + String(36), ForeignKey("jobs.job_id", ondelete="CASCADE"), nullable=False + ) + page_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + shard_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + global_signals: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + Index("idx_document_page_plan_job", "job_id"), + Index("idx_document_page_plan_created", "created_at"), + ) diff --git a/packages/shared-python/shared/models/database/parse_agent.py b/packages/shared-python/shared/models/database/parse_agent.py new file mode 100644 index 000000000..430957c43 --- /dev/null +++ b/packages/shared-python/shared/models/database/parse_agent.py @@ -0,0 +1,63 @@ +"""Parse-side agent trace models.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Dict, Optional + +from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String +from sqlalchemy.orm import Mapped, mapped_column + +from shared.core.database import Base +from shared.utils.utc_now import utc_now_naive + + +class ParseRun(Base): + __tablename__ = "parse_runs" + + run_id: Mapped[str] = mapped_column(String(36), primary_key=True) + job_id: Mapped[str] = mapped_column( + String(36), ForeignKey("jobs.job_id", ondelete="CASCADE"), nullable=False + ) + kind: Mapped[str] = mapped_column(String(32), nullable=False, default="profile") + final_status: Mapped[str] = mapped_column(String(32), nullable=False) + rounds_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + total_tokens: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + total_latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + summary: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + started_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + finished_at: Mapped[Optional[datetime]] = mapped_column( + DateTime, default=utc_now_naive, nullable=True + ) + + __table_args__ = ( + Index("idx_parse_runs_job_kind", "job_id", "kind"), + Index("idx_parse_runs_started", "started_at"), + ) + + +class ParseStep(Base): + __tablename__ = "parse_steps" + + step_id: Mapped[str] = mapped_column(String(36), primary_key=True) + run_id: Mapped[str] = mapped_column( + String(36), ForeignKey("parse_runs.run_id", ondelete="CASCADE"), nullable=False + ) + round_index: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + actor: Mapped[str] = mapped_column(String(64), nullable=False) + action_type: Mapped[str] = mapped_column(String(64), nullable=False) + tool_name: Mapped[Optional[str]] = mapped_column(String(64), nullable=True) + tool_args: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + observation: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + tokens_used: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime, default=utc_now_naive, nullable=False + ) + + __table_args__ = ( + Index("idx_parse_steps_run_round", "run_id", "round_index"), + Index("idx_parse_steps_tool", "tool_name"), + ) 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 79a00e97a..e8d5f4135 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 @@ -176,14 +176,14 @@ def _should_use_ali_pool(self) -> bool: return False return _is_ali_model(self.default_model) - def _make_ali_pool_call( + def _make_ali_pool_raw_call( self, model: str, all_messages: List[ChatCompletionMessageParam], temperature: float, max_tokens: int, api_kwargs: Dict[str, Any], - ) -> tuple[str, LLMUsage]: + ) -> tuple[Any, LLMUsage]: """Acquire a token, make the call, and retry inline on 429.""" from shared.services.ai.ali_quota_manager import get_ali_quota_manager @@ -208,13 +208,12 @@ def _make_ali_pool_call( max_tokens=max_tokens, **api_kwargs, ) - choices = response.choices - if not choices: + if not response.choices: raise LLMServiceException( internal_message="AI returned empty result", provider=self.default_model, ) - return choices[0].message.content or "", _extract_usage(response) + return response, _extract_usage(response) except openai.RateLimitError as exc: retry_after = _parse_retry_after(exc) quota_manager.mark_rate_limited(lease.token_id, retry_after) @@ -253,8 +252,114 @@ def _make_ali_pool_call( provider=self.default_model, ) + def _make_ali_pool_call( + self, + model: str, + all_messages: List[ChatCompletionMessageParam], + temperature: float, + max_tokens: int, + api_kwargs: Dict[str, Any], + ) -> tuple[str, LLMUsage]: + response, usage = self._make_ali_pool_raw_call( + model=model, + all_messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + api_kwargs=api_kwargs, + ) + return response.choices[0].message.content or "", usage + # ------------------------------------------------------------------ + def chat_completion_raw_with_usage( + self, + messages: Union[str, List[ChatCompletionMessageParam]], + model: Optional[str] = None, + temperature: float = 0.1, + max_tokens: int = 4096, + top_p: Optional[float] = None, + timeout: Optional[int] = None, + **kwargs, + ) -> tuple[Any, LLMUsage]: + all_messages: List[ChatCompletionMessageParam] + if isinstance(messages, list): + all_messages = messages # type: ignore[assignment] + else: + all_messages = [{"role": "user", "content": str(messages)}] + + api_kwargs: Dict[str, Any] = {} + if top_p is not None: + api_kwargs["top_p"] = top_p + if timeout is not None: + api_kwargs["timeout"] = timeout + allowed_api_params = { + "n", "stop", "presence_penalty", "frequency_penalty", + "logit_bias", "user", "seed", "tools", "tool_choice", + "response_format", "logprobs", "top_logprobs", + } + for key, value in kwargs.items(): + if key in allowed_api_params: + api_kwargs[key] = value + + extra_body = api_kwargs.get("extra_body", {}) + if isinstance(extra_body, dict): + extra_body.setdefault("enable_thinking", False) + else: + extra_body = {"enable_thinking": False} + api_kwargs["extra_body"] = extra_body + + effective_model = model or self.default_model + if _should_mock_llm_calls(): + content = build_mock_chat_completion_response( + messages=all_messages, + model_name=effective_model, + ) + 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, + ) + + 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, + temperature=temperature, + max_tokens=max_tokens, + **api_kwargs, + ) + if not response.choices: + raise LLMServiceException( + internal_message="AI returned empty result", + provider=self.default_model, + ) + return response, _extract_usage(response) + except LLMServiceException: + raise + except Exception as exc: + logger.error( + "LLM raw request failed: model={model}, base_url={base_url}, error_chain={error_chain}", + model=effective_model, + base_url=client.base_url, + error_chain=_summarize_exception_chain(exc), + ) + raise LLMServiceException( + internal_message=f"API request failed: {_summarize_exception_chain(exc)}", + provider=self.default_model, + original_exception=exc, + ) from exc + def chat_completion_with_usage( self, messages: Union[str, List[ChatCompletionMessageParam]], diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 339e00d7d..e587294e7 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -339,13 +339,11 @@ def build_prompt(task, texts, query, **kwargs): # """ elif task == "eval-headings": - # COMPACT-input variant. Input is pre-compressed by `_compact_for_llm` + # COMPACT-input variant. Input is pre-compressed by `compact_for_llm` # so that consecutive body-text rows are folded into a single - # ``[N BODY LINES]`` placeholder row. The LLM therefore sees only: - # * heading CANDIDATES (integer id, real heading text), and - # * PLACEHOLDER rows (id with "-" or a single collapsed id, heading - # "[N BODY LINES]", level "-") that carry positional / section-bulk - # signals. + # ``[N BODY LINES]`` placeholder row. The LLM sees only two columns + # (``id`` and ``heading``) — no preliminary level estimate is provided, + # so the model assigns levels purely from structural/semantic analysis. temperature = 0 top_p = 0.01 max_depth = kwargs["paras"]["max_depth"] @@ -354,38 +352,44 @@ def build_prompt(task, texts, query, **kwargs): if toc_context: toc_section = f""" - ***Important Reference: Table of Contents (TOC)*** - The following is the document's table of contents with predefined levels. - Use it as a prior when assigning levels to CANDIDATE rows: + ***CONFIRMED Structure: Table of Contents (TOC)*** ''' {toc_context} ''' - - If a candidate's heading matches a TOC entry, use the TOC's predefined level. - - If a candidate appears to be a sub-section of a TOC entry, assign a deeper level. - - If a candidate does NOT appear in the TOC, it can ONLY be either body text - (level = -1) or a sub-section deeper than the nearest TOC heading above it. + RULES for using the TOC: + 1. MUST TRUST the TOC levels as ground truth. If a candidate heading matches or + closely corresponds to a TOC entry, you MUST assign it the same level as + the TOC entry. Do NOT override or re-interpret the TOC's level assignment. + 2. Candidates that appear between two TOC entries should be assigned a level + DEEPER than the TOC entry they fall under (they are sub-sections not listed + in the TOC). + 3. A candidate that does NOT correspond to any TOC entry can only be: + - Body text (level = -1), OR + - A sub-section with a level deeper than its nearest TOC heading above it. + 4. The TOC provides the SKELETON of the document. Your job is to fill in the + gaps for candidates not covered by the TOC, while strictly preserving the + TOC's structure. """ else: toc_section = "" prompt = f""" You are a document structure auditing expert. The input you receive is a - COMPACT skeleton of a document. Body-text lines have already been collapsed for - you so that every row is one of two kinds: + COMPACT skeleton of a document. Body-text lines have already been collapsed + for you so that every row is one of two kinds: - 1) HEADING CANDIDATE — ``id`` is an integer. ``heading`` is the candidate text. - ``level`` is a preliminary estimate: a positive integer (1 = shallowest, deeper = larger) - or the string "Not Sure" (undetermined). These rows — and ONLY these — are the ones you must evaluate. + 1) HEADING CANDIDATE — ``id`` is an integer, ``heading`` is the candidate + text. These — and ONLY these — are the rows you can evaluate. - 2) PLACEHOLDER — ``id`` is ALWAYS a hyphenated range "start-end" (for a - single-line it is "N-N", e.g. "56-56"); ``heading`` is "[N BODY LINES]" - where N is the number of body lines folded here; ``level`` is the literal "-". + 2) PLACEHOLDER — ``id`` is ALWAYS a range "start-end" (for a single-line + it is "N-N", e.g. "56-56"); ``heading`` is "[N BODY LINES]" + where N is the number of body lines folded here. Placeholders are positional markers that tell you how many body lines sit between adjacent candidates. Use them as context ONLY. - Data to be adjusted: + Data to be evaluated: ''' {texts} ''' @@ -394,10 +398,8 @@ def build_prompt(task, texts, query, **kwargs): ***Hard rules about placeholders*** - Placeholders are NEVER candidates. Do not output them. - - Every ``id`` in your output MUST be a single integer; never emit an id - containing a hyphen ("-"). Never emit the level string "-". - - Use N in ``[N BODY LINES]`` as a "section bulk" signal when applying - the rules below (Rule 6 in particular). + - Every ``id`` in your output MUST be a single integer; never emit an id containing a hyphen. + - Use N in ``[N BODY LINES]`` as a "section bulk" signal when applying rules below (Rule 2 in particular). ***Process in TWO steps:*** @@ -407,15 +409,14 @@ def build_prompt(task, texts, query, **kwargs): - Decimal numbering: "1.", "1.1", "1.1.1" → depth increases with dot count - Enumeration styles: "一、" "（一）" "1、" "①" "1 " → shallower to deeper with increasing numbers - Chapter/section keywords: "Chapter X", "Part X", "第X章", "第X节" - - Upper case / lower case differences in candidate headings - Clear semantic granularities or groups of themes + - Upper case / lower case differences Rank these patterns from shallowest to deepest to form a pattern → level mapping. Placeholder rows MUST NOT influence this scan. **STEP 2 — Assign a level to every candidate (rules in priority order)** - A candidate whose preliminary ``level`` is "Not Sure" or any positive - integer is **always** open to revision. Pure body text has already been - folded into placeholders, but a candidate **can still be** demoted to level = -1. + Your task is to determine each candidate's heading level from scratch + based on its text, context, and the patterns discovered in STEP 1. Rule 0 — Global consistency: Candidates sharing the same structural pattern or semantic granularity SHOULD receive the @@ -423,7 +424,7 @@ def build_prompt(task, texts, query, **kwargs): shares one level; every "X.Y.Z" shares a different, deeper level.) Rule 1 — Parent-child continuity and no level skipping: - A heading, compared to candidates before it, may stay at the same level, + A heading, compared to candidates before it, may stay at the same level, or go ONE level deeper than its nearest valid ancestor heading. However, jumps such as level 1 → level 3 are **always invalid**. @@ -431,19 +432,17 @@ def build_prompt(task, texts, query, **kwargs): A candidate WITHOUT any structural/numbering marker can still be a heading, but ONLY when ALL of the following hold: a) The text is short and title-like — no sentence-ending punctuation. - b) It is NOT a broken fragment that continues into the next row + b) It is NOT a broken fragment that continues into the next row. c) In the input sequence it is IMMEDIATELY followed by a placeholder ``[N BODY LINES]``, or by another candidate with finer granularity. This is the "section bulk" signal — the row introduces a body block or a subsection group. - When Rule-2 is satisfied, pick a level consistent with Rule 1 + When Rule-2 is satisfied, pick a level consistent with Rule 1. Rule 3 — Body text demotion (candidate → -1): - Demote a candidate to level = -1 when it clearly does NOT serve as a section title. - In compact input, the strongest demotion cues are: - - Two CANDIDATE rows appear adjacent with NO placeholder between them - - The text contains equations and math symbols such as + = - × ÷. - - The text is exactly "Figure/Image", demote it to level = -1. - - The text is an isolated broken phrase, fragment, data value, or caption-like snippet (e.g. "Table 3-2", "Figure 4" + Demote a candidate to level = -1 when it clearly does NOT serve as a + section title. Strongest demotion cues are: + - Two CANDIDATE rows appear adjacent with NO placeholder between them. + - The text is an isolated broken phrase, fragment, data value, or caption-like snippet (e.g. "Table 3-2", "Figure 4"). Rule 4 — Normalise to start at level 1: The shallowest (the most coarse granularity) heading found MUST be assigned level 1.