feat(doc-agent): Modular Document Anatomy Agent — bootstrap, planner, executor, persist

.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

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

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")

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": []},
                 }
             ],
         )

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",
 ]

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"]

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"]

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"]

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"]

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),
+        }