From b026f6bbb1437cf034bd0ddb5ba05b2cf377687c Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 25 Aug 2026 12:32:50 +0800 Subject: [PATCH 01/14] Reduce duplicate watchdog log noise --- apps/worker/app/core/visibility_recovery_watchdog.py | 7 ++++--- .../shared/services/redis/periodic_task_lock.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/apps/worker/app/core/visibility_recovery_watchdog.py b/apps/worker/app/core/visibility_recovery_watchdog.py index 0d2d3cd5b..fa3b3febc 100644 --- a/apps/worker/app/core/visibility_recovery_watchdog.py +++ b/apps/worker/app/core/visibility_recovery_watchdog.py @@ -46,9 +46,10 @@ def run_visibility_recovery_watchdog() -> None: while True: try: result: VisibilityRecoveryResult = restore_expired_reservations() - logger.bind(**result).info( - "Expired Celery reservation recovery sweep attempted" - ) + if result["status"] == "attempted": + logger.bind(**result).debug( + "Expired Celery reservation recovery sweep attempted" + ) except Exception: logger.exception("Expired Celery reservation recovery sweep failed") finally: diff --git a/packages/shared-python/shared/services/redis/periodic_task_lock.py b/packages/shared-python/shared/services/redis/periodic_task_lock.py index 47859cb14..3d07d38d8 100644 --- a/packages/shared-python/shared/services/redis/periodic_task_lock.py +++ b/packages/shared-python/shared/services/redis/periodic_task_lock.py @@ -105,7 +105,7 @@ def periodic_task_lock( if acquired: logger.debug(f"periodic_task_lock: acquired for task='{task_name}', ttl={ttl}s") else: - logger.info( + logger.debug( f"periodic_task_lock: task='{task_name}' already running in this " "window (duplicate Beat firing) — skipping" ) From 3e6d940f919307fbf52a3ce3f41bacabadd3dad2 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 17:26:51 +0800 Subject: [PATCH 02/14] fix: bound retrieval snapshot SQL loads --- apps/api/app/core/exception_handlers.py | 18 +- .../test_exception_handlers_contract.py | 55 ++ ...etrieval_snapshot_large_corpus_contract.py | 475 ++++++++++++++++++ .../shared/core/response/ErrorCode.py | 3 + .../shared/services/retrieval/nav_snapshot.py | 196 +++++--- .../shared/tests/test_nav_snapshot.py | 9 +- 6 files changed, 668 insertions(+), 88 deletions(-) create mode 100644 apps/api/tests/contract/test_exception_handlers_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py diff --git a/apps/api/app/core/exception_handlers.py b/apps/api/app/core/exception_handlers.py index 1df2f085b..7d5a72680 100644 --- a/apps/api/app/core/exception_handlers.py +++ b/apps/api/app/core/exception_handlers.py @@ -48,7 +48,7 @@ """ import uuid -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from typing import List, cast from fastapi import FastAPI, HTTPException, Request @@ -88,7 +88,10 @@ def _get_request_id(request: Request) -> str: async def knowhere_exception_handler( - request: Request, exc: KnowhereException + request: Request, + exc: KnowhereException, + *, + response_headers: Mapping[str, str] | None = None, ) -> JSONResponse: """ This handler enforces the separation between: @@ -119,7 +122,10 @@ async def knowhere_exception_handler( exc.logging(request_id=request_id) # Always include request ID header for client-side correlation - headers = {"X-Request-ID": request_id} + # Preserve framework-provided headers such as ``Allow: POST`` for 405 + # responses and ``WWW-Authenticate`` for authentication challenges. + headers = dict(response_headers or {}) + headers["X-Request-ID"] = request_id retry_after = exc.details.get("retry_after") if retry_after: headers["Retry-After"] = str(retry_after) @@ -208,7 +214,11 @@ async def http_exception_handler(request: Request, exc: HTTPException) -> JSONRe ) # Delegate to central handler - return await knowhere_exception_handler(request, knowhere_exc) + return await knowhere_exception_handler( + request, + knowhere_exc, + response_headers=exc.headers, + ) async def validation_exception_handler( diff --git a/apps/api/tests/contract/test_exception_handlers_contract.py b/apps/api/tests/contract/test_exception_handlers_contract.py new file mode 100644 index 000000000..40eff1064 --- /dev/null +++ b/apps/api/tests/contract/test_exception_handlers_contract.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import sys +from pathlib import Path +from typing import cast + +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient + +from tests.support.import_environment import ( + configure_import_environment, + ensure_import_paths, +) + +configure_import_environment() +ensure_import_paths() + + +def _prepare_api_app_imports() -> None: + api_root = str(Path(__file__).resolve().parents[2]) + if api_root in sys.path: + sys.path.remove(api_root) + sys.path.insert(0, api_root) + + +def _create_post_only_app() -> FastAPI: + _prepare_api_app_imports() + + from app.core.exception_handlers import setup_exception_handlers + + app = FastAPI() + + @app.post("/v2/retrieval/query") + async def query_retrieval() -> dict[str, bool]: + return {"ok": True} + + setup_exception_handlers(app) + return app + + +async def test_get_to_post_only_route_returns_method_not_allowed() -> None: + app = _create_post_only_app() + transport = ASGITransport(app=app) + + async with AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get("/v2/retrieval/query") + + assert response.status_code == 405 + assert response.headers["allow"] == "POST" + + response_json = cast(dict[str, object], response.json()) + error = cast(dict[str, object], response_json["error"]) + assert response_json["success"] is False + assert error["code"] == "METHOD_NOT_ALLOWED" + assert error["message"] == "Method not allowed" diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py new file mode 100644 index 000000000..5bbb2c4f5 --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -0,0 +1,475 @@ +from __future__ import annotations + +import time +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from collections.abc import AsyncGenerator +from typing import Any +from uuid import uuid4 + +import pytest +from httpx import AsyncClient + +from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job_result import JobResult +from shared.services.retrieval.nav_snapshot import load_nav_snapshot +from sqlalchemy import select, text +from sqlalchemy.engine import Row +from sqlalchemy.exc import DBAPIError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from tests.support.contract_database import ContractDatabase +from shared.testing.contract_runtime import get_contract_database_url + + +_USER_ID = "local-dev-user" +_DOCUMENT_COUNT = 100 +_CHUNKS_PER_DOCUMENT = 600 +_SECTIONS_PER_DOCUMENT = 8 +_CONTENT_BYTES = 2048 +_TOTAL_CHUNKS = _DOCUMENT_COUNT * _CHUNKS_PER_DOCUMENT +LegacySnapshotRow = Row[ + tuple[ + str, + str, + str | None, + str, + str, + int, + str, + str | None, + dict[str, object], + str | None, + str | None, + ] +] + + +@asynccontextmanager +async def _contract_db_session() -> AsyncGenerator[AsyncSession, None]: + engine = create_async_engine(get_contract_database_url(), future=True) + session_factory = async_sessionmaker(engine, expire_on_commit=False) + try: + async with session_factory() as session: + yield session + finally: + await engine.dispose() + + +async def _seed_large_retrieval_corpus(namespace: str) -> None: + await ContractDatabase.execute( + """ + INSERT INTO jobs ( + job_id, user_id, job_type, status, source_type, version, + webhook_enabled, created_at, updated_at, credits_charged, billing_status + ) + SELECT + 'job_lg_' || i, + :user_id, + 'document_ingestion', + 'done', + 'file', + 0, + false, + NOW(), + NOW(), + 0, + 'skipped' + FROM generate_series(1, :document_count) AS values(i) + """, + {"user_id": _USER_ID, "document_count": _DOCUMENT_COUNT}, + ) + await ContractDatabase.execute( + """ + INSERT INTO documents ( + document_id, user_id, namespace, status, current_job_result_id, + source_file_name, parse_track, created_at, updated_at + ) + SELECT + 'doc_lg_' || i, + :user_id, + :namespace, + 'active', + NULL, + 'large-corpus-' || i || '.pdf', + 'chunk', + NOW(), + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + { + "user_id": _USER_ID, + "namespace": namespace, + "document_count": _DOCUMENT_COUNT, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO job_results ( + id, job_id, document_id, delivery_mode, created_at, updated_at + ) + SELECT + 'result_lg_' || i, + 'job_lg_' || i, + 'doc_lg_' || i, + 'inline', + NOW(), + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + {"document_count": _DOCUMENT_COUNT}, + ) + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = 'result_lg_' || i + FROM generate_series(1, :document_count) AS values(i) + WHERE documents.document_id = 'doc_lg_' || i + """, + {"document_count": _DOCUMENT_COUNT}, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_sections ( + section_id, user_id, namespace, document_id, job_result_id, + section_path, section_title, section_level, sort_order, created_at + ) + SELECT + 'section_lg_' || i || '_' || section_number, + :user_id, + :namespace, + 'doc_lg_' || i, + 'result_lg_' || i, + 'large-corpus-' || i || '/section/' || section_number, + 'section-' || section_number, + 1, + section_number, + NOW() + FROM generate_series(1, :document_count) AS values(i) + CROSS JOIN generate_series(1, :sections_per_document) AS sections(section_number) + """, + { + "user_id": _USER_ID, + "namespace": namespace, + "document_count": _DOCUMENT_COUNT, + "sections_per_document": _SECTIONS_PER_DOCUMENT, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_chunks ( + id, chunk_id, user_id, namespace, document_id, job_result_id, + section_id, chunk_type, content, content_lexical_text, + path_lexical_text, content_search_text, path_search_text, + term_search_text, source_chunk_path, chunk_metadata, sort_order, + created_at + ) + SELECT + 'dchunk_lg_' || document_number || '_' || chunk_number, + 'chunk_lg_' || document_number || '_' || chunk_number, + :user_id, + :namespace, + 'doc_lg_' || document_number, + 'result_lg_' || document_number, + 'section_lg_' || document_number || '_' || section_number, + 'text', + repeat(md5(document_number::text || ':' || chunk_number::text), 64), + repeat(md5(document_number::text || ':' || chunk_number::text), 64), + 'large-corpus-' || document_number || '/section/' || section_number || '/' || chunk_number, + repeat(md5(document_number::text || ':' || chunk_number::text), 64), + 'large-corpus-' || document_number || '/section/' || section_number || '/' || chunk_number, + repeat(md5(document_number::text || ':' || chunk_number::text), 64), + 'large-corpus-' || document_number || '/section/' || section_number || '/' || chunk_number, + json_build_object( + 'tokens', ARRAY['benchmark', 'retrieval', 'document', document_number::text], + 'keywords', ARRAY['benchmark', 'production-shaped'], + 'summary', repeat('payload ', 32) + )::json, + chunk_number, + NOW() + FROM generate_series(1, :document_count) AS documents(document_number) + CROSS JOIN generate_series(1, :chunks_per_document) AS chunks(chunk_number) + CROSS JOIN LATERAL ( + SELECT ((chunk_number - 1) % :sections_per_document) + 1 AS section_number + ) AS section_values + """, + { + "user_id": _USER_ID, + "namespace": namespace, + "document_count": _DOCUMENT_COUNT, + "chunks_per_document": _CHUNKS_PER_DOCUMENT, + "sections_per_document": _SECTIONS_PER_DOCUMENT, + }, + ) + + +async def _load_legacy_rows( + namespace: str, + *, + statement_timeout_ms: int | None = None, +) -> list[LegacySnapshotRow]: + stmt = ( + select( + Document.document_id, + DocumentChunk.chunk_id, + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.sort_order, + DocumentChunk.source_chunk_path, + DocumentChunk.file_path, + DocumentChunk.chunk_metadata, + DocumentSection.section_path, + JobResult.job_id, + ) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .outerjoin(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == _USER_ID) + .where(Document.namespace == namespace) + .where(Document.status == "active") + .order_by(Document.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id) + ) + async with _contract_db_session() as db: + if statement_timeout_ms is not None: + await db.execute( + text("SELECT set_config('statement_timeout', :timeout_value, true)"), + {"timeout_value": f"{statement_timeout_ms}ms"}, + ) + rows: list[LegacySnapshotRow] = list((await db.execute(stmt)).all()) + return rows + + +async def _explain_snapshot_query(namespace: str) -> dict[str, Any]: + row = await ContractDatabase.fetch_one( + """ + EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) + SELECT + d.document_id, + c.chunk_id, + c.section_id, + c.chunk_type, + c.content, + c.sort_order, + c.source_chunk_path, + c.file_path, + c.chunk_metadata, + s.section_path, + r.job_id + FROM documents AS d + JOIN document_chunks AS c + ON c.document_id = d.document_id + AND c.job_result_id = d.current_job_result_id + LEFT JOIN document_sections AS s ON s.section_id = c.section_id + LEFT JOIN job_results AS r ON r.id = c.job_result_id + WHERE d.user_id = :user_id + AND d.namespace = :namespace + AND d.status = 'active' + ORDER BY d.document_id, c.sort_order, c.chunk_id + """, + {"user_id": _USER_ID, "namespace": namespace}, + ) + if row is None: + raise AssertionError("EXPLAIN returned no plan") + plan_value = next(iter(row.values())) + if not isinstance(plan_value, list) or not plan_value: + raise AssertionError("EXPLAIN returned an unexpected plan shape") + plan = plan_value[0] + if not isinstance(plan, dict): + raise AssertionError("EXPLAIN plan root is not an object") + return plan + + +async def _explain_revision_query() -> dict[str, Any]: + row = await ContractDatabase.fetch_one( + """ + EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) + SELECT + chunk_id, + section_id, + chunk_type, + content, + sort_order, + source_chunk_path, + file_path, + chunk_metadata, + id + FROM document_chunks + WHERE document_id = 'doc_lg_1' + AND job_result_id = 'result_lg_1' + ORDER BY sort_order, chunk_id, id + LIMIT 2000 + """ + ) + if row is None: + raise AssertionError("revision EXPLAIN returned no plan") + plan_value = next(iter(row.values())) + if not isinstance(plan_value, list) or not plan_value: + raise AssertionError("revision EXPLAIN returned an unexpected plan shape") + plan = plan_value[0] + if not isinstance(plan, dict): + raise AssertionError("revision EXPLAIN plan root is not an object") + return plan + + +async def _measure_seeded_payload(namespace: str) -> int: + row = await ContractDatabase.fetch_one( + """ + SELECT COALESCE( + SUM(pg_column_size(content) + pg_column_size(chunk_metadata)), 0 + ) AS payload_bytes + FROM document_chunks + WHERE user_id = :user_id AND namespace = :namespace + """, + {"user_id": _USER_ID, "namespace": namespace}, + ) + if row is None or not isinstance(row.get("payload_bytes"), int): + raise AssertionError("payload measurement returned an unexpected value") + return row["payload_bytes"] + + +def _plan_nodes(plan_node: dict[str, Any]) -> list[str]: + summary = ( + f"{plan_node.get('Node Type')}" + f"[{plan_node.get('Relation Name', '')}" + f"/{plan_node.get('Index Name', '')}]" + f"={plan_node.get('Actual Total Time')}ms" + ) + children = plan_node.get("Plans", []) + if not isinstance(children, list): + return [summary] + nodes = [summary] + for child in children: + if isinstance(child, dict): + nodes.extend(_plan_nodes(child)) + return nodes + + +async def test_large_snapshot_keeps_all_retrieval_inputs_after_sql_optimization( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + namespace = f"large-corpus-{uuid4().hex[:8]}" + async with developer_api_client_factory(): + await _seed_large_retrieval_corpus(namespace) + payload_bytes = await _measure_seeded_payload(namespace) + print( + "seeded payload: " + f"{payload_bytes / (1024 * 1024):.1f}MiB " + f"content_bytes={_CONTENT_BYTES} " + f"sections_per_document={_SECTIONS_PER_DOCUMENT}" + ) + + explain_plan = await _explain_snapshot_query(namespace) + explain_root = explain_plan["Plan"] + assert isinstance(explain_root, dict) + print( + "snapshot EXPLAIN: " + f"node={explain_root.get('Node Type')} " + f"time={explain_plan.get('Execution Time')}ms " + f"shared_hit={explain_root.get('Shared Hit Blocks')} " + f"shared_read={explain_root.get('Shared Read Blocks')} " + f"nodes={' -> '.join(_plan_nodes(explain_root))}" + ) + revision_plan = await _explain_revision_query() + revision_root = revision_plan["Plan"] + assert isinstance(revision_root, dict) + print( + "revision EXPLAIN: " + f"time={revision_plan.get('Execution Time')}ms " + f"nodes={' -> '.join(_plan_nodes(revision_root))}" + ) + await ContractDatabase.execute( + """ + CREATE INDEX IF NOT EXISTS idx_benchmark_chunks_revision_order + ON document_chunks (document_id, job_result_id, sort_order, chunk_id, id) + """ + ) + indexed_explain_plan = await _explain_snapshot_query(namespace) + indexed_root = indexed_explain_plan["Plan"] + assert isinstance(indexed_root, dict) + print( + "snapshot EXPLAIN with candidate index: " + f"node={indexed_root.get('Node Type')} " + f"time={indexed_explain_plan.get('Execution Time')}ms " + f"nodes={' -> '.join(_plan_nodes(indexed_root))}" + ) + indexed_revision_plan = await _explain_revision_query() + indexed_revision_root = indexed_revision_plan["Plan"] + assert isinstance(indexed_revision_root, dict) + print( + "revision EXPLAIN with candidate index: " + f"time={indexed_revision_plan.get('Execution Time')}ms " + f"nodes={' -> '.join(_plan_nodes(indexed_revision_root))}" + ) + await ContractDatabase.execute( + "DROP INDEX IF EXISTS idx_benchmark_chunks_revision_order" + ) + + legacy_started = time.perf_counter() + legacy_rows = await _load_legacy_rows(namespace) + legacy_elapsed = time.perf_counter() - legacy_started + + optimized_started = time.perf_counter() + async with _contract_db_session() as db: + snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + ) + optimized_elapsed = time.perf_counter() - optimized_started + + with pytest.raises(DBAPIError, match="statement timeout"): + await _load_legacy_rows(namespace, statement_timeout_ms=100) + + async with _contract_db_session() as db: + await db.execute(text("SET LOCAL statement_timeout = 100")) + bounded_snapshot = await load_nav_snapshot( + db, + user_id=_USER_ID, + namespace=namespace, + ) + + assert len(legacy_rows) == _TOTAL_CHUNKS + assert len(snapshot.document_ids) == _DOCUMENT_COUNT + assert len(bounded_snapshot.document_ids) == _DOCUMENT_COUNT + + optimized_chunks = { + chunk_id: (document_id, chunk) + for document_id in snapshot.document_ids + for section_id in snapshot.provider.children(document_id) + for chunk in snapshot.provider.self_units(section_id) + for chunk_id in [chunk.chunk_id] + } + assert len(optimized_chunks) == _TOTAL_CHUNKS + + legacy_by_chunk = { + str(row[1]): row + for row in legacy_rows + } + for chunk_id, (document_id, chunk) in optimized_chunks.items(): + legacy_row = legacy_by_chunk[chunk_id] + assert document_id == str(legacy_row[0]) + assert chunk.section_id == str(legacy_row[2]) + assert chunk.chunk_type == str(legacy_row[3]) + assert chunk.content == str(legacy_row[4]) + assert chunk.sort_order == int(legacy_row[5]) + assert chunk.source_chunk_path == str(legacy_row[6]) + assert chunk.file_path == str(legacy_row[7] or "") + assert chunk.metadata == (legacy_row[8] if isinstance(legacy_row[8], dict) else {}) + + reference = snapshot.chunk_ref_index[f"{document_id}:{chunk_id}"] + assert reference["document_id"] == document_id + assert reference["section_path"] == str(legacy_row[9] or "") + assert reference["chunk_type"] == str(legacy_row[3]) + assert reference["file_path"] == (str(legacy_row[7]) if legacy_row[7] else None) + assert reference["job_id"] == str(legacy_row[10]) + + print( + f"large snapshot benchmark: legacy={legacy_elapsed:.3f}s " + f"optimized={optimized_elapsed:.3f}s chunks={_TOTAL_CHUNKS}" + ) diff --git a/packages/shared-python/shared/core/response/ErrorCode.py b/packages/shared-python/shared/core/response/ErrorCode.py index 4c9820936..cf5dd2f8f 100644 --- a/packages/shared-python/shared/core/response/ErrorCode.py +++ b/packages/shared-python/shared/core/response/ErrorCode.py @@ -55,6 +55,7 @@ class ErrorCode(str, Enum): ) PERMISSION_DENIED = "PERMISSION_DENIED" # 403 - Caller lacks permission NOT_FOUND = "NOT_FOUND" # 404 - Resource does not exist + METHOD_NOT_ALLOWED = "METHOD_NOT_ALLOWED" # 405 - HTTP method is not supported ABORTED = "ABORTED" # 409 - Concurrency conflict ALREADY_EXISTS = "ALREADY_EXISTS" # 409 - Resource already exists RESOURCE_EXHAUSTED = ( @@ -89,6 +90,7 @@ class ErrorCodeMapper: ErrorCode.PAYMENT_REQUIRED: 402, ErrorCode.PERMISSION_DENIED: 403, ErrorCode.NOT_FOUND: 404, + ErrorCode.METHOD_NOT_ALLOWED: 405, ErrorCode.ABORTED: 409, ErrorCode.ALREADY_EXISTS: 409, ErrorCode.RESOURCE_EXHAUSTED: 429, @@ -110,6 +112,7 @@ class ErrorCodeMapper: 402: ErrorCode.PAYMENT_REQUIRED, 403: ErrorCode.PERMISSION_DENIED, 404: ErrorCode.NOT_FOUND, + 405: ErrorCode.METHOD_NOT_ALLOWED, 409: ErrorCode.ALREADY_EXISTS, 422: ErrorCode.INVALID_ARGUMENT, # Pydantic validation 429: ErrorCode.RESOURCE_EXHAUSTED, diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 964d698b7..011405cf7 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -8,9 +8,9 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Optional +from typing import Any -from sqlalchemy import select +from sqlalchemy import literal, select, tuple_ from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection @@ -24,6 +24,11 @@ from shared.services.retrieval.search.section_filters import is_excluded_section +# Keep each payload query bounded under the API's 30-second statement timeout. +# Keyset pagination avoids the increasingly expensive OFFSET scans. +_CHUNK_BATCH_SIZE = 2_000 + + @dataclass(frozen=True) class NavSnapshot: """In-memory corpus for one map-nav episode.""" @@ -94,10 +99,27 @@ async def load_nav_snapshot( ) document_titles: dict[str, str] = {} - for document_id, source_file_name, _job in doc_rows: + current_job_result_ids: set[str] = set() + document_revisions: list[tuple[str, str]] = [] + for document_id, source_file_name, current_job_result_id in doc_rows: did = str(document_id) title = str(source_file_name or "").strip() or did document_titles[did] = title + if current_job_result_id: + job_result_id = str(current_job_result_id) + current_job_result_ids.add(job_result_id) + document_revisions.append((did, job_result_id)) + + job_result_rows = await db.execute( + select(JobResult.id, JobResult.job_id).where( + JobResult.id.in_(list(current_job_result_ids)) + ) + ) + job_id_by_result_id = { + str(job_result_id): str(job_id) + for job_result_id, job_id in job_result_rows.all() + if job_result_id and job_id + } sections_by_doc, section_path_by_id = await _load_sections( db, @@ -108,11 +130,10 @@ async def load_nav_snapshot( ) units_by_doc, chunk_ref_index = await _load_chunks( db, - user_id=user_id, - namespace=namespace, - exclude_document_ids=excluded_docs, + document_revisions=document_revisions, exclude_sections=excluded_secs, section_path_by_id=section_path_by_id, + job_id_by_result_id=job_id_by_result_id, ) # Keep only documents that still have sections after exclude filters. @@ -196,87 +217,100 @@ async def _load_sections( async def _load_chunks( db: AsyncSession, *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], + document_revisions: list[tuple[str, str]], exclude_sections: list[dict[str, str]], section_path_by_id: dict[str, str], + job_id_by_result_id: dict[str, str], ) -> tuple[dict[str, list[UnitRow]], dict[str, dict[str, Any]]]: - stmt = ( - select( - Document.document_id, - DocumentChunk.chunk_id, - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.sort_order, - DocumentChunk.source_chunk_path, - DocumentChunk.file_path, - DocumentChunk.chunk_metadata, - DocumentSection.section_path, - JobResult.job_id, - ) - .join( - DocumentChunk, - (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id), - ) - .outerjoin( - DocumentSection, - DocumentSection.section_id == DocumentChunk.section_id, - ) - .outerjoin(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .order_by(Document.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id) - ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) - by_doc: dict[str, list[UnitRow]] = {} ref_index: dict[str, dict[str, Any]] = {} - for row in (await db.execute(stmt)).all(): - document_id = str(row[0]) - chunk_id = str(row[1] or "").strip() - section_id = str(row[2]) if row[2] else None - # Prefer joined section_path; fall back to kept section map. - section_path: Optional[str] = str(row[9]) if row[9] is not None else None - if section_path is None and section_id: - section_path = section_path_by_id.get(section_id) - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - # Drop units whose section was filtered out of the tree. - if section_id and section_id not in section_path_by_id: - continue + for document_id, job_result_id in document_revisions: + last_key: tuple[int, str, str] | None = None + while True: + stmt = ( + select( + DocumentChunk.chunk_id, + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.sort_order, + DocumentChunk.source_chunk_path, + DocumentChunk.file_path, + DocumentChunk.chunk_metadata, + DocumentChunk.id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .order_by( + DocumentChunk.sort_order, + DocumentChunk.chunk_id, + DocumentChunk.id, + ) + .limit(_CHUNK_BATCH_SIZE) + ) + if last_key is not None: + stmt = stmt.where( + tuple_( + DocumentChunk.sort_order, + DocumentChunk.chunk_id, + DocumentChunk.id, + ) + > tuple_( + literal(last_key[0]), + literal(last_key[1]), + literal(last_key[2]), + ) + ) - unit = UnitRow( - chunk_id=chunk_id, - section_id=section_id, - chunk_type=str(row[3] or "text"), - content=str(row[4] or ""), - sort_order=int(row[5] or 0), - source_chunk_path=str(row[6] or ""), - file_path=str(row[7] or ""), - metadata=_as_meta(row[8]), - ) - by_doc.setdefault(document_id, []).append(unit) - if chunk_id: - meta = { - "document_id": document_id, - "section_path": section_path, - "chunk_type": unit.chunk_type, - "file_path": unit.file_path or None, - "job_id": str(row[10]) if row[10] else None, - } - # Bare chunk_id (last-wins) plus doc-scoped key so the same - # chunk_id can appear under multiple documents. - ref_index[chunk_id] = meta - ref_index[f"{document_id}:{chunk_id}"] = meta + rows = (await db.execute(stmt)).all() + if not rows: + break + for row in rows: + chunk_id = str(row[0] or "").strip() + section_id = str(row[1]) if row[1] else None + section_path: str | None = ( + section_path_by_id.get(section_id) if section_id else None + ) + if is_excluded_section( + document_id=document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + if section_id and section_id not in section_path_by_id: + continue + + unit = UnitRow( + chunk_id=chunk_id, + section_id=section_id, + chunk_type=str(row[2] or "text"), + content=str(row[3] or ""), + sort_order=int(row[4] or 0), + source_chunk_path=str(row[5] or ""), + file_path=str(row[6] or ""), + metadata=_as_meta(row[7]), + ) + by_doc.setdefault(document_id, []).append(unit) + if chunk_id: + meta = { + "document_id": document_id, + "section_path": section_path, + "chunk_type": unit.chunk_type, + "file_path": unit.file_path or None, + "job_id": job_id_by_result_id.get(job_result_id), + } + # Bare chunk_id (last-wins) plus doc-scoped key so the same + # chunk_id can appear under multiple documents. + ref_index[chunk_id] = meta + ref_index[f"{document_id}:{chunk_id}"] = meta + last_row = rows[-1] + last_key = ( + int(last_row[4] or 0), + str(last_row[0] or ""), + str(last_row[8]), + ) + if len(rows) < _CHUNK_BATCH_SIZE: + break return by_doc, ref_index diff --git a/packages/shared-python/shared/tests/test_nav_snapshot.py b/packages/shared-python/shared/tests/test_nav_snapshot.py index 340c70687..f8e094c26 100644 --- a/packages/shared-python/shared/tests/test_nav_snapshot.py +++ b/packages/shared-python/shared/tests/test_nav_snapshot.py @@ -95,13 +95,16 @@ def test_build_nav_snapshot_rejects_empty_corpus() -> None: ) -def test_load_nav_snapshot_joins_current_revision_only() -> None: - """Section/chunk loaders must bind rows to Document.current_job_result_id.""" +def test_load_nav_snapshot_filters_current_revision_only() -> None: + """Snapshot loading must select and query only each current revision.""" import inspect from shared.services.retrieval import nav_snapshot as nav_snapshot_mod + loader_src = "".join(inspect.getsource(nav_snapshot_mod.load_nav_snapshot).split()) sections_src = "".join(inspect.getsource(nav_snapshot_mod._load_sections).split()) chunks_src = "".join(inspect.getsource(nav_snapshot_mod._load_chunks).split()) + assert "Document.current_job_result_id" in loader_src assert "DocumentSection.job_result_id==Document.current_job_result_id" in sections_src - assert "DocumentChunk.job_result_id==Document.current_job_result_id" in chunks_src + assert "DocumentChunk.document_id==document_id" in chunks_src + assert "DocumentChunk.job_result_id==job_result_id" in chunks_src From 235bfc32d6b1927f3a533321235fa0f2498fab7e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 18:18:11 +0800 Subject: [PATCH 03/14] fix: make retrieval snapshots revision-consistent --- ...c1d2e3f4_add_chunk_snapshot_order_index.py | 34 ++ ...etrieval_snapshot_large_corpus_contract.py | 377 +++++++++++++++++- .../tests/migrations/test_schema_contract.py | 21 + .../shared/models/database/document.py | 8 + .../shared/services/retrieval/nav_snapshot.py | 202 +++++----- 5 files changed, 535 insertions(+), 107 deletions(-) create mode 100644 apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py diff --git a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py new file mode 100644 index 000000000..474dd7cee --- /dev/null +++ b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py @@ -0,0 +1,34 @@ +"""Add the chunk snapshot pagination order index. + +Revision ID: fbf0c1d2e3f4 +Revises: f0d85d209e68, fbe1c2d3e4f5 +Create Date: 2026-08-27 00:00:00.000000 +""" + +from __future__ import annotations + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "fbf0c1d2e3f4" +down_revision: Union[str, Sequence[str], None] = ( + "f0d85d209e68", + "fbe1c2d3e4f5", +) +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + """ + CREATE INDEX IF NOT EXISTS idx_document_chunks_revision_snapshot_order + ON document_chunks (document_id, job_result_id, sort_order, chunk_id, id) + """ + ) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS idx_document_chunks_revision_snapshot_order") diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index 5bbb2c4f5..b465ea72d 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -1,10 +1,9 @@ from __future__ import annotations import time -from collections.abc import Callable +from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager -from collections.abc import AsyncGenerator -from typing import Any +from typing import Any, cast from uuid import uuid4 import pytest @@ -13,10 +12,11 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult from shared.services.retrieval.nav_snapshot import load_nav_snapshot -from sqlalchemy import select, text +from sqlalchemy import Executable, Result, select, text from sqlalchemy.engine import Row from sqlalchemy.exc import DBAPIError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.sql.selectable import Select from tests.support.contract_database import ContractDatabase from shared.testing.contract_runtime import get_contract_database_url @@ -44,6 +44,43 @@ ] +class _PublishingSession: + def __init__( + self, + session: AsyncSession, + publish_revision: Callable[[], Awaitable[None]], + ) -> None: + self._session = session + self._publish_revision = publish_revision + self._has_published = False + + async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: + result = await self._session.execute(statement) + if not self._has_published: + self._has_published = True + await self._publish_revision() + return cast(Result[tuple[object, ...]], result) + + +class _CountingSession: + def __init__(self, session: AsyncSession) -> None: + self._session = session + self.chunk_query_count = 0 + + async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: + if isinstance(statement, Select): + selected_tables = { + getattr(table, "name", "") + for table in statement.get_final_froms() + } + else: + selected_tables = set() + if "document_chunks" in selected_tables: + self.chunk_query_count += 1 + result = await self._session.execute(statement) + return cast(Result[tuple[object, ...]], result) + + @asynccontextmanager async def _contract_db_session() -> AsyncGenerator[AsyncSession, None]: engine = create_async_engine(get_contract_database_url(), future=True) @@ -202,6 +239,319 @@ async def _seed_large_retrieval_corpus(namespace: str) -> None: ) +async def _seed_republished_document(namespace: str) -> tuple[str, str]: + identifier = uuid4().hex[:8] + document_id = f"doc_race_{identifier}" + old_result_id = f"result_race_old_{identifier}" + new_result_id = f"result_race_new_{identifier}" + await ContractDatabase.execute( + """ + INSERT INTO jobs ( + job_id, user_id, job_type, status, source_type, version, + webhook_enabled, created_at, updated_at, credits_charged, billing_status + ) VALUES + (:old_job_id, :user_id, 'document_ingestion', 'done', 'file', 0, + false, NOW(), NOW(), 0, 'skipped'), + (:new_job_id, :user_id, 'document_ingestion', 'done', 'file', 0, + false, NOW(), NOW(), 0, 'skipped') + """, + { + "old_job_id": f"job_race_old_{identifier}", + "new_job_id": f"job_race_new_{identifier}", + "user_id": _USER_ID, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO documents ( + document_id, user_id, namespace, status, current_job_result_id, + source_file_name, parse_track, created_at, updated_at + ) VALUES ( + :document_id, :user_id, :namespace, 'active', NULL, + 'republished.pdf', 'chunk', NOW(), NOW() + ) + """, + { + "document_id": document_id, + "user_id": _USER_ID, + "namespace": namespace, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO job_results ( + id, job_id, document_id, delivery_mode, created_at, updated_at + ) VALUES + (:old_result_id, :old_job_id, :document_id, 'inline', NOW(), NOW()), + (:new_result_id, :new_job_id, :document_id, 'inline', NOW(), NOW()) + """, + { + "old_result_id": old_result_id, + "new_result_id": new_result_id, + "old_job_id": f"job_race_old_{identifier}", + "new_job_id": f"job_race_new_{identifier}", + "document_id": document_id, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_sections ( + section_id, user_id, namespace, document_id, job_result_id, + section_path, section_title, section_level, sort_order, created_at + ) VALUES + (:old_section_id, :user_id, :namespace, :document_id, :old_result_id, + 'republished.pdf/old', 'old', 1, 1, NOW()), + (:new_section_id, :user_id, :namespace, :document_id, :new_result_id, + 'republished.pdf/new', 'new', 1, 1, NOW()) + """, + { + "old_section_id": f"section_race_old_{identifier}", + "new_section_id": f"section_race_new_{identifier}", + "user_id": _USER_ID, + "namespace": namespace, + "document_id": document_id, + "old_result_id": old_result_id, + "new_result_id": new_result_id, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_chunks ( + id, chunk_id, user_id, namespace, document_id, job_result_id, + section_id, chunk_type, content, source_chunk_path, + chunk_metadata, sort_order, created_at + ) VALUES + (:old_id, :old_chunk_id, :user_id, :namespace, :document_id, + :old_result_id, :old_section_id, 'text', 'old content', + 'republished.pdf/old/chunk', '{}'::json, 1, NOW()), + (:new_id, :new_chunk_id, :user_id, :namespace, :document_id, + :new_result_id, :new_section_id, 'text', 'new content', + 'republished.pdf/new/chunk', '{}'::json, 1, NOW()) + """, + { + "old_id": f"dchunk_race_old_{identifier}", + "new_id": f"dchunk_race_new_{identifier}", + "old_chunk_id": f"chunk_race_old_{identifier}", + "new_chunk_id": f"chunk_race_new_{identifier}", + "user_id": _USER_ID, + "namespace": namespace, + "document_id": document_id, + "old_result_id": old_result_id, + "new_result_id": new_result_id, + "old_section_id": f"section_race_old_{identifier}", + "new_section_id": f"section_race_new_{identifier}", + }, + ) + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = :old_result_id + WHERE document_id = :document_id + """, + {"old_result_id": old_result_id, "document_id": document_id}, + ) + return document_id, new_result_id + + +async def _seed_many_small_documents(namespace: str, *, document_count: int) -> None: + identifier = uuid4().hex[:8] + await ContractDatabase.execute( + """ + INSERT INTO jobs ( + job_id, user_id, job_type, status, source_type, version, + webhook_enabled, created_at, updated_at, credits_charged, billing_status + ) + SELECT + 'job_batch_' || :identifier || '_' || i, + :user_id, + 'document_ingestion', + 'done', + 'file', + 0, + false, + NOW(), + NOW(), + 0, + 'skipped' + FROM generate_series(1, :document_count) AS values(i) + """, + { + "identifier": identifier, + "user_id": _USER_ID, + "document_count": document_count, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO documents ( + document_id, user_id, namespace, status, current_job_result_id, + source_file_name, parse_track, created_at, updated_at + ) + SELECT + 'doc_batch_' || :identifier || '_' || i, + :user_id, + :namespace, + 'active', + NULL, + 'batch-' || i || '.pdf', + 'chunk', + NOW(), + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + { + "identifier": identifier, + "user_id": _USER_ID, + "namespace": namespace, + "document_count": document_count, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO job_results ( + id, job_id, document_id, delivery_mode, created_at, updated_at + ) + SELECT + 'result_batch_' || :identifier || '_' || i, + 'job_batch_' || :identifier || '_' || i, + 'doc_batch_' || :identifier || '_' || i, + 'inline', + NOW(), + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + {"identifier": identifier, "document_count": document_count}, + ) + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = 'result_batch_' || :identifier || '_' || i + FROM generate_series(1, :document_count) AS values(i) + WHERE document_id = 'doc_batch_' || :identifier || '_' || i + """, + {"identifier": identifier, "document_count": document_count}, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_sections ( + section_id, user_id, namespace, document_id, job_result_id, + section_path, section_title, section_level, sort_order, created_at + ) + SELECT + 'section_batch_' || :identifier || '_' || i, + :user_id, + :namespace, + 'doc_batch_' || :identifier || '_' || i, + 'result_batch_' || :identifier || '_' || i, + 'batch-' || i || '.pdf/section', + 'section', + 1, + 1, + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + { + "identifier": identifier, + "user_id": _USER_ID, + "namespace": namespace, + "document_count": document_count, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_chunks ( + id, chunk_id, user_id, namespace, document_id, job_result_id, + section_id, chunk_type, content, source_chunk_path, + chunk_metadata, sort_order, created_at + ) + SELECT + 'dchunk_batch_' || :identifier || '_' || i, + 'chunk_batch_' || :identifier || '_' || i, + :user_id, + :namespace, + 'doc_batch_' || :identifier || '_' || i, + 'result_batch_' || :identifier || '_' || i, + 'section_batch_' || :identifier || '_' || i, + 'text', + 'content-' || i, + 'batch-' || i || '.pdf/section/chunk', + '{}'::json, + 1, + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + { + "identifier": identifier, + "user_id": _USER_ID, + "namespace": namespace, + "document_count": document_count, + }, + ) + + +async def test_snapshot_keeps_sections_and_chunks_on_the_captured_revision( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + namespace = f"revision-race-{uuid4().hex[:8]}" + async with developer_api_client_factory(): + document_id, new_result_id = await _seed_republished_document(namespace) + + async def publish_new_revision() -> None: + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = :new_result_id + WHERE document_id = :document_id + """, + {"new_result_id": new_result_id, "document_id": document_id}, + ) + + async with _contract_db_session() as db: + publishing_db = cast( + AsyncSession, + _PublishingSession(db, publish_new_revision), + ) + snapshot = await load_nav_snapshot( + publishing_db, + user_id=_USER_ID, + namespace=namespace, + ) + + section_ids = list(snapshot.provider.children(document_id)) + chunks = [ + chunk + for section_id in section_ids + for chunk in snapshot.provider.self_units(section_id) + ] + assert [chunk.content for chunk in chunks] == ["old content"] + assert snapshot.chunk_ref_index[chunks[0].chunk_id]["section_path"] == ( + "republished.pdf/old" + ) + + +async def test_snapshot_batches_chunks_across_many_small_documents( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_count = 25 + namespace = f"batch-documents-{uuid4().hex[:8]}" + async with developer_api_client_factory(): + await _seed_many_small_documents(namespace, document_count=document_count) + async with _contract_db_session() as db: + counting_db = _CountingSession(db) + snapshot = await load_nav_snapshot( + cast(AsyncSession, counting_db), + user_id=_USER_ID, + namespace=namespace, + ) + + assert len(snapshot.document_ids) == document_count + assert counting_db.chunk_query_count == 1 + + async def _load_legacy_rows( namespace: str, *, @@ -226,12 +576,16 @@ async def _load_legacy_rows( (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id), ) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .outerjoin( + DocumentSection, DocumentSection.section_id == DocumentChunk.section_id + ) .outerjoin(JobResult, JobResult.id == DocumentChunk.job_result_id) .where(Document.user_id == _USER_ID) .where(Document.namespace == namespace) .where(Document.status == "active") - .order_by(Document.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id) + .order_by( + Document.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id + ) ) async with _contract_db_session() as db: if statement_timeout_ms is not None: @@ -427,7 +781,7 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_sql_optimization( await _load_legacy_rows(namespace, statement_timeout_ms=100) async with _contract_db_session() as db: - await db.execute(text("SET LOCAL statement_timeout = 100")) + await db.execute(text("SET LOCAL statement_timeout = 5000")) bounded_snapshot = await load_nav_snapshot( db, user_id=_USER_ID, @@ -447,10 +801,7 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_sql_optimization( } assert len(optimized_chunks) == _TOTAL_CHUNKS - legacy_by_chunk = { - str(row[1]): row - for row in legacy_rows - } + legacy_by_chunk = {str(row[1]): row for row in legacy_rows} for chunk_id, (document_id, chunk) in optimized_chunks.items(): legacy_row = legacy_by_chunk[chunk_id] assert document_id == str(legacy_row[0]) @@ -460,7 +811,9 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_sql_optimization( assert chunk.sort_order == int(legacy_row[5]) assert chunk.source_chunk_path == str(legacy_row[6]) assert chunk.file_path == str(legacy_row[7] or "") - assert chunk.metadata == (legacy_row[8] if isinstance(legacy_row[8], dict) else {}) + assert chunk.metadata == ( + legacy_row[8] if isinstance(legacy_row[8], dict) else {} + ) reference = snapshot.chunk_ref_index[f"{document_id}:{chunk_id}"] assert reference["document_id"] == document_id diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 830a2cbd1..7cccaa115 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -199,6 +199,27 @@ def test_should_seed_v2_job_polling_system_limit( assert result["description"] == "Job queries - prevent polling" +def test_should_index_document_chunks_in_snapshot_pagination_order( + migrated_head_engine: Engine, +) -> None: + with migrated_head_engine.begin() as connection: + index_definition = connection.execute( + text( + """ + SELECT indexdef + FROM pg_indexes + WHERE schemaname = current_schema() + AND tablename = 'document_chunks' + AND indexname = 'idx_document_chunks_revision_snapshot_order' + """ + ) + ).scalar_one() + + assert "(document_id, job_result_id, sort_order, chunk_id, id)" in str( + index_definition + ) + + def test_api_standalone_mode_should_create_auth_user_table_before_migrations( standalone_alembic_engine: Engine, ) -> None: diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index c9af71407..f9681e178 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -202,6 +202,14 @@ class DocumentChunk(Base): Index("idx_document_chunks_scope", "user_id", "namespace"), Index("idx_document_chunks_chunk_id", "chunk_id"), Index("idx_document_chunks_doc_revision", "document_id", "job_result_id"), + Index( + "idx_document_chunks_revision_snapshot_order", + "document_id", + "job_result_id", + "sort_order", + "chunk_id", + "id", + ), Index("idx_document_chunks_section", "section_id"), Index( "idx_chunk_content_search_tsv", diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 011405cf7..7d85b5241 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -123,9 +123,7 @@ async def load_nav_snapshot( sections_by_doc, section_path_by_id = await _load_sections( db, - user_id=user_id, - namespace=namespace, - exclude_document_ids=excluded_docs, + document_revisions=document_revisions, exclude_sections=excluded_secs, ) units_by_doc, chunk_ref_index = await _load_chunks( @@ -159,14 +157,13 @@ async def load_nav_snapshot( async def _load_sections( db: AsyncSession, *, - user_id: str, - namespace: str, - exclude_document_ids: list[str], + document_revisions: list[tuple[str, str]], exclude_sections: list[dict[str, str]], ) -> tuple[dict[str, list[SectionRow]], dict[str, str]]: + # Captured pairs replace DocumentSection.job_result_id == Document.current_job_result_id. stmt = ( select( - Document.document_id, + DocumentSection.document_id, DocumentSection.section_id, DocumentSection.parent_section_id, DocumentSection.section_path, @@ -175,18 +172,18 @@ async def _load_sections( DocumentSection.summary, DocumentSection.sort_order, ) - .join( - DocumentSection, - (DocumentSection.document_id == Document.document_id) - & (DocumentSection.job_result_id == Document.current_job_result_id), + .where( + tuple_( + DocumentSection.document_id, + DocumentSection.job_result_id, + ).in_(document_revisions) + ) + .order_by( + DocumentSection.document_id, + DocumentSection.sort_order, + DocumentSection.section_id, ) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == "active") - .order_by(Document.document_id, DocumentSection.sort_order, DocumentSection.section_id) ) - if exclude_document_ids: - stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) by_doc: dict[str, list[SectionRow]] = {} path_by_id: dict[str, str] = {} @@ -222,95 +219,110 @@ async def _load_chunks( section_path_by_id: dict[str, str], job_id_by_result_id: dict[str, str], ) -> tuple[dict[str, list[UnitRow]], dict[str, dict[str, Any]]]: + # Captured pairs replace DocumentChunk.document_id == document_id and DocumentChunk.job_result_id == job_result_id. by_doc: dict[str, list[UnitRow]] = {} ref_index: dict[str, dict[str, Any]] = {} - for document_id, job_result_id in document_revisions: - last_key: tuple[int, str, str] | None = None - while True: - stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.sort_order, - DocumentChunk.source_chunk_path, - DocumentChunk.file_path, - DocumentChunk.chunk_metadata, - DocumentChunk.id, + last_key: tuple[str, str, int, str, str] | None = None + while True: + stmt = ( + select( + DocumentChunk.document_id, + DocumentChunk.job_result_id, + DocumentChunk.chunk_id, + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.sort_order, + DocumentChunk.source_chunk_path, + DocumentChunk.file_path, + DocumentChunk.chunk_metadata, + DocumentChunk.id, + ) + .where( + tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( + document_revisions ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .order_by( + ) + .order_by( + DocumentChunk.document_id, + DocumentChunk.job_result_id, + DocumentChunk.sort_order, + DocumentChunk.chunk_id, + DocumentChunk.id, + ) + .limit(_CHUNK_BATCH_SIZE) + ) + if last_key is not None: + stmt = stmt.where( + tuple_( + DocumentChunk.document_id, + DocumentChunk.job_result_id, DocumentChunk.sort_order, DocumentChunk.chunk_id, DocumentChunk.id, ) - .limit(_CHUNK_BATCH_SIZE) - ) - if last_key is not None: - stmt = stmt.where( - tuple_( - DocumentChunk.sort_order, - DocumentChunk.chunk_id, - DocumentChunk.id, - ) - > tuple_( - literal(last_key[0]), - literal(last_key[1]), - literal(last_key[2]), - ) + > tuple_( + literal(last_key[0]), + literal(last_key[1]), + literal(last_key[2]), + literal(last_key[3]), + literal(last_key[4]), ) + ) - rows = (await db.execute(stmt)).all() - if not rows: - break - for row in rows: - chunk_id = str(row[0] or "").strip() - section_id = str(row[1]) if row[1] else None - section_path: str | None = ( - section_path_by_id.get(section_id) if section_id else None - ) - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - if section_id and section_id not in section_path_by_id: - continue + rows = (await db.execute(stmt)).all() + if not rows: + break + for row in rows: + document_id = str(row[0]) + job_result_id = str(row[1]) + chunk_id = str(row[2] or "").strip() + section_id = str(row[3]) if row[3] else None + section_path: str | None = ( + section_path_by_id.get(section_id) if section_id else None + ) + if is_excluded_section( + document_id=document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + if section_id and section_id not in section_path_by_id: + continue - unit = UnitRow( - chunk_id=chunk_id, - section_id=section_id, - chunk_type=str(row[2] or "text"), - content=str(row[3] or ""), - sort_order=int(row[4] or 0), - source_chunk_path=str(row[5] or ""), - file_path=str(row[6] or ""), - metadata=_as_meta(row[7]), - ) - by_doc.setdefault(document_id, []).append(unit) - if chunk_id: - meta = { - "document_id": document_id, - "section_path": section_path, - "chunk_type": unit.chunk_type, - "file_path": unit.file_path or None, - "job_id": job_id_by_result_id.get(job_result_id), - } - # Bare chunk_id (last-wins) plus doc-scoped key so the same - # chunk_id can appear under multiple documents. - ref_index[chunk_id] = meta - ref_index[f"{document_id}:{chunk_id}"] = meta - last_row = rows[-1] - last_key = ( - int(last_row[4] or 0), - str(last_row[0] or ""), - str(last_row[8]), + unit = UnitRow( + chunk_id=chunk_id, + section_id=section_id, + chunk_type=str(row[4] or "text"), + content=str(row[5] or ""), + sort_order=int(row[6] or 0), + source_chunk_path=str(row[7] or ""), + file_path=str(row[8] or ""), + metadata=_as_meta(row[9]), ) - if len(rows) < _CHUNK_BATCH_SIZE: - break + by_doc.setdefault(document_id, []).append(unit) + if chunk_id: + meta = { + "document_id": document_id, + "section_path": section_path, + "chunk_type": unit.chunk_type, + "file_path": unit.file_path or None, + "job_id": job_id_by_result_id.get(job_result_id), + } + # Bare chunk_id (last-wins) plus doc-scoped key so the same + # chunk_id can appear under multiple documents. + ref_index[chunk_id] = meta + ref_index[f"{document_id}:{chunk_id}"] = meta + last_row = rows[-1] + last_key = ( + str(last_row[0]), + str(last_row[1]), + int(last_row[6] or 0), + str(last_row[2] or ""), + str(last_row[10]), + ) + if len(rows) < _CHUNK_BATCH_SIZE: + break return by_doc, ref_index From 96c5794a6263f255bddd6c802ff89c0c9613fa63 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 18:55:08 +0800 Subject: [PATCH 04/14] perf: group revision snapshot queries --- ...etrieval_snapshot_large_corpus_contract.py | 2 + .../shared/services/retrieval/nav_snapshot.py | 189 +++++++++--------- 2 files changed, 100 insertions(+), 91 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index b465ea72d..65237e0ad 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -549,6 +549,8 @@ async def test_snapshot_batches_chunks_across_many_small_documents( ) assert len(snapshot.document_ids) == document_count + # Twenty-five tiny revisions should fit into one bounded SQL query, + # rather than one global query or one query per document. assert counting_db.chunk_query_count == 1 diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 7d85b5241..93e4771c1 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -27,6 +27,7 @@ # Keep each payload query bounded under the API's 30-second statement timeout. # Keyset pagination avoids the increasingly expensive OFFSET scans. _CHUNK_BATCH_SIZE = 2_000 +_REVISION_GROUP_SIZE = 32 @dataclass(frozen=True) @@ -219,110 +220,116 @@ async def _load_chunks( section_path_by_id: dict[str, str], job_id_by_result_id: dict[str, str], ) -> tuple[dict[str, list[UnitRow]], dict[str, dict[str, Any]]]: - # Captured pairs replace DocumentChunk.document_id == document_id and DocumentChunk.job_result_id == job_result_id. + # Captured pairs replace DocumentChunk.document_id == document_id and + # DocumentChunk.job_result_id == job_result_id. by_doc: dict[str, list[UnitRow]] = {} ref_index: dict[str, dict[str, Any]] = {} - last_key: tuple[str, str, int, str, str] | None = None - while True: - stmt = ( - select( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.chunk_id, - DocumentChunk.section_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.sort_order, - DocumentChunk.source_chunk_path, - DocumentChunk.file_path, - DocumentChunk.chunk_metadata, - DocumentChunk.id, - ) - .where( - tuple_(DocumentChunk.document_id, DocumentChunk.job_result_id).in_( - document_revisions + for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): + revision_group = document_revisions[ + group_start : group_start + _REVISION_GROUP_SIZE + ] + last_key: tuple[str, str, int, str, str] | None = None + while True: + stmt = ( + select( + DocumentChunk.document_id, + DocumentChunk.job_result_id, + DocumentChunk.chunk_id, + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.sort_order, + DocumentChunk.source_chunk_path, + DocumentChunk.file_path, + DocumentChunk.chunk_metadata, + DocumentChunk.id, ) - ) - .order_by( - DocumentChunk.document_id, - DocumentChunk.job_result_id, - DocumentChunk.sort_order, - DocumentChunk.chunk_id, - DocumentChunk.id, - ) - .limit(_CHUNK_BATCH_SIZE) - ) - if last_key is not None: - stmt = stmt.where( - tuple_( + .where( + tuple_( + DocumentChunk.document_id, + DocumentChunk.job_result_id, + ).in_(revision_group) + ) + .order_by( DocumentChunk.document_id, DocumentChunk.job_result_id, DocumentChunk.sort_order, DocumentChunk.chunk_id, DocumentChunk.id, ) - > tuple_( - literal(last_key[0]), - literal(last_key[1]), - literal(last_key[2]), - literal(last_key[3]), - literal(last_key[4]), - ) + .limit(_CHUNK_BATCH_SIZE) ) + if last_key is not None: + stmt = stmt.where( + tuple_( + DocumentChunk.document_id, + DocumentChunk.job_result_id, + DocumentChunk.sort_order, + DocumentChunk.chunk_id, + DocumentChunk.id, + ) + > tuple_( + literal(last_key[0]), + literal(last_key[1]), + literal(last_key[2]), + literal(last_key[3]), + literal(last_key[4]), + ) + ) - rows = (await db.execute(stmt)).all() - if not rows: - break - for row in rows: - document_id = str(row[0]) - job_result_id = str(row[1]) - chunk_id = str(row[2] or "").strip() - section_id = str(row[3]) if row[3] else None - section_path: str | None = ( - section_path_by_id.get(section_id) if section_id else None - ) - if is_excluded_section( - document_id=document_id, - section_path=section_path, - exclude_sections=exclude_sections, - ): - continue - if section_id and section_id not in section_path_by_id: - continue + rows = (await db.execute(stmt)).all() + if not rows: + break + for row in rows: + document_id = str(row[0]) + job_result_id = str(row[1]) + chunk_id = str(row[2] or "").strip() + section_id = str(row[3]) if row[3] else None + section_path: str | None = ( + section_path_by_id.get(section_id) if section_id else None + ) + if is_excluded_section( + document_id=document_id, + section_path=section_path, + exclude_sections=exclude_sections, + ): + continue + if section_id and section_id not in section_path_by_id: + continue - unit = UnitRow( - chunk_id=chunk_id, - section_id=section_id, - chunk_type=str(row[4] or "text"), - content=str(row[5] or ""), - sort_order=int(row[6] or 0), - source_chunk_path=str(row[7] or ""), - file_path=str(row[8] or ""), - metadata=_as_meta(row[9]), + unit = UnitRow( + chunk_id=chunk_id, + section_id=section_id, + chunk_type=str(row[4] or "text"), + content=str(row[5] or ""), + sort_order=int(row[6] or 0), + source_chunk_path=str(row[7] or ""), + file_path=str(row[8] or ""), + metadata=_as_meta(row[9]), + ) + by_doc.setdefault(document_id, []).append(unit) + if chunk_id: + meta = { + "document_id": document_id, + "section_path": section_path, + "chunk_type": unit.chunk_type, + "file_path": unit.file_path or None, + "job_id": job_id_by_result_id.get(job_result_id), + } + # Bare chunk_id (last-wins) plus doc-scoped key so the same + # chunk_id can appear under multiple documents. + ref_index[chunk_id] = meta + ref_index[f"{document_id}:{chunk_id}"] = meta + last_row = rows[-1] + last_key = ( + str(last_row[0]), + str(last_row[1]), + int(last_row[6] or 0), + str(last_row[2] or ""), + str(last_row[10]), ) - by_doc.setdefault(document_id, []).append(unit) - if chunk_id: - meta = { - "document_id": document_id, - "section_path": section_path, - "chunk_type": unit.chunk_type, - "file_path": unit.file_path or None, - "job_id": job_id_by_result_id.get(job_result_id), - } - # Bare chunk_id (last-wins) plus doc-scoped key so the same - # chunk_id can appear under multiple documents. - ref_index[chunk_id] = meta - ref_index[f"{document_id}:{chunk_id}"] = meta - last_row = rows[-1] - last_key = ( - str(last_row[0]), - str(last_row[1]), - int(last_row[6] or 0), - str(last_row[2] or ""), - str(last_row[10]), - ) - if len(rows) < _CHUNK_BATCH_SIZE: - break + if len(rows) < _CHUNK_BATCH_SIZE: + break return by_doc, ref_index From 94a519583bcdbd567edf0dc8422e9f9e065476ea Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 19:19:46 +0800 Subject: [PATCH 05/14] test: tighten retrieval snapshot contracts --- ...etrieval_snapshot_large_corpus_contract.py | 87 ++++++++++--------- .../shared/services/retrieval/nav_snapshot.py | 19 ++-- .../shared/tests/test_nav_snapshot.py | 15 ---- 3 files changed, 61 insertions(+), 60 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index 65237e0ad..bdcdde5ef 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -3,7 +3,7 @@ import time from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager -from typing import Any, cast +from typing import TypeAlias, cast from uuid import uuid4 import pytest @@ -11,7 +11,7 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult -from shared.services.retrieval.nav_snapshot import load_nav_snapshot +from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot from sqlalchemy import Executable, Result, select, text from sqlalchemy.engine import Row from sqlalchemy.exc import DBAPIError @@ -26,7 +26,12 @@ _CHUNKS_PER_DOCUMENT = 600 _SECTIONS_PER_DOCUMENT = 8 _CONTENT_BYTES = 2048 +_LEGACY_TIMEOUT_MS = 1 _TOTAL_CHUNKS = _DOCUMENT_COUNT * _CHUNKS_PER_DOCUMENT +JsonValue: TypeAlias = ( + None | bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"] +) +JsonObject: TypeAlias = dict[str, JsonValue] LegacySnapshotRow = Row[ tuple[ str, @@ -47,7 +52,7 @@ class _PublishingSession: def __init__( self, - session: AsyncSession, + session: SnapshotSession, publish_revision: Callable[[], Awaitable[None]], ) -> None: self._session = session @@ -63,7 +68,7 @@ async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: class _CountingSession: - def __init__(self, session: AsyncSession) -> None: + def __init__(self, session: SnapshotSession) -> None: self._session = session self.chunk_query_count = 0 @@ -509,10 +514,7 @@ async def publish_new_revision() -> None: ) async with _contract_db_session() as db: - publishing_db = cast( - AsyncSession, - _PublishingSession(db, publish_new_revision), - ) + publishing_db = _PublishingSession(db, publish_new_revision) snapshot = await load_nav_snapshot( publishing_db, user_id=_USER_ID, @@ -543,7 +545,7 @@ async def test_snapshot_batches_chunks_across_many_small_documents( async with _contract_db_session() as db: counting_db = _CountingSession(db) snapshot = await load_nav_snapshot( - cast(AsyncSession, counting_db), + counting_db, user_id=_USER_ID, namespace=namespace, ) @@ -599,7 +601,7 @@ async def _load_legacy_rows( return rows -async def _explain_snapshot_query(namespace: str) -> dict[str, Any]: +async def _explain_snapshot_query(namespace: str) -> JsonObject: row = await ContractDatabase.fetch_one( """ EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) @@ -639,7 +641,7 @@ async def _explain_snapshot_query(namespace: str) -> dict[str, Any]: return plan -async def _explain_revision_query() -> dict[str, Any]: +async def _explain_revision_query() -> JsonObject: row = await ContractDatabase.fetch_one( """ EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) @@ -687,7 +689,7 @@ async def _measure_seeded_payload(namespace: str) -> int: return row["payload_bytes"] -def _plan_nodes(plan_node: dict[str, Any]) -> list[str]: +def _plan_nodes(plan_node: JsonObject) -> list[str]: summary = ( f"{plan_node.get('Node Type')}" f"[{plan_node.get('Relation Name', '')}" @@ -739,32 +741,34 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_sql_optimization( f"time={revision_plan.get('Execution Time')}ms " f"nodes={' -> '.join(_plan_nodes(revision_root))}" ) - await ContractDatabase.execute( - """ - CREATE INDEX IF NOT EXISTS idx_benchmark_chunks_revision_order - ON document_chunks (document_id, job_result_id, sort_order, chunk_id, id) - """ - ) - indexed_explain_plan = await _explain_snapshot_query(namespace) - indexed_root = indexed_explain_plan["Plan"] - assert isinstance(indexed_root, dict) - print( - "snapshot EXPLAIN with candidate index: " - f"node={indexed_root.get('Node Type')} " - f"time={indexed_explain_plan.get('Execution Time')}ms " - f"nodes={' -> '.join(_plan_nodes(indexed_root))}" - ) - indexed_revision_plan = await _explain_revision_query() - indexed_revision_root = indexed_revision_plan["Plan"] - assert isinstance(indexed_revision_root, dict) - print( - "revision EXPLAIN with candidate index: " - f"time={indexed_revision_plan.get('Execution Time')}ms " - f"nodes={' -> '.join(_plan_nodes(indexed_revision_root))}" - ) - await ContractDatabase.execute( - "DROP INDEX IF EXISTS idx_benchmark_chunks_revision_order" - ) + try: + await ContractDatabase.execute( + """ + CREATE INDEX IF NOT EXISTS idx_benchmark_chunks_revision_order + ON document_chunks (document_id, job_result_id, sort_order, chunk_id, id) + """ + ) + indexed_explain_plan = await _explain_snapshot_query(namespace) + indexed_root = indexed_explain_plan["Plan"] + assert isinstance(indexed_root, dict) + print( + "snapshot EXPLAIN with candidate index: " + f"node={indexed_root.get('Node Type')} " + f"time={indexed_explain_plan.get('Execution Time')}ms " + f"nodes={' -> '.join(_plan_nodes(indexed_root))}" + ) + indexed_revision_plan = await _explain_revision_query() + indexed_revision_root = indexed_revision_plan["Plan"] + assert isinstance(indexed_revision_root, dict) + print( + "revision EXPLAIN with candidate index: " + f"time={indexed_revision_plan.get('Execution Time')}ms " + f"nodes={' -> '.join(_plan_nodes(indexed_revision_root))}" + ) + finally: + await ContractDatabase.execute( + "DROP INDEX IF EXISTS idx_benchmark_chunks_revision_order" + ) legacy_started = time.perf_counter() legacy_rows = await _load_legacy_rows(namespace) @@ -779,8 +783,13 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_sql_optimization( ) optimized_elapsed = time.perf_counter() - optimized_started + # A 1 ms budget makes the old unbounded statement deterministically + # cancel without asserting on machine-dependent elapsed wall time. with pytest.raises(DBAPIError, match="statement timeout"): - await _load_legacy_rows(namespace, statement_timeout_ms=100) + await _load_legacy_rows( + namespace, + statement_timeout_ms=_LEGACY_TIMEOUT_MS, + ) async with _contract_db_session() as db: await db.execute(text("SET LOCAL statement_timeout = 5000")) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 93e4771c1..6206ea2af 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -8,10 +8,10 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any +from typing import Any, Protocol -from sqlalchemy import literal, select, tuple_ -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy import Executable, literal, select, tuple_ +from sqlalchemy.engine import Result from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult @@ -30,6 +30,13 @@ _REVISION_GROUP_SIZE = 32 +class SnapshotSession(Protocol): + """Minimal database interface required by the snapshot loader.""" + + async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: + ... + + @dataclass(frozen=True) class NavSnapshot: """In-memory corpus for one map-nav episode.""" @@ -71,7 +78,7 @@ def build_nav_snapshot( async def load_nav_snapshot( - db: AsyncSession, + db: SnapshotSession, *, user_id: str, namespace: str, @@ -156,7 +163,7 @@ async def load_nav_snapshot( async def _load_sections( - db: AsyncSession, + db: SnapshotSession, *, document_revisions: list[tuple[str, str]], exclude_sections: list[dict[str, str]], @@ -213,7 +220,7 @@ async def _load_sections( async def _load_chunks( - db: AsyncSession, + db: SnapshotSession, *, document_revisions: list[tuple[str, str]], exclude_sections: list[dict[str, str]], diff --git a/packages/shared-python/shared/tests/test_nav_snapshot.py b/packages/shared-python/shared/tests/test_nav_snapshot.py index f8e094c26..f8bf76f61 100644 --- a/packages/shared-python/shared/tests/test_nav_snapshot.py +++ b/packages/shared-python/shared/tests/test_nav_snapshot.py @@ -93,18 +93,3 @@ def test_build_nav_snapshot_rejects_empty_corpus() -> None: units_by_doc={}, chunk_ref_index={}, ) - - -def test_load_nav_snapshot_filters_current_revision_only() -> None: - """Snapshot loading must select and query only each current revision.""" - import inspect - - from shared.services.retrieval import nav_snapshot as nav_snapshot_mod - - loader_src = "".join(inspect.getsource(nav_snapshot_mod.load_nav_snapshot).split()) - sections_src = "".join(inspect.getsource(nav_snapshot_mod._load_sections).split()) - chunks_src = "".join(inspect.getsource(nav_snapshot_mod._load_chunks).split()) - assert "Document.current_job_result_id" in loader_src - assert "DocumentSection.job_result_id==Document.current_job_result_id" in sections_src - assert "DocumentChunk.document_id==document_id" in chunks_src - assert "DocumentChunk.job_result_id==job_result_id" in chunks_src From 8d3654ba04b2eacdf9669d87c94b2cfcf05cd313 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 19:50:14 +0800 Subject: [PATCH 06/14] fix: make snapshot index migration non-blocking --- apps/api/alembic/env.py | 14 +++++++++-- ...c1d2e3f4_add_chunk_snapshot_order_index.py | 23 +++++++++++++------ ...etrieval_snapshot_large_corpus_contract.py | 6 ++--- .../tests/migrations/test_schema_contract.py | 6 ++--- .../shared/services/retrieval/nav_snapshot.py | 6 +++-- 5 files changed, 38 insertions(+), 17 deletions(-) diff --git a/apps/api/alembic/env.py b/apps/api/alembic/env.py index 7b9e5022c..b158ddbe6 100644 --- a/apps/api/alembic/env.py +++ b/apps/api/alembic/env.py @@ -101,6 +101,10 @@ def run_migrations_offline() -> None: include_object=include_object, literal_binds=True, dialect_opts={"paramstyle": "named"}, + # Keep each migration in its own transaction so migrations that + # require an autocommit block (for example CREATE INDEX + # CONCURRENTLY) can safely commit only their own predecessor. + transaction_per_migration=True, # Pass through configured SSL connect args. connect_args=ssl_connect_args, ) @@ -121,11 +125,17 @@ def run_migrations_online() -> None: def run_with_connection(connection: Connection) -> None: if settings.API_STANDALONE_MODE_ENABLED: ensure_better_auth_user_table(connection) + # The standalone bootstrap query starts SQLAlchemy's implicit + # transaction before Alembic begins tracking migration + # transactions. End it so autocommit migrations remain valid. + connection.commit() context.configure( connection=connection, target_metadata=target_metadata, include_object=include_object, + # Required for migrations that use autocommit_block(). + transaction_per_migration=True, ) with context.begin_transaction(): @@ -136,7 +146,7 @@ def run_with_connection(connection: Connection) -> None: return if isinstance(configured_connection, Engine): - with configured_connection.begin() as connection: + with configured_connection.connect() as connection: run_with_connection(connection) return @@ -149,7 +159,7 @@ def run_with_connection(connection: Connection) -> None: connect_args=ssl_connect_args, ) - with connectable.begin() as connection: + with connectable.connect() as connection: run_with_connection(connection) diff --git a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py index 474dd7cee..0696d44eb 100644 --- a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py +++ b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py @@ -22,13 +22,22 @@ def upgrade() -> None: - op.execute( - """ - CREATE INDEX IF NOT EXISTS idx_document_chunks_revision_snapshot_order - ON document_chunks (document_id, job_result_id, sort_order, chunk_id, id) - """ - ) + # Index creation must not hold a write lock on document_chunks while the + # production corpus is being indexed. CONCURRENTLY cannot run inside the + # transaction Alembic normally opens, so switch to an autocommit block. + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE INDEX CONCURRENTLY IF NOT EXISTS + idx_document_chunks_revision_snapshot_order + ON document_chunks (document_id, job_result_id, sort_order, chunk_id, id) + """ + ) def downgrade() -> None: - op.execute("DROP INDEX IF EXISTS idx_document_chunks_revision_snapshot_order") + with op.get_context().autocommit_block(): + op.execute( + "DROP INDEX CONCURRENTLY IF EXISTS " + "idx_document_chunks_revision_snapshot_order" + ) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index bdcdde5ef..fa8774989 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -706,7 +706,7 @@ def _plan_nodes(plan_node: JsonObject) -> list[str]: return nodes -async def test_large_snapshot_keeps_all_retrieval_inputs_after_sql_optimization( +async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], @@ -834,6 +834,6 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_sql_optimization( assert reference["job_id"] == str(legacy_row[10]) print( - f"large snapshot benchmark: legacy={legacy_elapsed:.3f}s " - f"optimized={optimized_elapsed:.3f}s chunks={_TOTAL_CHUNKS}" + f"large snapshot load: legacy={legacy_elapsed:.3f}s " + f"bounded={optimized_elapsed:.3f}s chunks={_TOTAL_CHUNKS}" ) diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 7cccaa115..bdd2f5a57 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -31,9 +31,9 @@ def _build_alembic_command_config(*, engine: Engine) -> Config: def _upgrade_to_heads(*, engine: Engine) -> None: config = _build_alembic_command_config(engine=engine) - with engine.begin() as connection: - config.attributes["connection"] = connection - command.upgrade(config, "heads") + # Let Alembic create its own connection. This is required for migrations + # that use PostgreSQL autocommit (for example CREATE INDEX CONCURRENTLY). + command.upgrade(config, "heads") def _insert_job( diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 6206ea2af..0c2163f5a 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -25,8 +25,10 @@ # Keep each payload query bounded under the API's 30-second statement timeout. -# Keyset pagination avoids the increasingly expensive OFFSET scans. -_CHUNK_BATCH_SIZE = 2_000 +# Ten-thousand-row keyset pages avoid OFFSET scans while reducing database +# round trips for production-shaped corpora. The contract benchmark verifies +# this page size against the full 2 KiB content and metadata payload. +_CHUNK_BATCH_SIZE = 10_000 _REVISION_GROUP_SIZE = 32 From 7176e8203327c87ba5b9b8daa0891df02b8cbcf7 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 20:28:18 +0800 Subject: [PATCH 07/14] fix: harden snapshot equivalence and index recovery --- ...c1d2e3f4_add_chunk_snapshot_order_index.py | 24 ++++++ ...etrieval_snapshot_large_corpus_contract.py | 81 ++++++++++++------- .../shared/services/retrieval/nav_snapshot.py | 6 +- 3 files changed, 81 insertions(+), 30 deletions(-) diff --git a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py index 0696d44eb..b144ba4ee 100644 --- a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py +++ b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py @@ -10,6 +10,7 @@ from typing import Sequence, Union from alembic import op +from sqlalchemy import text revision: str = "fbf0c1d2e3f4" @@ -26,6 +27,29 @@ def upgrade() -> None: # production corpus is being indexed. CONCURRENTLY cannot run inside the # transaction Alembic normally opens, so switch to an autocommit block. with op.get_context().autocommit_block(): + invalid_index = op.get_bind().execute( + text( + """ + SELECT EXISTS ( + SELECT 1 + FROM pg_class AS index_class + JOIN pg_namespace AS index_namespace + ON index_namespace.oid = index_class.relnamespace + JOIN pg_index AS index_metadata + ON index_metadata.indexrelid = index_class.oid + WHERE index_namespace.nspname = current_schema() + AND index_class.relname = + 'idx_document_chunks_revision_snapshot_order' + AND NOT index_metadata.indisvalid + ) + """ + ) + ).scalar_one() + if invalid_index: + op.execute( + "DROP INDEX CONCURRENTLY IF EXISTS " + "idx_document_chunks_revision_snapshot_order" + ) op.execute( """ CREATE INDEX CONCURRENTLY IF NOT EXISTS diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index fa8774989..c26bb0ae2 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -11,7 +11,12 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult -from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot +from shared.services.retrieval.nav_snapshot import ( + SnapshotSession, + _CHUNK_BATCH_SIZE, + _REVISION_GROUP_SIZE, + load_nav_snapshot, +) from sqlalchemy import Executable, Result, select, text from sqlalchemy.engine import Row from sqlalchemy.exc import DBAPIError @@ -776,12 +781,24 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( optimized_started = time.perf_counter() async with _contract_db_session() as db: + counting_db = _CountingSession(db) snapshot = await load_nav_snapshot( - db, + counting_db, user_id=_USER_ID, namespace=namespace, ) optimized_elapsed = time.perf_counter() - optimized_started + expected_chunk_query_count = sum( + ( + min(_REVISION_GROUP_SIZE, _DOCUMENT_COUNT - group_start) + * _CHUNKS_PER_DOCUMENT + + _CHUNK_BATCH_SIZE + - 1 + ) + // _CHUNK_BATCH_SIZE + for group_start in range(0, _DOCUMENT_COUNT, _REVISION_GROUP_SIZE) + ) + assert counting_db.chunk_query_count == expected_chunk_query_count # A 1 ms budget makes the old unbounded statement deterministically # cancel without asserting on machine-dependent elapsed wall time. @@ -803,35 +820,45 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( assert len(snapshot.document_ids) == _DOCUMENT_COUNT assert len(bounded_snapshot.document_ids) == _DOCUMENT_COUNT - optimized_chunks = { - chunk_id: (document_id, chunk) + optimized_rows = [ + ( + document_id, + chunk.chunk_id, + chunk.section_id, + chunk.chunk_type, + chunk.content, + chunk.sort_order, + chunk.source_chunk_path, + chunk.file_path, + chunk.metadata, + snapshot.chunk_ref_index[f"{document_id}:{chunk.chunk_id}"][ + "section_path" + ], + snapshot.chunk_ref_index[f"{document_id}:{chunk.chunk_id}"]["job_id"], + ) for document_id in snapshot.document_ids for section_id in snapshot.provider.children(document_id) for chunk in snapshot.provider.self_units(section_id) - for chunk_id in [chunk.chunk_id] - } - assert len(optimized_chunks) == _TOTAL_CHUNKS - - legacy_by_chunk = {str(row[1]): row for row in legacy_rows} - for chunk_id, (document_id, chunk) in optimized_chunks.items(): - legacy_row = legacy_by_chunk[chunk_id] - assert document_id == str(legacy_row[0]) - assert chunk.section_id == str(legacy_row[2]) - assert chunk.chunk_type == str(legacy_row[3]) - assert chunk.content == str(legacy_row[4]) - assert chunk.sort_order == int(legacy_row[5]) - assert chunk.source_chunk_path == str(legacy_row[6]) - assert chunk.file_path == str(legacy_row[7] or "") - assert chunk.metadata == ( - legacy_row[8] if isinstance(legacy_row[8], dict) else {} + ] + optimized_rows.sort(key=lambda row: (row[0], row[5], row[1], row[2] or "")) + legacy_rows_projected = [ + ( + str(row[0]), + str(row[1]), + str(row[2]) if row[2] else None, + str(row[3]), + str(row[4]), + int(row[5]), + str(row[6]), + str(row[7] or ""), + row[8] if isinstance(row[8], dict) else {}, + str(row[9] or ""), + str(row[10]) if row[10] else None, ) - - reference = snapshot.chunk_ref_index[f"{document_id}:{chunk_id}"] - assert reference["document_id"] == document_id - assert reference["section_path"] == str(legacy_row[9] or "") - assert reference["chunk_type"] == str(legacy_row[3]) - assert reference["file_path"] == (str(legacy_row[7]) if legacy_row[7] else None) - assert reference["job_id"] == str(legacy_row[10]) + for row in legacy_rows + ] + assert len(optimized_rows) == _TOTAL_CHUNKS + assert optimized_rows == legacy_rows_projected print( f"large snapshot load: legacy={legacy_elapsed:.3f}s " diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 0c2163f5a..e8c327c38 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -25,9 +25,9 @@ # Keep each payload query bounded under the API's 30-second statement timeout. -# Ten-thousand-row keyset pages avoid OFFSET scans while reducing database -# round trips for production-shaped corpora. The contract benchmark verifies -# this page size against the full 2 KiB content and metadata payload. +# Ten-thousand-row keyset pages avoid OFFSET scans while keeping each payload +# statement bounded. The contract benchmark verifies this page size against +# the full 2 KiB content and metadata payload. _CHUNK_BATCH_SIZE = 10_000 _REVISION_GROUP_SIZE = 32 From 73baa77af7551dfc82518c2be00580873977f1dc Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 20:35:34 +0800 Subject: [PATCH 08/14] test: compare snapshot rows deterministically --- ...t_retrieval_snapshot_large_corpus_contract.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index c26bb0ae2..d1d9432c6 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -857,6 +857,22 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( ) for row in legacy_rows ] + def row_order_key(row: tuple[object, ...]) -> tuple[str, ...]: + return ( + str(row[0]), + str(row[5]), + str(row[1]), + str(row[2] or ""), + str(row[3]), + str(row[4]), + str(row[6]), + str(row[7] or ""), + str(row[8]), + str(row[9] or ""), + str(row[10] or ""), + ) + legacy_rows_projected.sort(key=row_order_key) + optimized_rows.sort(key=row_order_key) assert len(optimized_rows) == _TOTAL_CHUNKS assert optimized_rows == legacy_rows_projected From 3ba3e36b7b18e95f6af65adae6d8e2a3aef26205 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 20:49:19 +0800 Subject: [PATCH 09/14] test: split retrieval snapshot contracts --- ...st_retrieval_snapshot_batching_contract.py | 194 ++++++++++ ...retrieval_snapshot_consistency_contract.py | 182 +++++++++ ...etrieval_snapshot_large_corpus_contract.py | 354 +----------------- .../support/retrieval_snapshot_support.py | 19 + 4 files changed, 401 insertions(+), 348 deletions(-) create mode 100644 apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py create mode 100644 apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py create mode 100644 apps/api/tests/support/retrieval_snapshot_support.py diff --git a/apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py new file mode 100644 index 000000000..a7239100a --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_snapshot_batching_contract.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from typing import cast +from uuid import uuid4 + +from httpx import AsyncClient +from sqlalchemy import Executable, Result +from sqlalchemy.sql.selectable import Select + +from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot +from tests.support.retrieval_snapshot_support import contract_db_session +from tests.support.contract_database import ContractDatabase + + +_USER_ID = "local-dev-user" + + +class _CountingSession: + def __init__(self, session: SnapshotSession) -> None: + self._session = session + self.chunk_query_count = 0 + + async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: + selected_tables = ( + { + getattr(table, "name", "") + for table in statement.get_final_froms() + } + if isinstance(statement, Select) + else set() + ) + if "document_chunks" in selected_tables: + self.chunk_query_count += 1 + result = await self._session.execute(statement) + return cast(Result[tuple[object, ...]], result) + + +async def _seed_many_small_documents(namespace: str, *, document_count: int) -> None: + identifier = uuid4().hex[:8] + await ContractDatabase.execute( + """ + INSERT INTO jobs ( + job_id, user_id, job_type, status, source_type, version, + webhook_enabled, created_at, updated_at, credits_charged, billing_status + ) + SELECT + 'job_batch_' || :identifier || '_' || i, + :user_id, + 'document_ingestion', + 'done', + 'file', + 0, + false, + NOW(), + NOW(), + 0, + 'skipped' + FROM generate_series(1, :document_count) AS values(i) + """, + { + "identifier": identifier, + "user_id": _USER_ID, + "document_count": document_count, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO documents ( + document_id, user_id, namespace, status, current_job_result_id, + source_file_name, parse_track, created_at, updated_at + ) + SELECT + 'doc_batch_' || :identifier || '_' || i, + :user_id, + :namespace, + 'active', + NULL, + 'batch-' || i || '.pdf', + 'chunk', + NOW(), + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + { + "identifier": identifier, + "user_id": _USER_ID, + "namespace": namespace, + "document_count": document_count, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO job_results ( + id, job_id, document_id, delivery_mode, created_at, updated_at + ) + SELECT + 'result_batch_' || :identifier || '_' || i, + 'job_batch_' || :identifier || '_' || i, + 'doc_batch_' || :identifier || '_' || i, + 'inline', + NOW(), + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + {"identifier": identifier, "document_count": document_count}, + ) + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = 'result_batch_' || :identifier || '_' || i + FROM generate_series(1, :document_count) AS values(i) + WHERE document_id = 'doc_batch_' || :identifier || '_' || i + """, + {"identifier": identifier, "document_count": document_count}, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_sections ( + section_id, user_id, namespace, document_id, job_result_id, + section_path, section_title, section_level, sort_order, created_at + ) + SELECT + 'section_batch_' || :identifier || '_' || i, + :user_id, + :namespace, + 'doc_batch_' || :identifier || '_' || i, + 'result_batch_' || :identifier || '_' || i, + 'batch-' || i || '.pdf/section', + 'section', + 1, + 1, + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + { + "identifier": identifier, + "user_id": _USER_ID, + "namespace": namespace, + "document_count": document_count, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_chunks ( + id, chunk_id, user_id, namespace, document_id, job_result_id, + section_id, chunk_type, content, source_chunk_path, + chunk_metadata, sort_order, created_at + ) + SELECT + 'dchunk_batch_' || :identifier || '_' || i, + 'chunk_batch_' || :identifier || '_' || i, + :user_id, + :namespace, + 'doc_batch_' || :identifier || '_' || i, + 'result_batch_' || :identifier || '_' || i, + 'section_batch_' || :identifier || '_' || i, + 'text', + 'content-' || i, + 'batch-' || i || '.pdf/section/chunk', + '{}'::json, + 1, + NOW() + FROM generate_series(1, :document_count) AS values(i) + """, + { + "identifier": identifier, + "user_id": _USER_ID, + "namespace": namespace, + "document_count": document_count, + }, + ) + + +async def test_snapshot_batches_chunks_across_many_small_documents( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_count = 25 + namespace = f"batch-documents-{uuid4().hex[:8]}" + async with developer_api_client_factory(): + await _seed_many_small_documents(namespace, document_count=document_count) + async with contract_db_session() as db: + counting_db = _CountingSession(db) + snapshot = await load_nav_snapshot( + counting_db, + user_id=_USER_ID, + namespace=namespace, + ) + + assert len(snapshot.document_ids) == document_count + assert counting_db.chunk_query_count == 1 diff --git a/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py new file mode 100644 index 000000000..64aaf4c0b --- /dev/null +++ b/apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from contextlib import AbstractAsyncContextManager +from typing import cast +from uuid import uuid4 + +from httpx import AsyncClient +from sqlalchemy import Executable, Result + +from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot +from tests.support.retrieval_snapshot_support import contract_db_session +from tests.support.contract_database import ContractDatabase + + +_USER_ID = "local-dev-user" + + +class _PublishingSession: + def __init__( + self, + session: SnapshotSession, + publish_revision: Callable[[], Awaitable[None]], + ) -> None: + self._session = session + self._publish_revision = publish_revision + self._has_published = False + + async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: + result = await self._session.execute(statement) + if not self._has_published: + self._has_published = True + await self._publish_revision() + return cast(Result[tuple[object, ...]], result) + + +async def _seed_republished_document(namespace: str) -> tuple[str, str]: + identifier = uuid4().hex[:8] + document_id = f"doc_race_{identifier}" + old_result_id = f"result_race_old_{identifier}" + new_result_id = f"result_race_new_{identifier}" + await ContractDatabase.execute( + """ + INSERT INTO jobs ( + job_id, user_id, job_type, status, source_type, version, + webhook_enabled, created_at, updated_at, credits_charged, billing_status + ) VALUES + (:old_job_id, :user_id, 'document_ingestion', 'done', 'file', 0, + false, NOW(), NOW(), 0, 'skipped'), + (:new_job_id, :user_id, 'document_ingestion', 'done', 'file', 0, + false, NOW(), NOW(), 0, 'skipped') + """, + { + "old_job_id": f"job_race_old_{identifier}", + "new_job_id": f"job_race_new_{identifier}", + "user_id": _USER_ID, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO documents ( + document_id, user_id, namespace, status, current_job_result_id, + source_file_name, parse_track, created_at, updated_at + ) VALUES ( + :document_id, :user_id, :namespace, 'active', NULL, + 'republished.pdf', 'chunk', NOW(), NOW() + ) + """, + {"document_id": document_id, "user_id": _USER_ID, "namespace": namespace}, + ) + await ContractDatabase.execute( + """ + INSERT INTO job_results ( + id, job_id, document_id, delivery_mode, created_at, updated_at + ) VALUES + (:old_result_id, :old_job_id, :document_id, 'inline', NOW(), NOW()), + (:new_result_id, :new_job_id, :document_id, 'inline', NOW(), NOW()) + """, + { + "old_result_id": old_result_id, + "new_result_id": new_result_id, + "old_job_id": f"job_race_old_{identifier}", + "new_job_id": f"job_race_new_{identifier}", + "document_id": document_id, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_sections ( + section_id, user_id, namespace, document_id, job_result_id, + section_path, section_title, section_level, sort_order, created_at + ) VALUES + (:old_section_id, :user_id, :namespace, :document_id, :old_result_id, + 'republished.pdf/old', 'old', 1, 1, NOW()), + (:new_section_id, :user_id, :namespace, :document_id, :new_result_id, + 'republished.pdf/new', 'new', 1, 1, NOW()) + """, + { + "old_section_id": f"section_race_old_{identifier}", + "new_section_id": f"section_race_new_{identifier}", + "user_id": _USER_ID, + "namespace": namespace, + "document_id": document_id, + "old_result_id": old_result_id, + "new_result_id": new_result_id, + }, + ) + await ContractDatabase.execute( + """ + INSERT INTO document_chunks ( + id, chunk_id, user_id, namespace, document_id, job_result_id, + section_id, chunk_type, content, source_chunk_path, + chunk_metadata, sort_order, created_at + ) VALUES + (:old_id, :old_chunk_id, :user_id, :namespace, :document_id, + :old_result_id, :old_section_id, 'text', 'old content', + 'republished.pdf/old/chunk', '{}'::json, 1, NOW()), + (:new_id, :new_chunk_id, :user_id, :namespace, :document_id, + :new_result_id, :new_section_id, 'text', 'new content', + 'republished.pdf/new/chunk', '{}'::json, 1, NOW()) + """, + { + "old_id": f"dchunk_race_old_{identifier}", + "new_id": f"dchunk_race_new_{identifier}", + "old_chunk_id": f"chunk_race_old_{identifier}", + "new_chunk_id": f"chunk_race_new_{identifier}", + "user_id": _USER_ID, + "namespace": namespace, + "document_id": document_id, + "old_result_id": old_result_id, + "new_result_id": new_result_id, + "old_section_id": f"section_race_old_{identifier}", + "new_section_id": f"section_race_new_{identifier}", + }, + ) + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = :old_result_id + WHERE document_id = :document_id + """, + {"old_result_id": old_result_id, "document_id": document_id}, + ) + return document_id, new_result_id + + +async def test_snapshot_keeps_sections_and_chunks_on_the_captured_revision( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + namespace = f"revision-race-{uuid4().hex[:8]}" + async with developer_api_client_factory(): + document_id, new_result_id = await _seed_republished_document(namespace) + + async def publish_new_revision() -> None: + await ContractDatabase.execute( + """ + UPDATE documents + SET current_job_result_id = :new_result_id + WHERE document_id = :document_id + """, + {"new_result_id": new_result_id, "document_id": document_id}, + ) + + async with contract_db_session() as db: + snapshot = await load_nav_snapshot( + _PublishingSession(db, publish_new_revision), + user_id=_USER_ID, + namespace=namespace, + ) + + section_ids = list(snapshot.provider.children(document_id)) + chunks = [ + chunk + for section_id in section_ids + for chunk in snapshot.provider.self_units(section_id) + ] + assert [chunk.content for chunk in chunks] == ["old content"] + assert snapshot.chunk_ref_index[chunks[0].chunk_id]["section_path"] == ( + "republished.pdf/old" + ) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index d1d9432c6..ce0dec324 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -1,8 +1,8 @@ from __future__ import annotations import time -from collections.abc import AsyncGenerator, Awaitable, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager from typing import TypeAlias, cast from uuid import uuid4 @@ -20,10 +20,9 @@ from sqlalchemy import Executable, Result, select, text from sqlalchemy.engine import Row from sqlalchemy.exc import DBAPIError -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.sql.selectable import Select +from tests.support.retrieval_snapshot_support import contract_db_session from tests.support.contract_database import ContractDatabase -from shared.testing.contract_runtime import get_contract_database_url _USER_ID = "local-dev-user" @@ -54,24 +53,6 @@ ] -class _PublishingSession: - def __init__( - self, - session: SnapshotSession, - publish_revision: Callable[[], Awaitable[None]], - ) -> None: - self._session = session - self._publish_revision = publish_revision - self._has_published = False - - async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: - result = await self._session.execute(statement) - if not self._has_published: - self._has_published = True - await self._publish_revision() - return cast(Result[tuple[object, ...]], result) - - class _CountingSession: def __init__(self, session: SnapshotSession) -> None: self._session = session @@ -91,17 +72,6 @@ async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: return cast(Result[tuple[object, ...]], result) -@asynccontextmanager -async def _contract_db_session() -> AsyncGenerator[AsyncSession, None]: - engine = create_async_engine(get_contract_database_url(), future=True) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - try: - async with session_factory() as session: - yield session - finally: - await engine.dispose() - - async def _seed_large_retrieval_corpus(namespace: str) -> None: await ContractDatabase.execute( """ @@ -249,318 +219,6 @@ async def _seed_large_retrieval_corpus(namespace: str) -> None: ) -async def _seed_republished_document(namespace: str) -> tuple[str, str]: - identifier = uuid4().hex[:8] - document_id = f"doc_race_{identifier}" - old_result_id = f"result_race_old_{identifier}" - new_result_id = f"result_race_new_{identifier}" - await ContractDatabase.execute( - """ - INSERT INTO jobs ( - job_id, user_id, job_type, status, source_type, version, - webhook_enabled, created_at, updated_at, credits_charged, billing_status - ) VALUES - (:old_job_id, :user_id, 'document_ingestion', 'done', 'file', 0, - false, NOW(), NOW(), 0, 'skipped'), - (:new_job_id, :user_id, 'document_ingestion', 'done', 'file', 0, - false, NOW(), NOW(), 0, 'skipped') - """, - { - "old_job_id": f"job_race_old_{identifier}", - "new_job_id": f"job_race_new_{identifier}", - "user_id": _USER_ID, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO documents ( - document_id, user_id, namespace, status, current_job_result_id, - source_file_name, parse_track, created_at, updated_at - ) VALUES ( - :document_id, :user_id, :namespace, 'active', NULL, - 'republished.pdf', 'chunk', NOW(), NOW() - ) - """, - { - "document_id": document_id, - "user_id": _USER_ID, - "namespace": namespace, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO job_results ( - id, job_id, document_id, delivery_mode, created_at, updated_at - ) VALUES - (:old_result_id, :old_job_id, :document_id, 'inline', NOW(), NOW()), - (:new_result_id, :new_job_id, :document_id, 'inline', NOW(), NOW()) - """, - { - "old_result_id": old_result_id, - "new_result_id": new_result_id, - "old_job_id": f"job_race_old_{identifier}", - "new_job_id": f"job_race_new_{identifier}", - "document_id": document_id, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_sections ( - section_id, user_id, namespace, document_id, job_result_id, - section_path, section_title, section_level, sort_order, created_at - ) VALUES - (:old_section_id, :user_id, :namespace, :document_id, :old_result_id, - 'republished.pdf/old', 'old', 1, 1, NOW()), - (:new_section_id, :user_id, :namespace, :document_id, :new_result_id, - 'republished.pdf/new', 'new', 1, 1, NOW()) - """, - { - "old_section_id": f"section_race_old_{identifier}", - "new_section_id": f"section_race_new_{identifier}", - "user_id": _USER_ID, - "namespace": namespace, - "document_id": document_id, - "old_result_id": old_result_id, - "new_result_id": new_result_id, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_chunks ( - id, chunk_id, user_id, namespace, document_id, job_result_id, - section_id, chunk_type, content, source_chunk_path, - chunk_metadata, sort_order, created_at - ) VALUES - (:old_id, :old_chunk_id, :user_id, :namespace, :document_id, - :old_result_id, :old_section_id, 'text', 'old content', - 'republished.pdf/old/chunk', '{}'::json, 1, NOW()), - (:new_id, :new_chunk_id, :user_id, :namespace, :document_id, - :new_result_id, :new_section_id, 'text', 'new content', - 'republished.pdf/new/chunk', '{}'::json, 1, NOW()) - """, - { - "old_id": f"dchunk_race_old_{identifier}", - "new_id": f"dchunk_race_new_{identifier}", - "old_chunk_id": f"chunk_race_old_{identifier}", - "new_chunk_id": f"chunk_race_new_{identifier}", - "user_id": _USER_ID, - "namespace": namespace, - "document_id": document_id, - "old_result_id": old_result_id, - "new_result_id": new_result_id, - "old_section_id": f"section_race_old_{identifier}", - "new_section_id": f"section_race_new_{identifier}", - }, - ) - await ContractDatabase.execute( - """ - UPDATE documents - SET current_job_result_id = :old_result_id - WHERE document_id = :document_id - """, - {"old_result_id": old_result_id, "document_id": document_id}, - ) - return document_id, new_result_id - - -async def _seed_many_small_documents(namespace: str, *, document_count: int) -> None: - identifier = uuid4().hex[:8] - await ContractDatabase.execute( - """ - INSERT INTO jobs ( - job_id, user_id, job_type, status, source_type, version, - webhook_enabled, created_at, updated_at, credits_charged, billing_status - ) - SELECT - 'job_batch_' || :identifier || '_' || i, - :user_id, - 'document_ingestion', - 'done', - 'file', - 0, - false, - NOW(), - NOW(), - 0, - 'skipped' - FROM generate_series(1, :document_count) AS values(i) - """, - { - "identifier": identifier, - "user_id": _USER_ID, - "document_count": document_count, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO documents ( - document_id, user_id, namespace, status, current_job_result_id, - source_file_name, parse_track, created_at, updated_at - ) - SELECT - 'doc_batch_' || :identifier || '_' || i, - :user_id, - :namespace, - 'active', - NULL, - 'batch-' || i || '.pdf', - 'chunk', - NOW(), - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - { - "identifier": identifier, - "user_id": _USER_ID, - "namespace": namespace, - "document_count": document_count, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO job_results ( - id, job_id, document_id, delivery_mode, created_at, updated_at - ) - SELECT - 'result_batch_' || :identifier || '_' || i, - 'job_batch_' || :identifier || '_' || i, - 'doc_batch_' || :identifier || '_' || i, - 'inline', - NOW(), - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - {"identifier": identifier, "document_count": document_count}, - ) - await ContractDatabase.execute( - """ - UPDATE documents - SET current_job_result_id = 'result_batch_' || :identifier || '_' || i - FROM generate_series(1, :document_count) AS values(i) - WHERE document_id = 'doc_batch_' || :identifier || '_' || i - """, - {"identifier": identifier, "document_count": document_count}, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_sections ( - section_id, user_id, namespace, document_id, job_result_id, - section_path, section_title, section_level, sort_order, created_at - ) - SELECT - 'section_batch_' || :identifier || '_' || i, - :user_id, - :namespace, - 'doc_batch_' || :identifier || '_' || i, - 'result_batch_' || :identifier || '_' || i, - 'batch-' || i || '.pdf/section', - 'section', - 1, - 1, - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - { - "identifier": identifier, - "user_id": _USER_ID, - "namespace": namespace, - "document_count": document_count, - }, - ) - await ContractDatabase.execute( - """ - INSERT INTO document_chunks ( - id, chunk_id, user_id, namespace, document_id, job_result_id, - section_id, chunk_type, content, source_chunk_path, - chunk_metadata, sort_order, created_at - ) - SELECT - 'dchunk_batch_' || :identifier || '_' || i, - 'chunk_batch_' || :identifier || '_' || i, - :user_id, - :namespace, - 'doc_batch_' || :identifier || '_' || i, - 'result_batch_' || :identifier || '_' || i, - 'section_batch_' || :identifier || '_' || i, - 'text', - 'content-' || i, - 'batch-' || i || '.pdf/section/chunk', - '{}'::json, - 1, - NOW() - FROM generate_series(1, :document_count) AS values(i) - """, - { - "identifier": identifier, - "user_id": _USER_ID, - "namespace": namespace, - "document_count": document_count, - }, - ) - - -async def test_snapshot_keeps_sections_and_chunks_on_the_captured_revision( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - namespace = f"revision-race-{uuid4().hex[:8]}" - async with developer_api_client_factory(): - document_id, new_result_id = await _seed_republished_document(namespace) - - async def publish_new_revision() -> None: - await ContractDatabase.execute( - """ - UPDATE documents - SET current_job_result_id = :new_result_id - WHERE document_id = :document_id - """, - {"new_result_id": new_result_id, "document_id": document_id}, - ) - - async with _contract_db_session() as db: - publishing_db = _PublishingSession(db, publish_new_revision) - snapshot = await load_nav_snapshot( - publishing_db, - user_id=_USER_ID, - namespace=namespace, - ) - - section_ids = list(snapshot.provider.children(document_id)) - chunks = [ - chunk - for section_id in section_ids - for chunk in snapshot.provider.self_units(section_id) - ] - assert [chunk.content for chunk in chunks] == ["old content"] - assert snapshot.chunk_ref_index[chunks[0].chunk_id]["section_path"] == ( - "republished.pdf/old" - ) - - -async def test_snapshot_batches_chunks_across_many_small_documents( - developer_api_client_factory: Callable[ - [], AbstractAsyncContextManager[AsyncClient] - ], -) -> None: - document_count = 25 - namespace = f"batch-documents-{uuid4().hex[:8]}" - async with developer_api_client_factory(): - await _seed_many_small_documents(namespace, document_count=document_count) - async with _contract_db_session() as db: - counting_db = _CountingSession(db) - snapshot = await load_nav_snapshot( - counting_db, - user_id=_USER_ID, - namespace=namespace, - ) - - assert len(snapshot.document_ids) == document_count - # Twenty-five tiny revisions should fit into one bounded SQL query, - # rather than one global query or one query per document. - assert counting_db.chunk_query_count == 1 - - async def _load_legacy_rows( namespace: str, *, @@ -596,7 +254,7 @@ async def _load_legacy_rows( Document.document_id, DocumentChunk.sort_order, DocumentChunk.chunk_id ) ) - async with _contract_db_session() as db: + async with contract_db_session() as db: if statement_timeout_ms is not None: await db.execute( text("SELECT set_config('statement_timeout', :timeout_value, true)"), @@ -780,7 +438,7 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( legacy_elapsed = time.perf_counter() - legacy_started optimized_started = time.perf_counter() - async with _contract_db_session() as db: + async with contract_db_session() as db: counting_db = _CountingSession(db) snapshot = await load_nav_snapshot( counting_db, @@ -808,7 +466,7 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( statement_timeout_ms=_LEGACY_TIMEOUT_MS, ) - async with _contract_db_session() as db: + async with contract_db_session() as db: await db.execute(text("SET LOCAL statement_timeout = 5000")) bounded_snapshot = await load_nav_snapshot( db, diff --git a/apps/api/tests/support/retrieval_snapshot_support.py b/apps/api/tests/support/retrieval_snapshot_support.py new file mode 100644 index 000000000..b06b55f15 --- /dev/null +++ b/apps/api/tests/support/retrieval_snapshot_support.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager + +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from shared.testing.contract_runtime import get_contract_database_url + + +@asynccontextmanager +async def contract_db_session() -> AsyncGenerator[AsyncSession, None]: + engine = create_async_engine(get_contract_database_url(), future=True) + session_factory = async_sessionmaker(engine, expire_on_commit=False) + try: + async with session_factory() as session: + yield session + finally: + await engine.dispose() From d40191f5863e39a0aea34ce721517cd2bf3c390c Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 27 Aug 2026 21:34:34 +0800 Subject: [PATCH 10/14] fix: harden snapshot migration contracts --- apps/api/alembic/env.py | 10 +- ...c1d2e3f4_add_chunk_snapshot_order_index.py | 113 +++++++--- ...etrieval_snapshot_large_corpus_contract.py | 200 +----------------- .../tests/migrations/test_schema_contract.py | 59 ++++++ 4 files changed, 147 insertions(+), 235 deletions(-) diff --git a/apps/api/alembic/env.py b/apps/api/alembic/env.py index b158ddbe6..8e3a5d5a7 100644 --- a/apps/api/alembic/env.py +++ b/apps/api/alembic/env.py @@ -123,12 +123,15 @@ def run_migrations_online() -> None: configured_connection = config.attributes.get("connection") def run_with_connection(connection: Connection) -> None: + caller_owned_transaction = connection.in_transaction() if settings.API_STANDALONE_MODE_ENABLED: ensure_better_auth_user_table(connection) # The standalone bootstrap query starts SQLAlchemy's implicit # transaction before Alembic begins tracking migration - # transactions. End it so autocommit migrations remain valid. - connection.commit() + # transactions. End only that transaction; never commit a + # transaction supplied by the caller. + if not caller_owned_transaction: + connection.commit() context.configure( connection=connection, @@ -136,6 +139,9 @@ def run_with_connection(connection: Connection) -> None: include_object=include_object, # Required for migrations that use autocommit_block(). transaction_per_migration=True, + # Concurrent DDL cannot run inside a transaction owned by the + # caller. Migrations use regular DDL for that compatibility path. + knowhere_external_transaction=caller_owned_transaction, ) with context.begin_transaction(): diff --git a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py index b144ba4ee..56b6f4e3f 100644 --- a/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py +++ b/apps/api/alembic/versions/fbf0c1d2e3f4_add_chunk_snapshot_order_index.py @@ -21,47 +21,90 @@ branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None +_INDEX_NAME = "idx_document_chunks_revision_snapshot_order" +_INDEX_COLUMNS = "(document_id, job_result_id, sort_order, chunk_id, id)" +_INDEX_METHOD_AND_COLUMNS = f"using btree {_INDEX_COLUMNS}" + + +def _read_index_definition() -> tuple[str, bool] | None: + row = op.get_bind().execute( + text( + """ + SELECT indexdef, index_metadata.indisvalid + FROM pg_indexes AS indexes + JOIN pg_class AS index_class + ON index_class.relname = indexes.indexname + JOIN pg_namespace AS index_namespace + ON index_namespace.oid = index_class.relnamespace + AND index_namespace.nspname = indexes.schemaname + JOIN pg_index AS index_metadata + ON index_metadata.indexrelid = index_class.oid + WHERE indexes.schemaname = current_schema() + AND indexes.tablename = 'document_chunks' + AND indexes.indexname = :index_name + """ + ), + {"index_name": _INDEX_NAME}, + ).first() + if row is None: + return None + return str(row[0]), bool(row[1]) + + +def _index_is_intended() -> bool: + definition = _read_index_definition() + if definition is None or not definition[1]: + return False + normalized_definition = " ".join(definition[0].lower().split()) + return ( + normalized_definition.startswith("create index ") + and _INDEX_METHOD_AND_COLUMNS in normalized_definition + and " where " not in normalized_definition + ) + + +def _drop_index(*, concurrently: bool) -> None: + concurrent_clause = "CONCURRENTLY " if concurrently else "" + op.execute(f"DROP INDEX {concurrent_clause}IF EXISTS {_INDEX_NAME}") + + +def _create_index(*, concurrently: bool) -> None: + concurrent_clause = "CONCURRENTLY " if concurrently else "" + op.execute( + f""" + CREATE INDEX {concurrent_clause}IF NOT EXISTS {_INDEX_NAME} + ON document_chunks {_INDEX_COLUMNS} + """ + ) + def upgrade() -> None: + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + if external_transaction: + if _read_index_definition() is not None and not _index_is_intended(): + _drop_index(concurrently=False) + if _read_index_definition() is None: + _create_index(concurrently=False) + return + # Index creation must not hold a write lock on document_chunks while the # production corpus is being indexed. CONCURRENTLY cannot run inside the # transaction Alembic normally opens, so switch to an autocommit block. with op.get_context().autocommit_block(): - invalid_index = op.get_bind().execute( - text( - """ - SELECT EXISTS ( - SELECT 1 - FROM pg_class AS index_class - JOIN pg_namespace AS index_namespace - ON index_namespace.oid = index_class.relnamespace - JOIN pg_index AS index_metadata - ON index_metadata.indexrelid = index_class.oid - WHERE index_namespace.nspname = current_schema() - AND index_class.relname = - 'idx_document_chunks_revision_snapshot_order' - AND NOT index_metadata.indisvalid - ) - """ - ) - ).scalar_one() - if invalid_index: - op.execute( - "DROP INDEX CONCURRENTLY IF EXISTS " - "idx_document_chunks_revision_snapshot_order" - ) - op.execute( - """ - CREATE INDEX CONCURRENTLY IF NOT EXISTS - idx_document_chunks_revision_snapshot_order - ON document_chunks (document_id, job_result_id, sort_order, chunk_id, id) - """ - ) + if _read_index_definition() is not None and not _index_is_intended(): + _drop_index(concurrently=True) + if _read_index_definition() is None: + _create_index(concurrently=True) def downgrade() -> None: - with op.get_context().autocommit_block(): - op.execute( - "DROP INDEX CONCURRENTLY IF EXISTS " - "idx_document_chunks_revision_snapshot_order" - ) + external_transaction = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + if external_transaction: + _drop_index(concurrently=False) + else: + with op.get_context().autocommit_block(): + _drop_index(concurrently=True) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index ce0dec324..a414e44b3 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -1,12 +1,10 @@ from __future__ import annotations -import time from collections.abc import Callable from contextlib import AbstractAsyncContextManager -from typing import TypeAlias, cast +from typing import cast from uuid import uuid4 -import pytest from httpx import AsyncClient from shared.models.database.document import Document, DocumentChunk, DocumentSection @@ -17,9 +15,8 @@ _REVISION_GROUP_SIZE, load_nav_snapshot, ) -from sqlalchemy import Executable, Result, select, text +from sqlalchemy import Executable, Result, select from sqlalchemy.engine import Row -from sqlalchemy.exc import DBAPIError from sqlalchemy.sql.selectable import Select from tests.support.retrieval_snapshot_support import contract_db_session from tests.support.contract_database import ContractDatabase @@ -29,13 +26,7 @@ _DOCUMENT_COUNT = 100 _CHUNKS_PER_DOCUMENT = 600 _SECTIONS_PER_DOCUMENT = 8 -_CONTENT_BYTES = 2048 -_LEGACY_TIMEOUT_MS = 1 _TOTAL_CHUNKS = _DOCUMENT_COUNT * _CHUNKS_PER_DOCUMENT -JsonValue: TypeAlias = ( - None | bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"] -) -JsonObject: TypeAlias = dict[str, JsonValue] LegacySnapshotRow = Row[ tuple[ str, @@ -221,8 +212,6 @@ async def _seed_large_retrieval_corpus(namespace: str) -> None: async def _load_legacy_rows( namespace: str, - *, - statement_timeout_ms: int | None = None, ) -> list[LegacySnapshotRow]: stmt = ( select( @@ -255,120 +244,10 @@ async def _load_legacy_rows( ) ) async with contract_db_session() as db: - if statement_timeout_ms is not None: - await db.execute( - text("SELECT set_config('statement_timeout', :timeout_value, true)"), - {"timeout_value": f"{statement_timeout_ms}ms"}, - ) rows: list[LegacySnapshotRow] = list((await db.execute(stmt)).all()) return rows -async def _explain_snapshot_query(namespace: str) -> JsonObject: - row = await ContractDatabase.fetch_one( - """ - EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) - SELECT - d.document_id, - c.chunk_id, - c.section_id, - c.chunk_type, - c.content, - c.sort_order, - c.source_chunk_path, - c.file_path, - c.chunk_metadata, - s.section_path, - r.job_id - FROM documents AS d - JOIN document_chunks AS c - ON c.document_id = d.document_id - AND c.job_result_id = d.current_job_result_id - LEFT JOIN document_sections AS s ON s.section_id = c.section_id - LEFT JOIN job_results AS r ON r.id = c.job_result_id - WHERE d.user_id = :user_id - AND d.namespace = :namespace - AND d.status = 'active' - ORDER BY d.document_id, c.sort_order, c.chunk_id - """, - {"user_id": _USER_ID, "namespace": namespace}, - ) - if row is None: - raise AssertionError("EXPLAIN returned no plan") - plan_value = next(iter(row.values())) - if not isinstance(plan_value, list) or not plan_value: - raise AssertionError("EXPLAIN returned an unexpected plan shape") - plan = plan_value[0] - if not isinstance(plan, dict): - raise AssertionError("EXPLAIN plan root is not an object") - return plan - - -async def _explain_revision_query() -> JsonObject: - row = await ContractDatabase.fetch_one( - """ - EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) - SELECT - chunk_id, - section_id, - chunk_type, - content, - sort_order, - source_chunk_path, - file_path, - chunk_metadata, - id - FROM document_chunks - WHERE document_id = 'doc_lg_1' - AND job_result_id = 'result_lg_1' - ORDER BY sort_order, chunk_id, id - LIMIT 2000 - """ - ) - if row is None: - raise AssertionError("revision EXPLAIN returned no plan") - plan_value = next(iter(row.values())) - if not isinstance(plan_value, list) or not plan_value: - raise AssertionError("revision EXPLAIN returned an unexpected plan shape") - plan = plan_value[0] - if not isinstance(plan, dict): - raise AssertionError("revision EXPLAIN plan root is not an object") - return plan - - -async def _measure_seeded_payload(namespace: str) -> int: - row = await ContractDatabase.fetch_one( - """ - SELECT COALESCE( - SUM(pg_column_size(content) + pg_column_size(chunk_metadata)), 0 - ) AS payload_bytes - FROM document_chunks - WHERE user_id = :user_id AND namespace = :namespace - """, - {"user_id": _USER_ID, "namespace": namespace}, - ) - if row is None or not isinstance(row.get("payload_bytes"), int): - raise AssertionError("payload measurement returned an unexpected value") - return row["payload_bytes"] - - -def _plan_nodes(plan_node: JsonObject) -> list[str]: - summary = ( - f"{plan_node.get('Node Type')}" - f"[{plan_node.get('Relation Name', '')}" - f"/{plan_node.get('Index Name', '')}]" - f"={plan_node.get('Actual Total Time')}ms" - ) - children = plan_node.get("Plans", []) - if not isinstance(children, list): - return [summary] - nodes = [summary] - for child in children: - if isinstance(child, dict): - nodes.extend(_plan_nodes(child)) - return nodes - - async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] @@ -377,67 +256,7 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( namespace = f"large-corpus-{uuid4().hex[:8]}" async with developer_api_client_factory(): await _seed_large_retrieval_corpus(namespace) - payload_bytes = await _measure_seeded_payload(namespace) - print( - "seeded payload: " - f"{payload_bytes / (1024 * 1024):.1f}MiB " - f"content_bytes={_CONTENT_BYTES} " - f"sections_per_document={_SECTIONS_PER_DOCUMENT}" - ) - - explain_plan = await _explain_snapshot_query(namespace) - explain_root = explain_plan["Plan"] - assert isinstance(explain_root, dict) - print( - "snapshot EXPLAIN: " - f"node={explain_root.get('Node Type')} " - f"time={explain_plan.get('Execution Time')}ms " - f"shared_hit={explain_root.get('Shared Hit Blocks')} " - f"shared_read={explain_root.get('Shared Read Blocks')} " - f"nodes={' -> '.join(_plan_nodes(explain_root))}" - ) - revision_plan = await _explain_revision_query() - revision_root = revision_plan["Plan"] - assert isinstance(revision_root, dict) - print( - "revision EXPLAIN: " - f"time={revision_plan.get('Execution Time')}ms " - f"nodes={' -> '.join(_plan_nodes(revision_root))}" - ) - try: - await ContractDatabase.execute( - """ - CREATE INDEX IF NOT EXISTS idx_benchmark_chunks_revision_order - ON document_chunks (document_id, job_result_id, sort_order, chunk_id, id) - """ - ) - indexed_explain_plan = await _explain_snapshot_query(namespace) - indexed_root = indexed_explain_plan["Plan"] - assert isinstance(indexed_root, dict) - print( - "snapshot EXPLAIN with candidate index: " - f"node={indexed_root.get('Node Type')} " - f"time={indexed_explain_plan.get('Execution Time')}ms " - f"nodes={' -> '.join(_plan_nodes(indexed_root))}" - ) - indexed_revision_plan = await _explain_revision_query() - indexed_revision_root = indexed_revision_plan["Plan"] - assert isinstance(indexed_revision_root, dict) - print( - "revision EXPLAIN with candidate index: " - f"time={indexed_revision_plan.get('Execution Time')}ms " - f"nodes={' -> '.join(_plan_nodes(indexed_revision_root))}" - ) - finally: - await ContractDatabase.execute( - "DROP INDEX IF EXISTS idx_benchmark_chunks_revision_order" - ) - - legacy_started = time.perf_counter() legacy_rows = await _load_legacy_rows(namespace) - legacy_elapsed = time.perf_counter() - legacy_started - - optimized_started = time.perf_counter() async with contract_db_session() as db: counting_db = _CountingSession(db) snapshot = await load_nav_snapshot( @@ -445,7 +264,6 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( user_id=_USER_ID, namespace=namespace, ) - optimized_elapsed = time.perf_counter() - optimized_started expected_chunk_query_count = sum( ( min(_REVISION_GROUP_SIZE, _DOCUMENT_COUNT - group_start) @@ -458,16 +276,7 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( ) assert counting_db.chunk_query_count == expected_chunk_query_count - # A 1 ms budget makes the old unbounded statement deterministically - # cancel without asserting on machine-dependent elapsed wall time. - with pytest.raises(DBAPIError, match="statement timeout"): - await _load_legacy_rows( - namespace, - statement_timeout_ms=_LEGACY_TIMEOUT_MS, - ) - async with contract_db_session() as db: - await db.execute(text("SET LOCAL statement_timeout = 5000")) bounded_snapshot = await load_nav_snapshot( db, user_id=_USER_ID, @@ -533,8 +342,3 @@ def row_order_key(row: tuple[object, ...]) -> tuple[str, ...]: optimized_rows.sort(key=row_order_key) assert len(optimized_rows) == _TOTAL_CHUNKS assert optimized_rows == legacy_rows_projected - - print( - f"large snapshot load: legacy={legacy_elapsed:.3f}s " - f"bounded={optimized_elapsed:.3f}s chunks={_TOTAL_CHUNKS}" - ) diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index bdd2f5a57..24047dccd 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -36,6 +36,19 @@ def _upgrade_to_heads(*, engine: Engine) -> None: command.upgrade(config, "heads") +def _upgrade_to_heads_with_external_connection(*, engine: Engine) -> None: + config = _build_alembic_command_config(engine=engine) + with engine.begin() as connection: + config.attributes["connection"] = connection + command.upgrade(config, "heads") + + +def _upgrade_to_snapshot_parents(*, engine: Engine) -> None: + config = _build_alembic_command_config(engine=engine) + command.upgrade(config, "f0d85d209e68") + command.upgrade(config, "fbe1c2d3e4f5") + + def _insert_job( connection: Connection, *, @@ -220,6 +233,52 @@ def test_should_index_document_chunks_in_snapshot_pagination_order( ) +def test_should_upgrade_with_a_caller_owned_connection( + alembic_engine: Engine, +) -> None: + _upgrade_to_heads_with_external_connection(engine=alembic_engine) + + +def test_standalone_upgrade_should_preserve_a_caller_owned_connection( + standalone_alembic_engine: Engine, +) -> None: + _upgrade_to_heads_with_external_connection(engine=standalone_alembic_engine) + + +def test_should_replace_a_same_named_index_with_the_wrong_definition( + alembic_engine: Engine, +) -> None: + _upgrade_to_snapshot_parents(engine=alembic_engine) + with alembic_engine.begin() as connection: + connection.execute( + text( + """ + CREATE INDEX idx_document_chunks_revision_snapshot_order + ON document_chunks (document_id) + """ + ) + ) + + _upgrade_to_heads(engine=alembic_engine) + + with alembic_engine.begin() as connection: + index_definition = connection.execute( + text( + """ + SELECT indexdef + FROM pg_indexes + WHERE schemaname = current_schema() + AND tablename = 'document_chunks' + AND indexname = 'idx_document_chunks_revision_snapshot_order' + """ + ) + ).scalar_one() + + assert "(document_id, job_result_id, sort_order, chunk_id, id)" in str( + index_definition + ) + + def test_api_standalone_mode_should_create_auth_user_table_before_migrations( standalone_alembic_engine: Engine, ) -> None: From 38e04c08bcd46a47c0ba12aad4bed7d6de3992d9 Mon Sep 17 00:00:00 2001 From: cqboy1993 <167045138+EricNGOntos@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:23:11 +0800 Subject: [PATCH 11/14] feat: TOC rehome, page_memory scope isolation, unified token tracking * feat: enhance TOC processing and page exclusion logic - Introduced new functions for rehoming calibrated forests and handling pending records in the TOC anchoring process. - Updated the `compute_fat_leaf_pages` function to exclude TOC pages from the fat span test, ensuring accurate page calculations. - Refactored various functions to utilize the new `pages_excluding_toc` utility for consistent TOC page exclusion across multiple components. - Removed the deprecated `TocPagePolicy` class, streamlining the codebase. - Added tests to validate the exclusion of TOC pages in fat leaf calculations. * refactor: improve TOC rehome logic and enhance hierarchy resolution - Updated the hierarchy locator to better handle rehome matches and leaf nodes, ensuring accurate page range resolution. - Introduced new functions for managing rehome ranges and structural leaf identification, improving the clarity and efficiency of the code. - Enhanced the TOC anchoring process to support global paged-leaf rehome, allowing for more robust handling of TOC structures. - Refactored tests to validate the new rehome logic and ensure correct behavior in various scenarios. * refactor: enhance page rendering and inspection logic - Updated the rendering logic to support a new naming convention for PNG files, allowing for better organization based on the presence of a prefix. - Refactored the inspection process to render pages serially before conducting concurrent inspections, improving efficiency and clarity in the workflow. - Adjusted the handling of rendered pages in the inspection function to ensure proper coverage and error handling. - Enhanced tests to validate the new rendering and inspection behavior, ensuring robustness in various scenarios. * fix: pass toc_pages into outline skeleton and import extract_toc_nodes Unblocks make check before PR sync. Co-authored-by: Cursor * feat: enhance token tracking and usage reporting in document processing - Integrated token tracking into the concurrent page inspection process, allowing for detailed tracking of token usage during parsing. - Updated the summarization functions for DOCX and Excel formats to include usage task metadata, improving the granularity of token usage reporting. - Refactored the debug scripts to better manage token ledger operations, including loading and recording token usage across different stages. - Enhanced unit tests to validate the new token tracking features and ensure accurate reporting of token usage in various scenarios. * fix: stub pending TOC calibration internals in rehome contract Patching _calibrate_pending_tocs by string path was unreliable in the full contract suite, so the real path tried to render a missing pdf and dropped the pending graft record. Align with toc_graft stubs instead. Co-authored-by: Cursor * fix: drop unused serialize imports from toc rehome contract Co-authored-by: Cursor --------- Co-authored-by: Cursor --- .../document_agent/calibration/scan.py | 75 +- .../structure/hierarchy_locator.py | 136 +++- .../document_agent/structure/toc_anchoring.py | 92 ++- .../document_agent/structure/toc_graft.py | 29 +- .../document_agent/structure/toc_rehome.py | 549 +++++++++++++++ .../document_agent/tools/inspect_pages.py | 46 +- .../app/services/document_agent/visual.py | 8 +- .../document_parser/formats/docx/parser.py | 7 +- .../formats/excel/table_parser.py | 1 + .../formats/pdf/pymupdf_subprocess.py | 25 +- .../services/page_memory/fine_hierarchy.py | 15 +- .../services/page_memory/memory_service.py | 56 +- .../services/page_memory/node_assembler.py | 21 +- .../app/services/page_memory/page_renderer.py | 23 +- apps/worker/scripts/_debug_token_ledger.py | 128 ++++ apps/worker/scripts/debug_text_track.py | 380 ++++++---- .../scripts/page_memory/_debug_pm_shared.py | 166 ++--- .../debug_pm_stage3_scope_fine_hierarchy.py | 279 +++++--- .../page_memory/debug_pm_stage4_assets.py | 32 +- .../debug_pm_stage5_tagging_finalize.py | 99 +-- .../scripts/page_memory/toc_page_policy.py | 33 - .../test_calibration_scan_contract.py | 97 +++ ...est_page_memory_fine_hierarchy_contract.py | 27 + ...est_page_memory_node_assembler_contract.py | 14 +- ...emory_page_renderer_scope_path_contract.py | 64 ++ .../test_parse_usage_task_contract.py | 67 ++ .../contract/test_toc_rehome_contract.py | 662 ++++++++++++++++++ .../test_debug_text_track_token_ledger.py | 67 ++ .../unit/test_processing_metadata_persist.py | 49 ++ .../shared/services/ai/summary/engine.py | 87 +-- 30 files changed, 2722 insertions(+), 612 deletions(-) create mode 100644 apps/worker/app/services/document_agent/structure/toc_rehome.py create mode 100644 apps/worker/scripts/_debug_token_ledger.py delete mode 100644 apps/worker/scripts/page_memory/toc_page_policy.py create mode 100644 apps/worker/tests/contract/test_page_memory_page_renderer_scope_path_contract.py create mode 100644 apps/worker/tests/contract/test_parse_usage_task_contract.py create mode 100644 apps/worker/tests/contract/test_toc_rehome_contract.py create mode 100644 apps/worker/tests/unit/test_debug_text_track_token_ledger.py diff --git a/apps/worker/app/services/document_agent/calibration/scan.py b/apps/worker/app/services/document_agent/calibration/scan.py index 7b773aa18..82a220f83 100644 --- a/apps/worker/app/services/document_agent/calibration/scan.py +++ b/apps/worker/app/services/document_agent/calibration/scan.py @@ -5,8 +5,10 @@ window, feeding each round's cursor into the next one, so a miss never re-opens pages that were already inspected. -Each round still covers ``window_schedule[i]`` pages, but pages are inspected -one-at-a-time concurrently (never batched into a single VLM call). +Each round still covers ``window_schedule[i]`` pages. PDF→PNG for the whole +window is rendered in one serial child-process call (same discipline as TOC +extract: never drive the gevent PyMuPDF pool from a ThreadPool). VLM inspect +then runs one page per call, concurrently. """ from __future__ import annotations @@ -25,6 +27,11 @@ ) from app.services.document_agent.manifest import ToolContext from app.services.document_agent.tools.inspect_pages import inspect_pages +from app.services.document_agent.visual import render_pages +from shared.services.ai.token_tracking import ( + bind_token_tracker, + get_current_token_tracker_root_id, +) DEFAULT_WINDOW_SCHEDULE: tuple[int, ...] = (2, 4, 6, 10) @@ -111,6 +118,7 @@ def _inspect_one_page( ctx: ToolContext, title: str, page: int, + rendered_page: dict[str, Any], ) -> PageInspectResult: result = inspect_pages( ctx, @@ -122,6 +130,7 @@ def _inspect_one_page( "folder_name": "calibration_scan", "prefix": "scan", "usage_task": "calibration.scan_title_forward", + "rendered_pages": [rendered_page], }, ) if result.status != "ok": @@ -142,14 +151,63 @@ def _inspect_pages_concurrent( title: str, pages: list[int], ) -> list[PageInspectResult]: - """Inspect each page alone; keep page order in the returned list.""" + """Serial window render, then concurrent single-page VLM inspect.""" + rendered = render_pages( + ctx, + pages, + folder_name="calibration_scan", + prefix="scan", + timeout=120, + ) + rendered_by_page = { + int(item["page"]): { + "page": int(item["page"]), + "png_path": str(item["png_path"]), + } + for item in rendered + if item.get("page") is not None and item.get("png_path") + } + missing = [page for page in pages if page not in rendered_by_page] + if missing: + return [ + PageInspectResult( + page=page, + found=False, + error=( + f"render failed for pages={missing}" + if page in missing + else "render incomplete" + ), + ) + for page in pages + ] + if len(pages) == 1: - return [_inspect_one_page(ctx=ctx, title=title, page=pages[0])] + page = pages[0] + return [ + _inspect_one_page( + ctx=ctx, + title=title, + page=page, + rendered_page=rendered_by_page[page], + ) + ] + + token_tracker_root_id = get_current_token_tracker_root_id() + + def _inspect_one_page_with_tracking(page: int) -> PageInspectResult: + with bind_token_tracker(token_tracker_root_id): + return _inspect_one_page( + ctx=ctx, + title=title, + page=page, + rendered_page=rendered_by_page[page], + ) by_page: dict[int, PageInspectResult] = {} with ThreadPoolExecutor(max_workers=len(pages)) as pool: futures = { - pool.submit(_inspect_one_page, ctx=ctx, title=title, page=page): page + pool.submit(_inspect_one_page_with_tracking, page): page for page in pages } for future in as_completed(futures): @@ -174,9 +232,10 @@ def scan_title_forward( """Scan forward from ``start_page`` until the title is found or rounds run out. Each round covers ``window_schedule[i]`` consecutive pages starting at the - cursor left by the previous round. Pages inside a round are inspected - concurrently, one page per VLM call; the earliest true page wins. A page - error without a hit is logged and the scan continues to the next window. + cursor left by the previous round. The window is rendered once serially, + then each page is inspected via VLM concurrently; the earliest true page + wins. A page error without a hit is logged and the scan continues to the + next window. """ scanned: list[int] = [] rounds: list[ScanRound] = [] diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index 512689944..99e6e67aa 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -136,11 +136,11 @@ def first_leaf_start_under( parent_titles: tuple[str, ...], match_overrides: dict[tuple[str, ...], TitleMatch], ) -> int | None: - """Min start page among located leaves under *node*; None if none located.""" + """Min start page among structural leaves; rehome attachments do not bound scope.""" min_page: int | None = None for leaf_path, _leaf in iter_leaf_title_nodes([node], parent_titles=parent_titles): match = match_overrides.get(leaf_path) - if match is None: + if match is None or _is_rehome_match(match): continue if min_page is None or match.page < min_page: min_page = match.page @@ -179,7 +179,12 @@ def resolve_hierarchy_page_ranges( match_overrides=match_overrides or {}, resolved=resolved, ) - return resolved + rehome_ranges = _attach_rehome_ranges( + nodes, + match_overrides=match_overrides or {}, + structural_ranges=resolved, + ) + return _ranges_in_tree_order(nodes, [*resolved, *rehome_ranges]) def coverage_by_path( @@ -240,10 +245,19 @@ def _resolve_siblings( match_overrides: dict[tuple[str, ...], TitleMatch], resolved: list[ResolvedHierarchyRange], ) -> None: + boundary_nodes = [ + node + for node in nodes + if not _is_rehome_leaf( + node, + (*parent_titles, node.title), + match_overrides, + ) + ] located: list[tuple[TitleNode, int, TitleMatch | None]] = [] lower_bound = parent_scope.start - for index, node in enumerate(nodes): + for index, node in enumerate(boundary_nodes): path_titles = (*parent_titles, node.title) pages = _allowed_pages_between(lower_bound, parent_scope.end, allowed_pages) match = _locate_match_for_node( @@ -259,9 +273,9 @@ def _resolve_siblings( located.append((node, start_page, match)) if match is not None: lower_bound = start_page - elif index + 1 < len(nodes): + elif index + 1 < len(boundary_nodes): next_match = _find_next_located_sibling( - nodes=nodes, + nodes=boundary_nodes, start_index=index + 1, lower_bound=lower_bound, parent_end=parent_scope.end, @@ -334,6 +348,114 @@ def _resolve_siblings( ) +def _is_rehome_match(match: TitleMatch | None) -> bool: + if match is None: + return False + return bool((match.evidence or {}).get("toc_rehome")) + + +def _is_rehome_leaf( + node: TitleNode, + path_titles: tuple[str, ...], + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> bool: + return not node.children and _is_rehome_match(match_overrides.get(path_titles)) + + +def _rehome_host_range( + *, + structural_ranges: list[ResolvedHierarchyRange], + physical_page: int, +) -> ResolvedHierarchyRange: + candidates = [ + item + for item in structural_ranges + if item.start_page <= physical_page <= item.end_page + ] + if not candidates: + raise ValueError( + "rehome leaf has no resolved host scope " + f"page={physical_page}" + ) + latest_start = max(item.start_page for item in candidates) + return next( + item + for item in reversed(candidates) + if item.start_page == latest_start + ) + + +def _attach_rehome_ranges( + nodes: list[TitleNode], + *, + match_overrides: dict[tuple[str, ...], TitleMatch], + structural_ranges: list[ResolvedHierarchyRange], + parent_titles: tuple[str, ...] = (), +) -> list[ResolvedHierarchyRange]: + attached: list[ResolvedHierarchyRange] = [] + for node in nodes: + path_titles = (*parent_titles, node.title) + if _is_rehome_leaf(node, path_titles, match_overrides): + match = match_overrides.get(path_titles) + if match is None: + raise ValueError(f"rehome leaf override missing path={path_titles!r}") + host_range = _rehome_host_range( + structural_ranges=structural_ranges, + physical_page=match.page, + ) + attached.append( + ResolvedHierarchyRange( + title=node.title, + level=node.level, + start_page=host_range.start_page, + end_page=host_range.end_page, + path_titles=path_titles, + match=match, + evidence={ + **_range_evidence(match), + "status": "rehome_attached", + "skeleton_kind": "rehome_attachment", + "scope_host_path": list(host_range.path_titles), + }, + ) + ) + continue + if node.children: + attached.extend( + _attach_rehome_ranges( + node.children, + match_overrides=match_overrides, + structural_ranges=structural_ranges, + parent_titles=path_titles, + ) + ) + return attached + + +def _ranges_in_tree_order( + nodes: list[TitleNode], + ranges: list[ResolvedHierarchyRange], +) -> list[ResolvedHierarchyRange]: + order = { + path: index + for index, (path, _node) in enumerate(_iter_title_nodes(nodes)) + } + return sorted(ranges, key=lambda item: order[item.path_titles]) + + +def _iter_title_nodes( + nodes: list[TitleNode], + *, + parent_titles: tuple[str, ...] = (), +) -> list[tuple[tuple[str, ...], TitleNode]]: + items: list[tuple[tuple[str, ...], TitleNode]] = [] + for node in nodes: + path_titles = (*parent_titles, node.title) + items.append((path_titles, node)) + items.extend(_iter_title_nodes(node.children, parent_titles=path_titles)) + return items + + def _locate_match_for_node( node: TitleNode, *, @@ -391,7 +513,7 @@ def _infer_start_from_descendant_overrides( min_match: TitleMatch | None = None for leaf_path, _leaf_node in leaves: m = match_overrides.get(leaf_path) - if m is None: + if m is None or _is_rehome_match(m): continue if m.page not in scope_pages: continue diff --git a/apps/worker/app/services/document_agent/structure/toc_anchoring.py b/apps/worker/app/services/document_agent/structure/toc_anchoring.py index 1e9ccae0b..3e8b57f8c 100644 --- a/apps/worker/app/services/document_agent/structure/toc_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -55,6 +55,10 @@ def run_toc_anchoring(ctx: ToolContext) -> None: beat the confirmed printed TOC pages on coverage, anchor from outline with physical overrides (no calibrate VLM). Otherwise keep the extracted TOC tree and run the existing VLM calibration path. + + After calibration (primary + pending), each forest runs global paged-leaf + rehome, then pending forests are classified/grafted. Blackboard writes happen + only after that sequence. """ from app.services.document_agent.calibration.orchestrator import ( anchor_hierarchy, @@ -115,6 +119,16 @@ def run_toc_anchoring(ctx: ToolContext) -> None: page_count=page_count, body_pages=body_pages, ) + + # scope-local TOC pre-pass: future hook (no-op this period) + resolve_nodes, skeleton_anchor = _rehome_calibrated_forest( + resolve_nodes, + skeleton_anchor, + toc_pages=toc_pages, + ) + _rehome_pending_records(pending_records, toc_pages=toc_pages) + + if pending_records: _assign_toc_relationships( root_anchor=skeleton_anchor, pending_records=pending_records, @@ -201,6 +215,7 @@ def _try_outline_anchoring_route( page_texts=page_texts, body_pages=body_pages, page_count=page_count, + toc_pages=toc_pages, ): return False @@ -221,6 +236,7 @@ def _write_outline_skeleton( page_texts: dict[int, str], body_pages: list[int], page_count: int, + toc_pages: list[int] | None = None, ) -> bool: """Anchor outline rows via physical overrides (no calibrate VLM). @@ -241,6 +257,12 @@ def _write_outline_skeleton( ctx=ctx, ) skeleton_anchor = replace(skeleton_anchor, source="pdf_outline") + # scope-local TOC pre-pass: future hook (no-op this period) + resolve_nodes, skeleton_anchor = _rehome_calibrated_forest( + resolve_nodes, + skeleton_anchor, + toc_pages=toc_pages, + ) ctx.blackboard.skeleton_anchor = serialize_skeleton_anchor(skeleton_anchor) ctx.blackboard.skeleton_nodes = [ serialize_title_node(node) for node in resolve_nodes @@ -256,6 +278,50 @@ def _write_outline_skeleton( return True +def _rehome_calibrated_forest( + nodes: list[TitleNode], + skeleton_anchor: SkeletonAnchor, + *, + toc_pages: list[int] | None = None, +) -> tuple[list[TitleNode], SkeletonAnchor]: + from app.services.document_agent.structure.toc_rehome import rehome_skeleton_forest + + rehomed_nodes, rehomed_anchor, _events = rehome_skeleton_forest( + nodes, + skeleton_anchor, + toc_pages=toc_pages, + ) + return rehomed_nodes, rehomed_anchor + + +def _rehome_pending_records( + pending_records: list[dict[str, Any]], + *, + toc_pages: list[int] | None = None, +) -> None: + from app.services.document_agent.structure.toc_rehome import rehome_skeleton_forest + + for record in pending_records: + nodes_raw = record.get("nodes") or [] + anchor_raw = record.get("skeleton_anchor") + if not isinstance(anchor_raw, dict) or not nodes_raw: + continue + nodes = [ + deserialize_title_node(node) + for node in nodes_raw + if isinstance(node, dict) + ] + if not nodes: + continue + rehomed_nodes, rehomed_anchor, _events = rehome_skeleton_forest( + nodes, + deserialize_skeleton_anchor(anchor_raw), + toc_pages=toc_pages, + ) + record["nodes"] = [serialize_title_node(node) for node in rehomed_nodes] + record["skeleton_anchor"] = serialize_skeleton_anchor(rehomed_anchor) + + def outline_physical_overrides( hierarchies: list[dict[str, Any]], ) -> dict[tuple[str, ...], TitleMatch]: @@ -348,9 +414,31 @@ def select_global_toc_hierarchies( return selected, pending, summary +def pages_excluding_toc(pages: Any, toc_pages: Any) -> list[int]: + """Return ``pages`` with TOC pages removed (stable order). + + Same exclusion source as TEXT-track MinerU ``exclude_pages`` and anchoring + ``body_pages``: ``toc_result.toc_pages``. + """ + excluded: set[int] = set() + for raw in toc_pages or []: + try: + excluded.add(int(raw)) + except (TypeError, ValueError): + continue + result: list[int] = [] + for raw in pages or []: + try: + page = int(raw) + except (TypeError, ValueError): + continue + if page not in excluded: + result.append(page) + return result + + def body_pages_excluding_toc(toc_pages: Any, page_count: int) -> list[int]: - excluded = {int(page) for page in (toc_pages or [])} - return [page for page in range(1, page_count + 1) if page not in excluded] + return pages_excluding_toc(range(1, page_count + 1), toc_pages) def pending_toc_body_scope( diff --git a/apps/worker/app/services/document_agent/structure/toc_graft.py b/apps/worker/app/services/document_agent/structure/toc_graft.py index be8991ce2..7e094aa77 100644 --- a/apps/worker/app/services/document_agent/structure/toc_graft.py +++ b/apps/worker/app/services/document_agent/structure/toc_graft.py @@ -199,7 +199,7 @@ def _graft_one_child( reason="no_covering_path", ) return primary_children - covering_node = _node_at_path(primary_children, covering) + covering_node = node_at_path(primary_children, covering) if covering_node is None: _record_skip( events=events, @@ -223,7 +223,7 @@ def _graft_one_child( events=events, ), ) - return _replace_node_at_path(primary_children, covering, updated_parent) + return replace_node_at_path(primary_children, covering, updated_parent) def _attach_new_child( @@ -238,7 +238,7 @@ def _attach_new_child( contained_overrides: dict[tuple[str, ...], TitleMatch], events: list[dict[str, Any]], ) -> list[TitleNode]: - grafted = _rebase_levels(child, (parent.level + 1) - child.level) + grafted = rebase_levels(child, (parent.level + 1) - child.level) new_path = (*parent_path, grafted.title) events.append( { @@ -248,14 +248,14 @@ def _attach_new_child( "start": start, } ) - _remap_overrides( + remap_overrides( contained_overrides=contained_overrides, primary_overrides=primary_overrides, old_prefix=contained_path, new_prefix=new_path, drop_root=False, ) - return _insert_by_start( + return insert_by_start( siblings=primary_children, new_node=grafted, start=start, @@ -301,7 +301,7 @@ def _record_skip( events.append(event) -def _remap_overrides( +def remap_overrides( *, contained_overrides: dict[tuple[str, ...], TitleMatch], primary_overrides: dict[tuple[str, ...], TitleMatch], @@ -309,6 +309,7 @@ def _remap_overrides( new_prefix: tuple[str, ...], drop_root: bool, ) -> None: + """Copy override entries under ``old_prefix`` onto ``new_prefix`` paths.""" prefix_len = len(old_prefix) for path, match in contained_overrides.items(): if path[:prefix_len] != old_prefix: @@ -320,15 +321,16 @@ def _remap_overrides( primary_overrides[new_path] = match -def _rebase_levels(node: TitleNode, delta: int) -> TitleNode: +def rebase_levels(node: TitleNode, delta: int) -> TitleNode: + """Shift ``node.level`` (and descendants) by ``delta``.""" return replace( node, level=node.level + delta, - children=[_rebase_levels(child, delta) for child in node.children], + children=[rebase_levels(child, delta) for child in node.children], ) -def _insert_by_start( +def insert_by_start( *, siblings: list[TitleNode], new_node: TitleNode, @@ -336,6 +338,7 @@ def _insert_by_start( parent_path: tuple[str, ...], primary_overrides: dict[tuple[str, ...], TitleMatch], ) -> list[TitleNode]: + """Insert ``new_node`` before the first sibling whose own page is ``> start``.""" insert_at = len(siblings) for index, sibling in enumerate(siblings): match = primary_overrides.get((*parent_path, sibling.title)) @@ -347,7 +350,8 @@ def _insert_by_start( return updated -def _node_at_path(nodes: list[TitleNode], path: tuple[str, ...]) -> TitleNode | None: +def node_at_path(nodes: list[TitleNode], path: tuple[str, ...]) -> TitleNode | None: + """Return the node at ``path`` under ``nodes``, or None if missing.""" current: list[TitleNode] = nodes node: TitleNode | None = None for title in path: @@ -358,11 +362,12 @@ def _node_at_path(nodes: list[TitleNode], path: tuple[str, ...]) -> TitleNode | return node -def _replace_node_at_path( +def replace_node_at_path( nodes: list[TitleNode], path: tuple[str, ...], new_node: TitleNode, ) -> list[TitleNode]: + """Replace the node at ``path`` with ``new_node`` (immutable list update).""" title, *rest = path updated: list[TitleNode] = [] for node in nodes: @@ -375,7 +380,7 @@ def _replace_node_at_path( updated.append( replace( node, - children=_replace_node_at_path(node.children, tuple(rest), new_node), + children=replace_node_at_path(node.children, tuple(rest), new_node), ) ) return updated diff --git a/apps/worker/app/services/document_agent/structure/toc_rehome.py b/apps/worker/app/services/document_agent/structure/toc_rehome.py new file mode 100644 index 000000000..090294fcc --- /dev/null +++ b/apps/worker/app/services/document_agent/structure/toc_rehome.py @@ -0,0 +1,549 @@ +"""Same-forest TOC rehome: global paged-leaf monotonic repair before classify.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any + +from loguru import logger + +from app.services.document_agent.structure.anchoring_primitives import SkeletonAnchor +from app.services.document_agent.structure.hierarchy_locator import TitleMatch, TitleNode +from app.services.document_agent.structure.toc_graft import ( + insert_by_start, + node_at_path, + rebase_levels, + replace_node_at_path, +) + +_LOG_PREFIX = "[profile.toc_rehome]" + + +@dataclass(frozen=True) +class RehomePlan: + source_path: tuple[str, ...] + physical_page: int + segment_index: int + source_toc_index: int + anchor_source_path: tuple[str, ...] + action: str = "moved" + + +@dataclass(frozen=True) +class RehomeResult: + nodes: list[TitleNode] + match_overrides: dict[tuple[str, ...], TitleMatch] + events: list[dict[str, Any]] + + +@dataclass(frozen=True) +class _LeafRef: + path: tuple[str, ...] + page: int + toc_index: int + + +def own_page( + path: tuple[str, ...], + match_overrides: dict[tuple[str, ...], TitleMatch], +) -> int | None: + """Node's own calibrated physical page; never inferred from descendants.""" + match = match_overrides.get(path) + if match is None or match.page is None: + return None + return int(match.page) + + +def rehome_forest( + nodes: list[TitleNode], + match_overrides: dict[tuple[str, ...], TitleMatch], + *, + toc_pages: list[int] | None = None, +) -> RehomeResult: + """Repair TOC-order backjumps among paged leaves within one calibrated forest.""" + overrides = dict(match_overrides) + toc_page_set = {int(page) for page in (toc_pages or [])} + plans = _collect_backjump_plans(nodes, overrides, toc_pages=toc_page_set) + if not plans: + return RehomeResult(nodes=list(nodes), match_overrides=overrides, events=[]) + + prune_plans = [plan for plan in plans if plan.action == "pruned"] + move_plans = [plan for plan in plans if plan.action == "moved"] + logger.info( + "{} plans={} prunes={} moves={}", + _LOG_PREFIX, + len(plans), + len(prune_plans), + len(move_plans), + ) + working = list(nodes) + events: list[dict[str, Any]] = [] + if prune_plans: + working, overrides, prune_events = _apply_prunes( + working, overrides, prune_plans + ) + events.extend(prune_events) + if move_plans: + working, overrides, move_events = _apply_plans( + working, overrides, move_plans + ) + events.extend(move_events) + return RehomeResult(nodes=working, match_overrides=overrides, events=events) + + +def rehome_skeleton_forest( + nodes: list[TitleNode], + anchor: SkeletonAnchor, + *, + toc_pages: list[int] | None = None, +) -> tuple[list[TitleNode], SkeletonAnchor, list[dict[str, Any]]]: + """Run ``rehome_forest`` and return an updated ``SkeletonAnchor``.""" + result = rehome_forest( + nodes, + anchor.match_overrides, + toc_pages=toc_pages, + ) + return ( + result.nodes, + replace(anchor, match_overrides=result.match_overrides), + result.events, + ) + + +def _collect_paged_leaves( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + parent_path: tuple[str, ...] = (), + *, + toc_index_start: int = 0, +) -> list[_LeafRef]: + """TOC-order list of paged leaves (nodes with no children and own page).""" + leaves: list[_LeafRef] = [] + toc_index = toc_index_start + for node in nodes: + path = (*parent_path, node.title) + if node.children: + child_leaves = _collect_paged_leaves( + list(node.children), + overrides, + path, + toc_index_start=toc_index, + ) + leaves.extend(child_leaves) + toc_index += len(child_leaves) + continue + page = own_page(path, overrides) + if page is None: + continue + leaves.append(_LeafRef(path=path, page=page, toc_index=toc_index)) + toc_index += 1 + return leaves + + +def _split_monotonic_segments(leaves: list[_LeafRef]) -> list[list[_LeafRef]]: + """Split TOC-order leaves at every physical-page backjump.""" + segments: list[list[_LeafRef]] = [] + current: list[_LeafRef] = [] + for leaf in leaves: + if current and leaf.page < current[-1].page: + segments.append(current) + current = [] + current.append(leaf) + if current: + segments.append(current) + return segments + + +def _nearest_leaf_in_segment( + *, + physical_page: int, + segment: list[_LeafRef], +) -> _LeafRef | None: + """Nearest leaf with ``page <= physical_page`` inside one fixed segment.""" + candidates = [ + leaf + for leaf in segment + if leaf.page <= physical_page + ] + if not candidates: + return None + return max(candidates, key=lambda leaf: (leaf.page, leaf.toc_index)) + + +def _first_segment_duplicate( + leaf: _LeafRef, + first_segment: list[_LeafRef], +) -> _LeafRef | None: + """First-segment leaf with the same hierarchy path and physical page, if any.""" + for candidate in first_segment: + if candidate.path == leaf.path and candidate.page == leaf.page: + return candidate + return None + + +def _collect_backjump_plans( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + *, + toc_pages: set[int], +) -> list[RehomePlan]: + """Plan every post-break segment against the first monotonic segment only.""" + segments = _split_monotonic_segments(_collect_paged_leaves(nodes, overrides)) + if len(segments) < 2: + return [] + first_segment = segments[0] + plans: list[RehomePlan] = [] + for segment_index in range(1, len(segments)): + for leaf in segments[segment_index]: + duplicate = _first_segment_duplicate(leaf, first_segment) + if duplicate is not None: + logger.info( + "{} prune same_path_page path={} page={} segment={}", + _LOG_PREFIX, + leaf.path, + leaf.page, + segment_index, + ) + plans.append( + RehomePlan( + source_path=leaf.path, + physical_page=leaf.page, + segment_index=segment_index, + source_toc_index=leaf.toc_index, + anchor_source_path=duplicate.path, + action="pruned", + ) + ) + continue + if leaf.page in toc_pages: + logger.info( + "{} skip toc_page path={} page={} segment={}", + _LOG_PREFIX, + leaf.path, + leaf.page, + segment_index, + ) + continue + anchor = _nearest_leaf_in_segment( + physical_page=leaf.page, + segment=first_segment, + ) + if anchor is None: + continue + plans.append( + RehomePlan( + source_path=leaf.path, + physical_page=leaf.page, + segment_index=segment_index, + source_toc_index=leaf.toc_index, + anchor_source_path=anchor.path, + action="moved", + ) + ) + return plans + + +def _apply_prunes( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + plans: list[RehomePlan], +) -> tuple[list[TitleNode], dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: + """Drop post-first-segment duplicates (same path + page as a first-segment leaf).""" + working = list(nodes) + events: list[dict[str, Any]] = [] + for plan in plans: + leaves = _collect_paged_leaves(working, overrides) + matches = [ + leaf + for leaf in leaves + if leaf.path == plan.source_path and leaf.page == plan.physical_page + ] + if len(matches) < 2: + raise ValueError( + f"{_LOG_PREFIX} prune miss path={plan.source_path!r} " + f"page={plan.physical_page} remaining={len(matches)}" + ) + target = matches[-1] + working, detached, detached_path = _detach_by_toc_index( + working, + overrides, + target_toc_index=target.toc_index, + ) + if detached is None or detached_path is None: + raise ValueError( + f"{_LOG_PREFIX} prune detach miss toc_index={target.toc_index}" + ) + remaining = [ + leaf + for leaf in _collect_paged_leaves(working, overrides) + if leaf.path == detached_path + ] + if not remaining: + overrides.pop(detached_path, None) + working, overrides = _drop_empty_ancestors( + working, + overrides, + detached_path[:-1], + ) + events.append( + { + "action": "pruned", + "source_path": list(detached_path), + "physical_page": plan.physical_page, + "segment_index": plan.segment_index, + "anchor_path": list(plan.anchor_source_path), + } + ) + logger.info( + "{} pruned {} page={} segment={}", + _LOG_PREFIX, + detached_path, + plan.physical_page, + plan.segment_index, + ) + return working, overrides, events + + +def _apply_plans( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + plans: list[RehomePlan], +) -> tuple[list[TitleNode], dict[tuple[str, ...], TitleMatch], list[dict[str, Any]]]: + working = list(nodes) + events: list[dict[str, Any]] = [] + current_paths = { + leaf.path: leaf.path for leaf in _collect_paged_leaves(working, overrides) + } + for plan in plans: + source_path = current_paths[plan.source_path] + anchor_path = current_paths[plan.anchor_source_path] + dest_parent = anchor_path[:-1] + if dest_parent and node_at_path(working, dest_parent) is None: + raise ValueError( + f"{_LOG_PREFIX} dest parent missing path={dest_parent!r}" + ) + + working, detached = _detach_node(working, source_path) + if detached is None: + raise ValueError( + f"{_LOG_PREFIX} detach miss source_path={source_path!r}" + ) + target_level = len(dest_parent) + 1 + detached = rebase_levels(detached, target_level - detached.level) + new_path = (*dest_parent, detached.title) + working = _insert_under_parent( + nodes=working, + parent_path=dest_parent, + new_node=detached, + start=plan.physical_page, + overrides=overrides, + ) + if new_path != source_path: + _move_override_prefix(overrides, source_path, new_path) + working, overrides = _drop_empty_ancestors( + working, + overrides, + source_path[:-1], + ) + current_paths[plan.source_path] = new_path + _mark_rehome_override( + overrides, + new_path, + segment_index=plan.segment_index, + ) + events.append( + { + "action": "moved", + "source_path": list(source_path), + "dest_parent_path": list(dest_parent), + "anchor_path": list(anchor_path), + "new_path": list(new_path), + "physical_page": plan.physical_page, + "segment_index": plan.segment_index, + } + ) + logger.info( + "{} moved {} -> {} after {} page={} segment={}", + _LOG_PREFIX, + source_path, + new_path, + anchor_path, + plan.physical_page, + plan.segment_index, + ) + return working, overrides, events + + +def _detach_by_toc_index( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + *, + target_toc_index: int, + parent_path: tuple[str, ...] = (), + toc_index_start: int = 0, +) -> tuple[list[TitleNode], TitleNode | None, tuple[str, ...] | None]: + """Detach the paged leaf whose TOC-order index matches ``target_toc_index``.""" + toc_index = toc_index_start + result: list[TitleNode] = [] + detached: TitleNode | None = None + detached_path: tuple[str, ...] | None = None + for node in nodes: + if detached is not None: + result.append(node) + continue + path = (*parent_path, node.title) + if node.children: + new_children, child_detached, child_path = _detach_by_toc_index( + list(node.children), + overrides, + target_toc_index=target_toc_index, + parent_path=path, + toc_index_start=toc_index, + ) + child_leaves = _collect_paged_leaves( + list(node.children), + overrides, + path, + toc_index_start=toc_index, + ) + toc_index += len(child_leaves) + if child_detached is not None: + detached = child_detached + detached_path = child_path + result.append(replace(node, children=new_children)) + else: + result.append(node) + continue + page = own_page(path, overrides) + if page is None: + result.append(node) + continue + if toc_index == target_toc_index: + detached = node + detached_path = path + toc_index += 1 + continue + result.append(node) + toc_index += 1 + return result, detached, detached_path + + +def _mark_rehome_override( + overrides: dict[tuple[str, ...], TitleMatch], + path: tuple[str, ...], + *, + segment_index: int, +) -> None: + match = overrides.get(path) + if match is None: + raise ValueError(f"{_LOG_PREFIX} moved override missing path={path!r}") + overrides[path] = replace( + match, + evidence={ + **dict(match.evidence or {}), + "toc_rehome": {"segment_index": segment_index}, + }, + ) + + +def _detach_node( + nodes: list[TitleNode], + path: tuple[str, ...], +) -> tuple[list[TitleNode], TitleNode | None]: + if not path: + return nodes, None + title, *rest = path + if not rest: + detached: TitleNode | None = None + kept: list[TitleNode] = [] + for node in nodes: + if detached is None and node.title == title: + detached = node + continue + kept.append(node) + return kept, detached + + updated: list[TitleNode] = [] + detached_node: TitleNode | None = None + for node in nodes: + if node.title != title: + updated.append(node) + continue + new_children, detached_node = _detach_node(list(node.children), tuple(rest)) + updated.append(replace(node, children=new_children)) + return updated, detached_node + + +def _drop_empty_ancestors( + nodes: list[TitleNode], + overrides: dict[tuple[str, ...], TitleMatch], + ancestor_path: tuple[str, ...], +) -> tuple[list[TitleNode], dict[tuple[str, ...], TitleMatch]]: + """Delete parents that became empty after children moved out.""" + if not ancestor_path: + return nodes, overrides + working = nodes + for depth in range(len(ancestor_path), 0, -1): + path = ancestor_path[:depth] + node = node_at_path(working, path) + if node is None: + continue + if node.children: + break + working, removed = _detach_node(working, path) + if removed is not None: + overrides.pop(path, None) + logger.info("{} drop empty shell path={}", _LOG_PREFIX, path) + return working, overrides + + +def _move_override_prefix( + overrides: dict[tuple[str, ...], TitleMatch], + old_prefix: tuple[str, ...], + new_prefix: tuple[str, ...], +) -> None: + if old_prefix == new_prefix: + return + prefix_len = len(old_prefix) + moved: dict[tuple[str, ...], TitleMatch] = {} + for path in list(overrides.keys()): + if path[:prefix_len] != old_prefix: + continue + match = overrides.pop(path) + moved[new_prefix + path[prefix_len:]] = match + overrides.update(moved) + + +def _insert_under_parent( + *, + nodes: list[TitleNode], + parent_path: tuple[str, ...], + new_node: TitleNode, + start: int, + overrides: dict[tuple[str, ...], TitleMatch], +) -> list[TitleNode]: + if not parent_path: + return insert_by_start( + siblings=nodes, + new_node=new_node, + start=start, + parent_path=(), + primary_overrides=overrides, + ) + parent = node_at_path(nodes, parent_path) + if parent is None: + raise ValueError( + f"{_LOG_PREFIX} insert parent missing path={parent_path!r}" + ) + new_children = insert_by_start( + siblings=list(parent.children), + new_node=new_node, + start=start, + parent_path=parent_path, + primary_overrides=overrides, + ) + return replace_node_at_path( + nodes, + parent_path, + replace(parent, children=new_children), + ) diff --git a/apps/worker/app/services/document_agent/tools/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py index e04cd8152..856aaa5c3 100644 --- a/apps/worker/app/services/document_agent/tools/inspect_pages.py +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -90,23 +90,41 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: latency_ms=int((time.monotonic() - start) * 1000), ) - from app.services.document_agent.visual import render_pages - folder_name = str(args.get("folder_name") or "inspect_pages") prefix = str(args.get("prefix") or "inspect") - rendered = render_pages( - ctx, - pages, - folder_name=folder_name, - prefix=prefix, - timeout=120, - ) - if not rendered: - return ToolResult( - status="error", - error="render failed", - latency_ms=int((time.monotonic() - start) * 1000), + pre_rendered = args.get("rendered_pages") + if isinstance(pre_rendered, list) and pre_rendered: + rendered = [ + {"page": int(item["page"]), "png_path": str(item["png_path"])} + for item in pre_rendered + if isinstance(item, dict) + and item.get("page") is not None + and item.get("png_path") + and int(item["page"]) in pages + ] + rendered = sorted(rendered, key=lambda item: int(item["page"])) + if len(rendered) != len(pages): + return ToolResult( + status="error", + error="rendered_pages missing coverage for requested pages", + latency_ms=int((time.monotonic() - start) * 1000), + ) + else: + from app.services.document_agent.visual import render_pages + + rendered = render_pages( + ctx, + pages, + folder_name=folder_name, + prefix=prefix, + timeout=120, ) + if not rendered: + return ToolResult( + status="error", + error="render failed", + latency_ms=int((time.monotonic() - start) * 1000), + ) model = ctx.settings.get("vlm_model") or os.environ.get("IMAGE_MODEL") if not model: diff --git a/apps/worker/app/services/document_agent/visual.py b/apps/worker/app/services/document_agent/visual.py index e09f34aa7..a6feb0ee1 100644 --- a/apps/worker/app/services/document_agent/visual.py +++ b/apps/worker/app/services/document_agent/visual.py @@ -76,7 +76,13 @@ def _render_pages_worker( 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" + # Empty prefix → page_memory style ``page-{n}.png``. + # Non-empty keeps legacy ``{prefix}_page_{n}.png`` for PROFILE tools. + png_name = ( + f"page-{page_num}.png" + if not prefix + else 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}) diff --git a/apps/worker/app/services/document_parser/formats/docx/parser.py b/apps/worker/app/services/document_parser/formats/docx/parser.py index f978ab77d..00ef0772c 100755 --- a/apps/worker/app/services/document_parser/formats/docx/parser.py +++ b/apps/worker/app/services/document_parser/formats/docx/parser.py @@ -432,7 +432,12 @@ def handle_table( from shared.services.ai.summary.engine import summarize # Tables are Contract B assets: title + summary + entities from HTML. - result = summarize(mode="asset", text=tb_html_str, max_keywords=3) + result = summarize( + mode="asset", + text=tb_html_str, + max_keywords=3, + usage_task="parser.docx.table", + ) llm_title = result.title or None tb_keywords = result.keywords_str() llm_summary = result.summary or None diff --git a/apps/worker/app/services/document_parser/formats/excel/table_parser.py b/apps/worker/app/services/document_parser/formats/excel/table_parser.py index 4aa4c5e78..367a56dd5 100644 --- a/apps/worker/app/services/document_parser/formats/excel/table_parser.py +++ b/apps/worker/app/services/document_parser/formats/excel/table_parser.py @@ -325,6 +325,7 @@ def _summarize_excel_table( mode="asset", text=table_html, max_keywords=3, + usage_task="parser.excel.table", ) return ( result.title or None, diff --git a/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py index 0964825b0..5b3b0c3dd 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py +++ b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py @@ -42,7 +42,6 @@ QUEUE_POLL_INTERVAL_SECONDS = 0.1 CHILD_EXIT_GRACE_SECONDS = 5 POST_RESULT_EXIT_GRACE_SECONDS = 5 -POST_KILL_JOIN_GRACE_SECONDS = 1 # Child can exit 0 before multiprocessing.Queue's feeder flushes; retry once. EMPTY_QUEUE_EXIT_RETRIES = 1 PROCESS_POOL_SIZE = read_pymupdf_max_concurrent() @@ -122,14 +121,16 @@ def _shutdown_process_pool() -> None: def _close_result_queue(result_queue: MultiprocessingQueue) -> None: """Release parent-side queue resources once the child result is no longer needed.""" try: - result_queue.close() + # Avoid joining the queue feeder thread after we already have the payload; + # otherwise parent/child teardown can stall each other. + result_queue.cancel_join_thread() except Exception as exc: - logger.debug(f"Failed to close PyMuPDF result queue: {exc}") + logger.debug(f"Failed to cancel PyMuPDF result queue join: {exc}") try: - result_queue.join_thread() + result_queue.close() except Exception as exc: - logger.debug(f"Failed to join PyMuPDF result queue thread: {exc}") + logger.debug(f"Failed to close PyMuPDF result queue: {exc}") def _is_empty_queue_exit(exc: Exception) -> bool: @@ -242,7 +243,8 @@ def _run_worker_in_spawned_process_once( ), ) - _close_result_queue(result_queue) + # Success path: wait for natural child exit before tearing down the queue. + # Do not kill — the payload is already in hand; forced kill only adds noise. proc.join(timeout=POST_RESULT_EXIT_GRACE_SECONDS) elapsed = time.monotonic() - t0 exit_lag = None @@ -250,17 +252,14 @@ def _run_worker_in_spawned_process_once( exit_lag = elapsed - (wait_result.received_at - t0) if proc.is_alive(): - proc.kill() - proc.join(timeout=POST_KILL_JOIN_GRACE_SECONDS) - elapsed = time.monotonic() - t0 - if wait_result.received_at is not None: - exit_lag = elapsed - (wait_result.received_at - t0) - logger.warning( + logger.debug( f"[pymupdf-subprocess] EXIT_DELAY pid={child_pid} fn={worker_fn.__name__} " f"elapsed={elapsed:.1f}s exit_lag={exit_lag:.1f}s " - f"— result returned before child exited; child killed after grace" + f"— result returned; leaving child to finish exit without kill" ) + _close_result_queue(result_queue) + result = wait_result.result or {} if not result.get("ok"): logger.error( diff --git a/apps/worker/app/services/page_memory/fine_hierarchy.py b/apps/worker/app/services/page_memory/fine_hierarchy.py index 57db6362b..287e8d542 100644 --- a/apps/worker/app/services/page_memory/fine_hierarchy.py +++ b/apps/worker/app/services/page_memory/fine_hierarchy.py @@ -107,6 +107,7 @@ def refine_fat_leaf_skeletons( def compute_fat_leaf_pages( skeletons: list[SectionSkeleton], min_pages: int, + toc_pages: list[int] | None = None, ) -> set[int]: """Compute the set of page indices belonging to fat-leaf sections. @@ -116,13 +117,21 @@ def compute_fat_leaf_pages( Uses exclusive end boundaries when sibling starts are visible in ``skeletons``. For single-leaf scopes the closed ``end_page`` is used, which includes the shared boundary page with the next leaf. + + ``toc_pages`` are excluded from both the fat span test and the returned + page set (same body-page rule as scope processing). """ + excluded = {int(page) for page in (toc_pages or [])} fat_pages: set[int] = set() for idx, skel in enumerate(skeletons): exclusive_end = _exclusive_end(skeletons, idx) - page_span = exclusive_end - skel.start_page + 1 - if page_span > min_pages: - fat_pages.update(range(skel.start_page, exclusive_end + 1)) + body_pages = [ + page + for page in range(skel.start_page, exclusive_end + 1) + if page not in excluded + ] + if len(body_pages) > min_pages: + fat_pages.update(body_pages) return fat_pages diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index abef06c4b..dfc2ac196 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -264,6 +264,7 @@ def _build_page_dataframe( ) ] + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.fine_hierarchy import build_next_title_by_path next_title_by_path = build_next_title_by_path(skeletons) @@ -272,9 +273,14 @@ def _build_page_dataframe( filename=filename, page_count=page_count, ) - coarse_pages_scope = _derive_hierarchy_page_scope( - skeletons=skeletons, - page_count=page_count, + toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) + coarse_pages_scope = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ), + toc_pages, ) _record_trace_stage( trace_recorder, @@ -308,6 +314,11 @@ def _build_page_dataframe( len(coarse_scopes), scope_concurrency, ) + if toc_pages: + logger.info( + "[page_memory] excluding TOC pages from scope processing: {}", + sorted({int(page) for page in toc_pages}), + ) scope_results: list[_ScopeRunResult] = [] asset_pages_remaining = ( _resolve_asset_max_pages(page_count, page_memory_config) @@ -329,6 +340,7 @@ def _build_page_dataframe( trace_recorder=trace_recorder, page_memory_config=page_memory_config, next_title_by_path=next_title_by_path, + toc_pages=toc_pages, ) if scope_concurrency <= 1 or len(coarse_scopes) <= 1: @@ -409,9 +421,12 @@ def _build_page_dataframe( # Shared per-page lookups for node-granularity assembly. raw_text_by_page: dict[int, str] = {} image_path_by_page: dict[int, str] = {} - final_pages_scope = _derive_hierarchy_page_scope( - skeletons=skeletons, - page_count=page_count, + final_pages_scope = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=skeletons, + page_count=page_count, + ), + toc_pages, ) for page in final_pages_scope: rend = render_map.get(page) @@ -426,6 +441,7 @@ def _build_page_dataframe( skeletons=skeletons, raw_text_by_page=raw_text_by_page, image_path_by_page=image_path_by_page, + output_dir=output_dir, kind_by_page=label_map, tag_by_page=tag_map, filename=filename, @@ -538,7 +554,9 @@ def _run_hierarchy_scope( trace_recorder: Any | None, page_memory_config: PageMemoryConfig, next_title_by_path: dict[str, str | None] | None = None, + toc_pages: list[int] | None = None, ) -> _ScopeRunResult: + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.fine_hierarchy import ( compute_fat_leaf_pages, refine_fat_leaf_skeletons, @@ -559,9 +577,12 @@ def _run_hierarchy_scope( page_count=page_count, strategy=scope.strategy, ) - coarse_pages = _derive_hierarchy_page_scope( - skeletons=scope_skeletons, - page_count=page_count, + coarse_pages = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=scope_skeletons, + page_count=page_count, + ), + toc_pages, ) logger.info( "[page_memory] scope {}/{} {} coarse ranges={}", @@ -582,7 +603,11 @@ def _run_hierarchy_scope( ) fine_min = page_memory_config.fine_min_pages - fat_leaf_pages = compute_fat_leaf_pages(scope_skeletons, min_pages=fine_min) + fat_leaf_pages = compute_fat_leaf_pages( + scope_skeletons, + min_pages=fine_min, + toc_pages=toc_pages, + ) if fat_leaf_pages: title_pages = sorted(fat_leaf_pages) with stage_timer("page_memory.title_render", page_count=len(title_pages)): @@ -590,6 +615,7 @@ def _run_hierarchy_scope( pdf_path=pdf_path, page_count=page_count, output_dir=output_dir, + scope_id=scope.scope_id, pages=title_pages, page_features=page_features, page_texts=page_texts, @@ -656,9 +682,12 @@ def _run_hierarchy_scope( }, ) - final_pages = _derive_hierarchy_page_scope( - skeletons=scope_skeletons, - page_count=page_count, + final_pages = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=scope_skeletons, + page_count=page_count, + ), + toc_pages, ) final_scope_summary = _summarize_tag_scope( skeletons=scope_skeletons, @@ -670,6 +699,7 @@ def _run_hierarchy_scope( pdf_path=pdf_path, page_count=page_count, output_dir=output_dir, + scope_id=scope.scope_id, pages=final_pages, page_features=page_features, page_texts=page_texts, diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index 7273a0932..c4dc46257 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -382,6 +382,7 @@ def build_node_rows( skeletons: list[SectionSkeleton], raw_text_by_page: dict[int, str], image_path_by_page: dict[int, str], + output_dir: str, kind_by_page: dict[int, str], tag_by_page: dict[int, PageTagResult], filename: str, @@ -469,6 +470,7 @@ def _summarize_one( "extra_metadata": _build_page_extra_metadata( pages=view.pages, image_path_by_page=image_path_by_page, + output_dir=output_dir, ), } rows.append(row) @@ -498,10 +500,12 @@ def _build_page_extra_metadata( *, pages: list[int], image_path_by_page: dict[int, str], + output_dir: str, ) -> dict[str, Any]: page_assets = _build_page_citation_assets( pages=pages, image_path_by_page=image_path_by_page, + output_dir=output_dir, ) if not page_assets: return {} @@ -512,6 +516,7 @@ def _build_page_citation_assets( *, pages: list[int], image_path_by_page: dict[int, str], + output_dir: str, ) -> list[dict[str, Any]]: assets: list[dict[str, Any]] = [] seen_pages: set[int] = set() @@ -522,7 +527,11 @@ def _build_page_citation_assets( image_path = image_path_by_page.get(page) if not image_path or not os.path.exists(image_path): continue - artifact_ref = _promote_page_citation_asset(page=page, image_path=image_path) + artifact_ref = _promote_page_citation_asset( + page=page, + image_path=image_path, + output_dir=output_dir, + ) if not artifact_ref: continue width, height = _read_image_dimensions(image_path) @@ -540,10 +549,14 @@ def _build_page_citation_assets( return assets -def _promote_page_citation_asset(*, page: int, image_path: str) -> str: +def _promote_page_citation_asset( + *, + page: int, + image_path: str, + output_dir: str, +) -> str: source_path = Path(image_path) - output_dir = source_path.parent.parent - target_dir = output_dir / "page_citation_assets" + target_dir = Path(output_dir) / "page_citation_assets" target_path = target_dir / f"page-{page}.png" artifact_ref = f"page_citation_assets/page-{page}.png" try: diff --git a/apps/worker/app/services/page_memory/page_renderer.py b/apps/worker/app/services/page_memory/page_renderer.py index 944d7083a..925c15df8 100644 --- a/apps/worker/app/services/page_memory/page_renderer.py +++ b/apps/worker/app/services/page_memory/page_renderer.py @@ -23,7 +23,7 @@ class PageRenderResult: """1-based page number.""" image_path: str - """Absolute path to full-resolution PNG (``pages/page-N.png``).""" + """Absolute path to full-resolution PNG (``pages/{scope_id}/page-{n}.png``).""" raw_text: str """PROFILE scan text for this page.""" @@ -43,6 +43,7 @@ def render_document_pages( pdf_path: str, page_count: int, output_dir: str, + scope_id: str, pages: list[int] | None = None, page_features: list[PageFeature] | None = None, page_texts: dict[int, str] | None = None, @@ -61,7 +62,11 @@ def render_document_pages( pages: Optional 1-based page subset to render. Defaults to every page. output_dir: - Root output directory; pages are written to ``output_dir/pages/``. + Root output directory; pages are written to + ``output_dir/pages/{scope_id}/page-{n}.png``. + scope_id: + Scope isolation key (e.g. ``p58-60``). Concurrent scopes must not + share a write path for the same page index. page_features: If available, dimensions are read from here (avoiding a second PyMuPDF open). Otherwise falls back to 0/0/False. @@ -81,6 +86,10 @@ def render_document_pages( list[PageRenderResult] One entry per page, ordered by page_index. """ + resolved_scope = str(scope_id or "").strip() + if not resolved_scope: + raise ValueError("render_document_pages requires a non-empty scope_id") + requested_pages = ( sorted({page for page in pages if 1 <= page <= page_count}) if pages is not None @@ -90,14 +99,16 @@ def render_document_pages( return [] texts = page_texts or {} + folder_name = f"pages/{resolved_scope}" # ── full-resolution PNGs ────────────────────────────────────────── + # Empty prefix → ``page-{n}.png`` (see visual._render_pages_worker). if ctx is not None: pngs = render_pages( ctx, requested_pages, - folder_name="pages", - prefix="page", + folder_name=folder_name, + prefix="", dpi=dpi, timeout=timeout, ) @@ -117,8 +128,8 @@ def render_document_pages( pngs = render_pages( tmp_ctx, requested_pages, - folder_name="pages", - prefix="page", + folder_name=folder_name, + prefix="", dpi=dpi, timeout=timeout, ) diff --git a/apps/worker/scripts/_debug_token_ledger.py b/apps/worker/scripts/_debug_token_ledger.py new file mode 100644 index 000000000..c0d81308e --- /dev/null +++ b/apps/worker/scripts/_debug_token_ledger.py @@ -0,0 +1,128 @@ +"""Shared token-ledger helpers for debug parse scripts.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from pathlib import Path +from typing import Any + +_NUMERIC_USAGE_FIELDS = ("prompt_tokens", "completion_tokens", "total_tokens", "calls") + + +def empty_token_usage() -> dict[str, Any]: + return { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "calls": 0, + "by_model": {}, + "by_task": {}, + } + + +def token_usage_delta(prev: dict[str, Any], cur: dict[str, Any]) -> dict[str, Any]: + numeric = _NUMERIC_USAGE_FIELDS + + def _sub(left: dict[str, Any], right: dict[str, Any]) -> dict[str, int]: + return { + field: int(right.get(field, 0)) - int(left.get(field, 0)) for field in numeric + } + + def _bucket(prev_bucket: dict[str, Any], cur_bucket: dict[str, Any]) -> dict[str, Any]: + merged: dict[str, Any] = {} + for key in set(prev_bucket) | set(cur_bucket): + prev_item = prev_bucket.get(key, {}) + cur_item = cur_bucket.get(key, {}) + if not isinstance(prev_item, dict) or not isinstance(cur_item, dict): + continue + entry = { + field: value for field, value in _sub(prev_item, cur_item).items() if value + } + prev_models = prev_item.get("models", {}) + cur_models = cur_item.get("models", {}) + if prev_models or cur_models: + models = _bucket(prev_models, cur_models) + if models: + entry["models"] = models + if entry: + merged[key] = entry + return merged + + delta = _sub(prev, cur) + for bucket_key in ("by_model", "by_task"): + bucket = _bucket(prev.get(bucket_key, {}), cur.get(bucket_key, {})) + if bucket: + delta[bucket_key] = bucket + return delta + + +def merge_token_usage(destination: dict[str, Any], source: dict[str, Any]) -> None: + for key, value in source.items(): + if isinstance(value, dict): + child = destination.setdefault(str(key), {}) + if isinstance(child, dict): + merge_token_usage(child, value) + elif isinstance(value, int | float) and not isinstance(value, bool): + destination[str(key)] = destination.get(str(key), 0) + value + + +def load_stage_ledger(path: Path, *, version: str = "1.0") -> dict[str, Any]: + if not path.exists(): + return {"version": version, "stages": {}} + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, dict): + return {"version": version, "stages": {}} + data.setdefault("version", version) + data.setdefault("stages", {}) + return data + + +def write_stage_ledger(path: Path, ledger: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temp_path = path.with_suffix(f"{path.suffix}.tmp") + temp_path.write_text(json.dumps(ledger, ensure_ascii=False, indent=2), encoding="utf-8") + temp_path.replace(path) + + +def record_stage_delta( + ledger: dict[str, Any], + *, + stage: str, + stage_keys: tuple[str, ...], + prev: dict[str, Any], + current: dict[str, Any], + out_path: Path, +) -> dict[str, Any]: + """Store per-stage delta and drop stale downstream stage keys.""" + stages = ledger.setdefault("stages", {}) + if not isinstance(stages, dict): + stages = {} + ledger["stages"] = stages + stage_index = stage_keys.index(stage) + for stale in stage_keys[stage_index + 1 :]: + stages.pop(stale, None) + stages[stage] = {"token_usage": token_usage_delta(prev, current)} + write_stage_ledger(out_path, ledger) + return deepcopy(current) + + +def aggregate_stage_deltas( + ledger: dict[str, Any], + stage_keys: tuple[str, ...], + *, + remainder: dict[str, Any] | None = None, +) -> dict[str, Any]: + usage = empty_token_usage() + stages = ledger.get("stages") + if isinstance(stages, dict): + for stage in stage_keys: + row = stages.get(stage) + if not isinstance(row, dict): + continue + raw_usage = row.get("token_usage") + if isinstance(raw_usage, dict): + merge_token_usage(usage, raw_usage) + if remainder: + merge_token_usage(usage, remainder) + return usage diff --git a/apps/worker/scripts/debug_text_track.py b/apps/worker/scripts/debug_text_track.py index d80c15bd9..e14176a1f 100644 --- a/apps/worker/scripts/debug_text_track.py +++ b/apps/worker/scripts/debug_text_track.py @@ -29,6 +29,7 @@ import shutil import sys import time +from copy import deepcopy from pathlib import Path from typing import Any @@ -36,6 +37,7 @@ ROOT = Path(__file__).resolve().parents[3] WORKER_ROOT = ROOT / "apps" / "worker" sys.path.insert(0, str(WORKER_ROOT)) +sys.path.insert(0, str(WORKER_ROOT / "scripts")) sys.path.insert(0, str(ROOT / "packages" / "shared-python")) from dotenv import load_dotenv @@ -47,8 +49,17 @@ from loguru import logger from shared.services.ai.token_tracking import ( - init_token_tracker, + cleanup_token_tracker, get_current_token_tracker, + init_token_tracker, +) + +from _debug_token_ledger import ( + aggregate_stage_deltas, + empty_token_usage, + load_stage_ledger, + record_stage_delta, + token_usage_delta, ) # ── Constants ─────────────────────────────────────────────────────────────── @@ -58,6 +69,8 @@ ) DEFAULT_SPACEX_PDF = Path("/Users/wuchengke/Desktop/temp/test_docs/spacex-s1.pdf") OUTPUT_ROOT = Path("~/.knowhere/_debug_parse").expanduser() +TOKEN_LEDGER_NAME = "token_ledger.json" +TOKEN_LEDGER_STAGES = ("profile", "mineru", "hierarchy", "full") def _write_json(path: Path, data: Any) -> None: @@ -66,6 +79,61 @@ def _write_json(path: Path, data: Any) -> None: logger.info(" → {}", path) +def _load_token_ledger(out_dir: Path) -> dict[str, Any]: + return load_stage_ledger(out_dir / TOKEN_LEDGER_NAME) + + +def _record_token_stage( + ledger: dict[str, Any], + stage: str, + *, + prev: dict[str, Any], + out_dir: Path, +) -> dict[str, Any]: + current = deepcopy(get_current_token_tracker() or {}) + record_stage_delta( + ledger, + stage=stage, + stage_keys=TOKEN_LEDGER_STAGES, + prev=prev, + current=current, + out_path=out_dir / TOKEN_LEDGER_NAME, + ) + logger.info(" → {}", out_dir / TOKEN_LEDGER_NAME) + return current + + +def _aggregate_token_ledger( + ledger: dict[str, Any], + *, + remainder: dict[str, Any] | None = None, +) -> dict[str, Any]: + return aggregate_stage_deltas( + ledger, + TOKEN_LEDGER_STAGES, + remainder=remainder, + ) + + +def _apply_token_usage_to_outputs( + out_dir: Path, + trace: dict[str, Any], + usage: dict[str, Any], +) -> None: + trace["token_usage"] = usage + _write_json(out_dir / "trace.json", trace) + manifest_path = out_dir / "manifest.json" + if not manifest_path.exists(): + return + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + if not isinstance(manifest, dict): + return + processing = manifest.setdefault("processing", {}) + if isinstance(processing, dict): + processing["token_usage"] = usage + _write_json(manifest_path, manifest) + + # ── Stage 1: Profile + Shard Plan (PDF only) ─────────────────────────────── def _stage_profile(pdf_path: str, filename: str, out_dir: Path, model: str | None): @@ -716,6 +784,8 @@ def main() -> int: out_dir.mkdir(parents=True, exist_ok=True) init_token_tracker() + ledger = _load_token_ledger(out_dir) + prev_usage = deepcopy(get_current_token_tracker() or empty_token_usage()) logger.info("█" * 70) logger.info(" TEXT-TRACK DEBUG: {}", filename) @@ -731,158 +801,188 @@ def main() -> int: } t_start = time.time() - # ── Format dispatch ────────────────────────────────────────────────────── - if ext == ".pdf": - # Stage 1: Profile - if args.reuse_profile: - anatomy = _load_anatomy_cache(out_dir, file_path, filename) - else: - anatomy, profile_elapsed, profile_meta = _stage_profile( - file_path, filename, out_dir, args.model + try: + # ── Format dispatch ────────────────────────────────────────────────── + if ext == ".pdf": + # Stage 1: Profile + if args.reuse_profile: + anatomy = _load_anatomy_cache(out_dir, file_path, filename) + else: + anatomy, profile_elapsed, profile_meta = _stage_profile( + file_path, filename, out_dir, args.model + ) + trace["stages"]["profile"] = { + "elapsed_s": round(profile_elapsed, 1), + **profile_meta, + } + prev_usage = _record_token_stage( + ledger, "profile", prev=prev_usage, out_dir=out_dir + ) + + if args.stop_at == "profile": + trace["stages"].setdefault("profile", {})["shard_count"] = len( + anatomy.shard_plan.shards + ) + logger.info("⏸️ Stopped at profile → {}", out_dir) + return 0 + + # Stage 2: MinerU extraction + if not args.reuse_mineru: + _shard_dirs, mineru_elapsed = _stage_mineru_pdf( + file_path, filename, out_dir, anatomy, + ) + trace["stages"]["mineru"] = { + "elapsed_s": round(mineru_elapsed, 1), + "shard_count": len(_shard_dirs), + } + prev_usage = _record_token_stage( + ledger, "mineru", prev=prev_usage, out_dir=out_dir + ) + else: + logger.info("⏩ Reusing cached MinerU shard dirs") + + if args.stop_at == "mineru": + logger.info("⏸️ Stopped at mineru → {}", out_dir) + return 0 + + # Stage 3: Heading prediction + hierarchy tree + if args.reuse_hierarchy: + merged_path = out_dir / "_shards" / "merged_lines.json" + if not merged_path.exists(): + raise FileNotFoundError(f"No cached hierarchy: {merged_path}") + logger.info("⏩ Reusing cached merged_lines: {}", merged_path) + merged_lines = json.loads(merged_path.read_text(encoding="utf-8")) + else: + merged_lines, hier_elapsed = _stage_hierarchy_pdf( + out_dir, anatomy, args.model, + ) + trace["stages"]["hierarchy"] = { + "elapsed_s": round(hier_elapsed, 1), + "merged_lines_count": len(merged_lines), + "heading_count": sum(1 for ln in merged_lines if ln.startswith("#")), + } + prev_usage = _record_token_stage( + ledger, "hierarchy", prev=prev_usage, out_dir=out_dir + ) + + if args.stop_at == "hierarchy": + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 + + # Stage 4: Full extraction + chunks, full_elapsed = _stage_full_pdf( + out_dir, filename, merged_lines, args.model ) - trace["stages"]["profile"] = { - "elapsed_s": round(profile_elapsed, 1), - **profile_meta, + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), } + prev_usage = _record_token_stage( + ledger, "full", prev=prev_usage, out_dir=out_dir + ) + + elif ext in (".docx", ".doc"): + if args.stop_at in ("profile", "mineru"): + logger.info("ℹ️ No profiling/MinerU for DOCX format. Nothing to do.") + return 0 - if args.stop_at == "profile": - trace["stages"].setdefault("profile", {})["shard_count"] = len( - anatomy.shard_plan.shards + parsed_df, hier_elapsed = _stage_hierarchy_docx( + file_path, filename, out_dir, args.model + ) + trace["stages"]["hierarchy"] = { + "elapsed_s": round(hier_elapsed, 1), + "row_count": len(parsed_df) if parsed_df is not None else 0, + } + prev_usage = _record_token_stage( + ledger, "hierarchy", prev=prev_usage, out_dir=out_dir ) - trace["token_usage"] = get_current_token_tracker() - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at profile → {}", out_dir) - return 0 - - # Stage 2: MinerU extraction - if not args.reuse_mineru: - _shard_dirs, mineru_elapsed = _stage_mineru_pdf( - file_path, filename, out_dir, anatomy, + + if args.stop_at == "hierarchy": + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 + + from app.services.document_parser.orchestration.postprocess import ( + apply_parse_postprocess, ) - trace["stages"]["mineru"] = { - "elapsed_s": round(mineru_elapsed, 1), - "shard_count": len(_shard_dirs), + parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) + chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), } - else: - logger.info("⏩ Reusing cached MinerU shard dirs") - - if args.stop_at == "mineru": - trace["token_usage"] = get_current_token_tracker() - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at mineru → {}", out_dir) - return 0 - - # Stage 3: Heading prediction + hierarchy tree - if args.reuse_hierarchy: - merged_path = out_dir / "_shards" / "merged_lines.json" - if not merged_path.exists(): - raise FileNotFoundError(f"No cached hierarchy: {merged_path}") - logger.info("⏩ Reusing cached merged_lines: {}", merged_path) - merged_lines = json.loads(merged_path.read_text(encoding="utf-8")) - else: - merged_lines, hier_elapsed = _stage_hierarchy_pdf( - out_dir, anatomy, args.model, + prev_usage = _record_token_stage( + ledger, "full", prev=prev_usage, out_dir=out_dir + ) + + elif ext in (".md", ".markdown"): + if args.stop_at in ("profile", "mineru"): + logger.info("ℹ️ No profiling/MinerU for Markdown format. Nothing to do.") + return 0 + + parsed_df, hier_elapsed = _stage_hierarchy_md( + file_path, filename, out_dir, args.model ) trace["stages"]["hierarchy"] = { "elapsed_s": round(hier_elapsed, 1), - "merged_lines_count": len(merged_lines), - "heading_count": sum(1 for ln in merged_lines if ln.startswith("#")), + "row_count": len(parsed_df) if parsed_df is not None else 0, } + prev_usage = _record_token_stage( + ledger, "hierarchy", prev=prev_usage, out_dir=out_dir + ) - if args.stop_at == "hierarchy": - trace["token_usage"] = get_current_token_tracker() - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at hierarchy → {}", out_dir) - return 0 - - # Stage 4: Full extraction - chunks, full_elapsed = _stage_full_pdf(out_dir, filename, merged_lines, args.model) - trace["stages"]["full"] = { - "elapsed_s": round(full_elapsed, 1), - "chunk_count": len(chunks), - } - - elif ext in (".docx", ".doc"): - if args.stop_at in ("profile", "mineru"): - logger.info("ℹ️ No profiling/MinerU for DOCX format. Nothing to do.") - return 0 - - # Stage 2: parse_docx - parsed_df, hier_elapsed = _stage_hierarchy_docx(file_path, filename, out_dir, args.model) - trace["stages"]["hierarchy"] = { - "elapsed_s": round(hier_elapsed, 1), - "row_count": len(parsed_df) if parsed_df is not None else 0, - } - - if args.stop_at == "hierarchy": - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at hierarchy → {}", out_dir) - return 0 - - # Stage 4: DataFrame → chunks - from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess - parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) - chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) - trace["stages"]["full"] = { - "elapsed_s": round(full_elapsed, 1), - "chunk_count": len(chunks), - } - - elif ext in (".md", ".markdown"): - if args.stop_at in ("profile", "mineru"): - logger.info("ℹ️ No profiling/MinerU for Markdown format. Nothing to do.") - return 0 - - # Stage 2: parse_md - parsed_df, hier_elapsed = _stage_hierarchy_md(file_path, filename, out_dir, args.model) - trace["stages"]["hierarchy"] = { - "elapsed_s": round(hier_elapsed, 1), - "row_count": len(parsed_df) if parsed_df is not None else 0, - } - - if args.stop_at == "hierarchy": - _write_json(out_dir / "trace.json", trace) - logger.info("⏸️ Stopped at hierarchy → {}", out_dir) - return 0 - - # Stage 4: DataFrame → chunks - from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess - parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) - chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) - trace["stages"]["full"] = { - "elapsed_s": round(full_elapsed, 1), - "chunk_count": len(chunks), - } + if args.stop_at == "hierarchy": + logger.info("⏸️ Stopped at hierarchy → {}", out_dir) + return 0 - else: - logger.error("Unsupported format: {}", ext) - return 1 - - # ── Optional DB publication ────────────────────────────────────────────── - if args.run_db: - from scripts._debug_publish import publish_debug_result_dir - publish_result = publish_debug_result_dir( - result_dir=out_dir, - source_file_name=filename, - chunks=chunks, - parse_track="text_track", - upload_assets=True, - ) - trace["stages"]["db_publish"] = { - "job_id": publish_result.job_id, - "document_id": publish_result.document_id, - } - - # ── Final trace ────────────────────────────────────────────────────────── - trace["total_elapsed_s"] = round(time.time() - t_start, 1) - trace["token_usage"] = get_current_token_tracker() - _write_json(out_dir / "trace.json", trace) + from app.services.document_parser.orchestration.postprocess import ( + apply_parse_postprocess, + ) + parsed_df = apply_parse_postprocess(str(out_dir), parsed_df) + chunks, full_elapsed = _finalize_df(out_dir, filename, parsed_df) + trace["stages"]["full"] = { + "elapsed_s": round(full_elapsed, 1), + "chunk_count": len(chunks), + } + prev_usage = _record_token_stage( + ledger, "full", prev=prev_usage, out_dir=out_dir + ) - logger.info("") - logger.info("═" * 70) - logger.info(" ✅ DONE in {:.1f}s → {}", time.time() - t_start, out_dir) - logger.info("═" * 70) - return 0 + else: + logger.error("Unsupported format: {}", ext) + return 1 + + # ── Optional DB publication ────────────────────────────────────────── + if args.run_db: + from scripts._debug_publish import publish_debug_result_dir + publish_result = publish_debug_result_dir( + result_dir=out_dir, + source_file_name=filename, + chunks=chunks, + parse_track="text_track", + upload_assets=True, + ) + trace["stages"]["db_publish"] = { + "job_id": publish_result.job_id, + "document_id": publish_result.document_id, + } + + trace["total_elapsed_s"] = round(time.time() - t_start, 1) + logger.info("") + logger.info("═" * 70) + logger.info(" ✅ DONE in {:.1f}s → {}", time.time() - t_start, out_dir) + logger.info("═" * 70) + return 0 + finally: + remainder = token_usage_delta( + prev_usage, + deepcopy(get_current_token_tracker() or {}), + ) + usage = _aggregate_token_ledger( + ledger, + remainder=remainder if remainder.get("calls") or remainder.get("total_tokens") else None, + ) + _apply_token_usage_to_outputs(out_dir, trace, usage) + cleanup_token_tracker() if __name__ == "__main__": diff --git a/apps/worker/scripts/page_memory/_debug_pm_shared.py b/apps/worker/scripts/page_memory/_debug_pm_shared.py index be93e6030..43c913d45 100644 --- a/apps/worker/scripts/page_memory/_debug_pm_shared.py +++ b/apps/worker/scripts/page_memory/_debug_pm_shared.py @@ -26,6 +26,7 @@ WORKER_ROOT = ROOT / "apps" / "worker" sys.path.insert(0, str(WORKER_ROOT)) sys.path.insert(0, str(ROOT / "packages" / "shared-python")) +sys.path.insert(0, str(WORKER_ROOT / "scripts")) from dotenv import load_dotenv @@ -50,10 +51,12 @@ serialize_scope_skeletons as _serialize_scope_skeletons, ) from shared.services.ai.token_tracking import ( + cleanup_token_tracker, init_token_tracker, get_current_token_tracker, ) from shared.services.ai.token_costing import build_token_cost_estimate +from _debug_token_ledger import empty_token_usage, merge_token_usage, token_usage_delta, write_stage_ledger # Re-exported for staged debug scripts. Listing them here marks the imports as # intentional so CodeQL does not treat them as unused. @@ -76,62 +79,17 @@ # ── Token cost tracker ──────────────────────────────────────────────────────── -def _usage_delta(prev: dict[str, Any], cur: dict[str, Any]) -> dict[str, Any]: - _NUM = ("prompt_tokens", "completion_tokens", "total_tokens", "calls") - - def _sub(a: dict, b: dict) -> dict: - return {f: int(b.get(f, 0)) - int(a.get(f, 0)) for f in _NUM} - - def _bucket(pa: dict, pb: dict) -> dict: - r: dict[str, Any] = {} - for k in set(pa) | set(pb): - pk, ck = pa.get(k, {}), pb.get(k, {}) - if not isinstance(pk, dict) or not isinstance(ck, dict): - continue - e = {f: v for f, v in _sub(pk, ck).items() if v} - pm, cm = pk.get("models", {}), ck.get("models", {}) - if pm or cm: - md = _bucket(pm, cm) - if md: - e["models"] = md - if e: - r[k] = e - return r - - d = _sub(prev, cur) - for bk in ("by_model", "by_task"): - bd = _bucket(prev.get(bk, {}), cur.get(bk, {})) - if bd: - d[bk] = bd - return d - - class TokenCostTracker: """Incremental token usage & cost tracker for debug pipeline stages.""" def __init__(self) -> None: self._dict = init_token_tracker() - self._root_gid = self._gid() self._prev: dict[str, Any] = deepcopy(self._dict) self._stages: list[dict[str, Any]] = [] - @staticmethod - def _gid() -> int: - from shared.services.ai.token_tracking import _current_greenlet_id - - return _current_greenlet_id() - - def register_child_thread(self) -> None: - from shared.services.ai.token_tracking import _root_ids, _lock - - gid = self._gid() - if gid != self._root_gid: - with _lock: - _root_ids[gid] = self._root_gid - def snapshot_stage(self, stage: str) -> None: cur = deepcopy(get_current_token_tracker() or {}) - delta = _usage_delta(self._prev, cur) + delta = token_usage_delta(self._prev, cur) self._stages.append({ "stage": stage, "prompt_tokens": delta.get("prompt_tokens", 0), @@ -953,20 +911,6 @@ def _stage_usage_snapshot(tracker: TokenCostTracker | None) -> dict[str, Any]: } -def _merge_token_usage( - destination: dict[str, Any], - source: dict[str, Any], -) -> None: - """Merge raw production token-tracker snapshots recursively.""" - for key, value in source.items(): - if isinstance(value, dict): - child = destination.setdefault(str(key), {}) - if isinstance(child, dict): - _merge_token_usage(child, value) - elif isinstance(value, int | float) and not isinstance(value, bool): - destination[str(key)] = destination.get(str(key), 0) + value - - def record_stage_cost( out_dir: Path, *, @@ -1008,13 +952,7 @@ def record_stage_cost( ledger["updated_at"] = updated_at path = stage_costs_path(out_dir) - path.parent.mkdir(parents=True, exist_ok=True) - temp_path = path.with_suffix(f"{path.suffix}.tmp") - temp_path.write_text( - json.dumps(ledger, ensure_ascii=False, indent=2), - encoding="utf-8", - ) - temp_path.replace(path) + write_stage_ledger(path, ledger) logger.info( " stage cost → {} ({} {:.1f}s / ${:.6f})", path, @@ -1034,14 +972,7 @@ def aggregate_stage_costs(ledger: dict[str, Any]) -> dict[str, Any]: by_pipeline_stage: dict[str, Any] = {} by_substage: list[dict[str, Any]] = [] merged_trace_stages: list[dict[str, Any]] = [] - usage = { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0, - "calls": 0, - "by_model": {}, - "by_task": {}, - } + usage = empty_token_usage() elapsed_s = 0.0 for stage_key in _COST_STAGE_KEYS: @@ -1060,7 +991,7 @@ def aggregate_stage_costs(ledger: dict[str, Any]) -> dict[str, Any]: raw_usage = row.get("token_usage") if isinstance(raw_usage, dict): # Numeric top-level fields were already added above. - _merge_token_usage( + merge_token_usage( usage, { key: value @@ -1200,38 +1131,41 @@ def stop_with_trace( extra_summary: dict[str, Any] | None = None, ) -> int: """Write TRACE.JSON, optionally recording this pipeline stage's cost ledger.""" - if pipeline_stage is not None and elapsed_s is not None: - record_stage_cost( - out_dir, - pipeline_stage=pipeline_stage, - elapsed_s=elapsed_s, - token_cost_tracker=token_cost_tracker, - trace_stages=stages, - stop_at=stop_at, - ) + try: + if pipeline_stage is not None and elapsed_s is not None: + record_stage_cost( + out_dir, + pipeline_stage=pipeline_stage, + elapsed_s=elapsed_s, + token_cost_tracker=token_cost_tracker, + trace_stages=stages, + stop_at=stop_at, + ) - aggregated = aggregate_stage_costs(load_stage_costs(out_dir)) - merged_stages = aggregated.get("stages") or stages - summary: dict[str, Any] = { - "page_count": page_count, - "scope_id": scope_id, - "rows_count": None, - "elapsed_s": aggregated.get("elapsed_s"), - "completed_pipeline_stages": aggregated.get("completed_pipeline_stages"), - "token_cost": aggregated.get("token_cost"), - } - if extra_summary: - summary.update(jsonable(extra_summary)) - - write_trace( - out_dir=out_dir, - stages=merged_stages, - final_status=final_status or f"stopped_at_{stop_at}", - summary=summary, - ) - remove_nested_doc_agent_trace(out_dir) - maybe_purge_debug_visuals(out_dir) - return 0 + aggregated = aggregate_stage_costs(load_stage_costs(out_dir)) + merged_stages = aggregated.get("stages") or stages + summary: dict[str, Any] = { + "page_count": page_count, + "scope_id": scope_id, + "rows_count": None, + "elapsed_s": aggregated.get("elapsed_s"), + "completed_pipeline_stages": aggregated.get("completed_pipeline_stages"), + "token_cost": aggregated.get("token_cost"), + } + if extra_summary: + summary.update(jsonable(extra_summary)) + + write_trace( + out_dir=out_dir, + stages=merged_stages, + final_status=final_status or f"stopped_at_{stop_at}", + summary=summary, + ) + remove_nested_doc_agent_trace(out_dir) + maybe_purge_debug_visuals(out_dir) + return 0 + finally: + cleanup_token_tracker() def remove_nested_doc_agent_trace(out_dir: Path) -> None: @@ -1509,10 +1443,17 @@ def build_debug_coarse_scopes( page_count: int, anatomy: Any | None = None, ) -> list[dict[str, Any]]: + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory._utils import build_hierarchy_scopes - from toc_page_policy import TocPagePolicy - policy = TocPagePolicy.from_anatomy(anatomy) + toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) + toc_page_set = set() + for raw in toc_pages: + try: + toc_page_set.add(int(raw)) + except (TypeError, ValueError): + continue scopes = build_hierarchy_scopes( skeletons=skeletons, filename=filename, @@ -1525,12 +1466,13 @@ def build_debug_coarse_scopes( "start_page": scope.start_page, "end_page": scope.end_page, "strategy": scope.strategy, - "processing_pages": policy.filter_processing_pages( - list(range(scope.start_page, scope.end_page + 1)) + "processing_pages": pages_excluding_toc( + list(range(scope.start_page, scope.end_page + 1)), + toc_pages, ), "excluded_toc_pages": sorted( page - for page in policy.pure_toc_pages + for page in toc_page_set if scope.start_page <= page <= scope.end_page ), } diff --git a/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py b/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py index 81a2e7e9d..bc885c352 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage3_scope_fine_hierarchy.py @@ -1,10 +1,9 @@ #!/usr/bin/env python3 # ruff: noqa: E402 -"""Stage 3: Coarse scopes + document page tagging + per-scope fine hierarchy. +"""Stage 3: Coarse scopes + per-scope title detection / fine hierarchy / page tag. -Builds coarse hierarchy scopes, writes ``scopes//skeletons.json``, renders -and tags each selected processing page once (global concurrency), then fans -tag subsets into scopes for fine hierarchy refinement. +Mirrors production ``_run_hierarchy_scope`` (without C5 assets; that is Stage 4): +fat-leaf title render → ``tag_page_titles`` → fine hierarchy → render → ``tag_pages``. Requires Stage 2 output: _doc_agent/pipeline_state.json (with skeletons), doc_profile.json (after production ``run_toc_anchoring``). @@ -65,28 +64,29 @@ def _resolve_scope_processing_pages( scope_meta: dict[str, Any], skeletons: list[Any], page_count: int, - toc_policy: Any, + toc_pages: list[int], ) -> tuple[list[int], list[int]]: + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc + processing_pages = [ int(page) for page in (scope_meta.get("processing_pages") or []) ] if not processing_pages: - processing_pages = toc_policy.filter_processing_pages( + processing_pages = pages_excluding_toc( _derive_hierarchy_page_scope( skeletons=skeletons, page_count=page_count, - ) + ), + toc_pages, ) excluded_toc_pages = [ int(page) for page in (scope_meta.get("excluded_toc_pages") or []) ] if not excluded_toc_pages: + start = int(scope_meta.get("start_page") or 1) + end = int(scope_meta.get("end_page") or page_count) excluded_toc_pages = sorted( - set(range( - int(scope_meta.get("start_page") or 1), - int(scope_meta.get("end_page") or page_count) + 1, - )) - & toc_policy.pure_toc_pages + page for page in toc_pages if start <= int(page) <= end ) return processing_pages, excluded_toc_pages @@ -96,34 +96,48 @@ def _run_fine_hierarchy_for_scope( scope_id: str, scope_dir: Path, out_dir: Path, + pdf_path: str, page_count: int, - rendered_by_page: dict[int, Any], - tags_by_page: dict[int, Any], + page_texts: dict[int, str], + page_features: list[Any], + page_labels: list[Any], + vlm_model: str | None, next_title_by_path: dict[str, str | None], - toc_policy: Any, + toc_pages: list[int], page_memory_config: Any, token_cost_tracker: TokenCostTracker | None = None, ) -> ScopeResult: - """Consume shared page tags and refine one coarse scope.""" + """Per-scope path aligned with production ``_run_hierarchy_scope`` (no C5).""" + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.fine_hierarchy import ( compute_fat_leaf_pages, refine_fat_leaf_skeletons, ) - from app.services.page_memory.memory_service import _resolve_hierarchy_model + from app.services.page_memory.memory_service import ( + _resolve_hierarchy_model, + _summarize_tag_scope, + _summarize_tags, + ) + from app.services.page_memory.page_plan import derive_page_processing_plan + from app.services.page_memory.page_renderer import render_document_pages + from app.services.page_memory.page_tagger import ( + PageTagResult, + tag_page_titles, + tag_pages, + ) scope_stages: list[dict[str, Any]] = [] - if token_cost_tracker is not None: - token_cost_tracker.register_child_thread() skel_path = scope_dir / "skeletons.json" require_file(skel_path, hint=f"Stage 3 should have created {skel_path}") scope_meta, active_skeletons = load_scope_skeletons_artifact(skel_path) + active_skeletons = sort_skeletons(active_skeletons) strategy = str(scope_meta.get("strategy") or "coarse_scope") processing_pages, excluded_toc_pages = _resolve_scope_processing_pages( scope_meta=scope_meta, skeletons=active_skeletons, page_count=page_count, - toc_policy=toc_policy, + toc_pages=toc_pages, ) scope_manifest = _scope_manifest( @@ -131,53 +145,66 @@ def _run_fine_hierarchy_for_scope( skeletons=active_skeletons, page_count=page_count, strategy=strategy, - processing_pages=processing_pages, - excluded_toc_pages=excluded_toc_pages, ) logger.info( - "🔬 [scope {}] {} skeletons p{}-{} processing={}", + "🔬 [scope {}] {} skeletons p{}-{} processing={} excluded_toc={}", scope_id, len(active_skeletons), scope_meta.get("start_page", "?"), scope_meta.get("end_page", "?"), processing_pages, + excluded_toc_pages, ) - if not processing_pages: - logger.info(" [scope {}] no processing pages after TOC exclusion", scope_id) - return ScopeResult( - scope_id=scope_id, - skeletons=active_skeletons, - tags=[], - assets_by_page={}, - rendered=[], - final_pages=[], - scope_manifest=scope_manifest, - trace_stages=scope_stages, - ) - - rendered = [ - rendered_by_page[page] - for page in processing_pages - if page in rendered_by_page - ] - tags = [ - tags_by_page[page] - for page in processing_pages - if page in tags_by_page - ] - fine_min = page_memory_config.fine_min_pages fat_leaf_pages = compute_fat_leaf_pages( active_skeletons, min_pages=fine_min, - exclude_pages=toc_policy.pure_toc_pages, + toc_pages=toc_pages, ) if fat_leaf_pages: + title_pages = sorted(fat_leaf_pages) + title_rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=str(out_dir), + scope_id=scope_id, + pages=title_pages, + page_features=page_features, + page_texts=page_texts, + ) + title_tags = [ + PageTagResult( + page_index=page, + summary="", + keywords=[], + strategy_used="title_detection_only", + ) + for page in title_pages + ] + title_tags = tag_page_titles( + pages=title_rendered, + tag_results=title_tags, + fat_leaf_pages=fat_leaf_pages, + vlm_model=vlm_model, + scan_direction=page_memory_config.scan_direction, + max_concurrent=page_memory_config.title_detection_concurrency, + ) + record_stage( + scope_stages, + "C3b.title_detection", + page_info=page_scope_info(title_pages), + variables={ + "scope_id": scope_id, + "tags": _summarize_tags(title_tags), + }, + ) + if token_cost_tracker is not None: + token_cost_tracker.snapshot_stage(f"C3b.title_detection:{scope_id}") active_skeletons = refine_fat_leaf_skeletons( coarse_skeletons=active_skeletons, - tag_results=tags, + tag_results=title_tags, fat_leaf_pages=fat_leaf_pages, next_title_by_path=next_title_by_path, model_name=_resolve_hierarchy_model(page_memory_config), @@ -200,8 +227,6 @@ def _run_fine_hierarchy_for_scope( skeletons=active_skeletons, page_count=page_count, strategy=f"{strategy}:refined", - processing_pages=processing_pages, - excluded_toc_pages=excluded_toc_pages, ) record_stage( scope_stages, "C4b.fine_hierarchy", @@ -215,6 +240,67 @@ def _run_fine_hierarchy_for_scope( if token_cost_tracker is not None: token_cost_tracker.snapshot_stage(f"C4b.fine_hierarchy:{scope_id}") + final_pages = pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=active_skeletons, + page_count=page_count, + ), + toc_pages, + ) + final_scope_summary = _summarize_tag_scope( + skeletons=active_skeletons, + page_count=page_count, + pages=final_pages, + ) + rendered = render_document_pages( + pdf_path=pdf_path, + page_count=page_count, + output_dir=str(out_dir), + scope_id=scope_id, + pages=final_pages, + page_features=page_features, + page_texts=page_texts, + ) + record_stage( + scope_stages, + "C1.render_pages", + page_info=page_scope_info([item.page_index for item in rendered]), + variables={ + "scope_id": scope_id, + "rendered_count": len(rendered), + "tag_scope": final_scope_summary, + }, + ) + + plans = derive_page_processing_plan( + page_count=page_count, + page_labels=page_labels, + page_features=page_features, + ) + final_page_set = set(final_pages) + plans = [plan for plan in plans if plan.page_index in final_page_set] + record_stage( + scope_stages, + "C2.page_plan", + page_info=page_scope_info([getattr(plan, "page_index", None) for plan in plans]), + variables={"scope_id": scope_id, "plan_count": len(plans)}, + ) + + tags = tag_pages( + pages=rendered, + plans=plans, + vlm_model=vlm_model, + max_concurrent=page_memory_config.tag_concurrency, + ) + record_stage( + scope_stages, + "C3.page_tagger", + page_info=page_scope_info([tag.page_index for tag in tags]), + variables={"scope_id": scope_id, "tags": _summarize_tags(tags)}, + ) + if token_cost_tracker is not None: + token_cost_tracker.snapshot_stage(f"C3.page_tagger:{scope_id}") + write_scope_artifacts( out_dir=out_dir, scope_id=scope_id, @@ -229,7 +315,7 @@ def _run_fine_hierarchy_for_scope( tags=tags, assets_by_page={}, rendered=rendered, - final_pages=processing_pages, + final_pages=final_pages, scope_manifest=scope_manifest, trace_stages=scope_stages, ) @@ -245,10 +331,9 @@ def main() -> int: args = parser.parse_args() from app.services.document_agent.pdf_text import read_page_texts + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.fine_hierarchy import build_next_title_by_path - from app.services.page_memory.memory_service import _render_and_tag_document_pages from app.services.page_memory.skeleton_extractor import SectionSkeleton - from toc_page_policy import TocPagePolicy from shared.models.schemas.page_memory_config import PageMemoryConfig pdf_path, filename, out_dir = resolve_paths(args) @@ -269,7 +354,8 @@ def main() -> int: page_count = anatomy.page_count page_features = anatomy.page_features if anatomy else [] page_labels = anatomy.page_labels if anatomy else [] - toc_policy = TocPagePolicy.from_anatomy(anatomy) + toc_result = getattr(anatomy, "toc_result", None) + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) page_memory_config = PageMemoryConfig.default() skeletons = load_pipeline_skeletons(state_path) @@ -307,10 +393,13 @@ def main() -> int: "start_page": 1, "end_page": page_count, "strategy": "fallback_root", - "processing_pages": toc_policy.filter_processing_pages( - list(range(1, page_count + 1)) + "processing_pages": pages_excluding_toc( + list(range(1, page_count + 1)), + toc_pages, + ), + "excluded_toc_pages": sorted( + {int(page) for page in toc_pages} ), - "excluded_toc_pages": sorted(toc_policy.pure_toc_pages), } ] logger.info(" no skeleton hierarchy → fallback Root scope p1-{}", page_count) @@ -373,11 +462,14 @@ def main() -> int: "start_page": pr_start, "end_page": pr_end, "strategy": "manual_page_range", - "processing_pages": toc_policy.filter_processing_pages( - requested_pages + "processing_pages": pages_excluding_toc( + requested_pages, + toc_pages, ), "excluded_toc_pages": sorted( - set(requested_pages) & toc_policy.pure_toc_pages + page + for page in toc_pages + if pr_start <= int(page) <= pr_end ), } ] @@ -450,66 +542,51 @@ def main() -> int: sum(1 for title in next_title_by_path.values() if title), ) - selected_processing_pages: set[int] = set() - scope_payloads: list[tuple[str, Path, list[int]]] = [] + scope_payloads: list[tuple[str, Path]] = [] + pages_needed: set[int] = set() for sid in scope_ids: scope_dir = scopes_dir / sid - scope_meta, scope_skeletons = load_scope_skeletons_artifact( + _scope_meta, scope_skeletons = load_scope_skeletons_artifact( scope_dir / "skeletons.json" ) - processing_pages, _excluded = _resolve_scope_processing_pages( - scope_meta=scope_meta, - skeletons=scope_skeletons, - page_count=page_count, - toc_policy=toc_policy, + pages_needed.update( + pages_excluding_toc( + _derive_hierarchy_page_scope( + skeletons=scope_skeletons, + page_count=page_count, + ), + toc_pages, + ) ) - selected_processing_pages.update(processing_pages) - scope_payloads.append((sid, scope_dir, processing_pages)) + scope_payloads.append((sid, scope_dir)) - processing_pages = sorted(selected_processing_pages) page_texts = read_page_texts( pdf_path, - processing_pages or list(range(1, page_count + 1)), + sorted(pages_needed) or list(range(1, page_count + 1)), timeout=300, ) logger.info( - " document page stage: {} processing pages (of {}), tag_concurrency={}", - len(processing_pages), + " page texts ready: {} pages (of {}), tag_concurrency={}", + len(page_texts), page_count, page_memory_config.tag_concurrency, ) vlm_model = getattr(args, "vlm_model", None) or os.environ.get("IMAGE_MODEL") - rendered_by_page, tags_by_page = _render_and_tag_document_pages( - pdf_path=pdf_path, - output_dir=str(out_dir), - page_count=page_count, - processing_pages=processing_pages, - page_texts=page_texts, - page_features=page_features, - page_labels=page_labels, - vlm_model=vlm_model, - toc_policy=toc_policy, - page_memory_config=page_memory_config, - trace_recorder=TraceStageAdapter(trace_stages), - ) - token_cost_tracker.snapshot_stage("C3.page_tagger") - logger.info( - " tagged {} unique pages; refining {} scopes", - len(tags_by_page), - len(scope_ids), - ) def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: return _run_fine_hierarchy_for_scope( scope_id=scope_id, scope_dir=scope_dir, out_dir=out_dir, + pdf_path=str(pdf_path), page_count=page_count, - rendered_by_page=rendered_by_page, - tags_by_page=tags_by_page, + page_texts=page_texts, + page_features=page_features, + page_labels=page_labels, + vlm_model=vlm_model, next_title_by_path=next_title_by_path, - toc_policy=toc_policy, + toc_pages=toc_pages, page_memory_config=page_memory_config, token_cost_tracker=token_cost_tracker, ) @@ -529,7 +606,7 @@ def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: sid, scope_dir, ) - for sid, scope_dir, _pages in scope_payloads + for sid, scope_dir in scope_payloads ] gevent.joinall(greenlets, raise_error=True) scope_results = [cast(ScopeResult, g.value) for g in greenlets] @@ -537,7 +614,7 @@ def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: logger.info(" serial fine hierarchy: {} scope(s)", len(scope_ids)) scope_results = [ _run_selected_scope(sid, scope_dir) - for sid, scope_dir, _pages in scope_payloads + for sid, scope_dir in scope_payloads ] for sr in scope_results: @@ -546,6 +623,10 @@ def _run_selected_scope(scope_id: str, scope_dir: Path) -> ScopeResult: merged_skeletons = sort_skeletons( [skel for sr in scope_results for skel in sr.skeletons] ) + tags_by_page: dict[int, Any] = {} + for sr in scope_results: + for tag in sr.tags: + tags_by_page[int(tag.page_index)] = tag merged_tags = [tags_by_page[page] for page in sorted(tags_by_page)] if not partial_run: write_top_level_artifacts( diff --git a/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py b/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py index 41b06b105..127121c2f 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage4_assets.py @@ -64,8 +64,10 @@ def _load_scope_asset_context( scope_id: str, scope_dir: Path, page_count: int, - toc_policy: Any, + toc_pages: list[int], ) -> _ScopeAssetContext | None: + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc + fine_hierarchy_path = scope_dir / "fine_hierarchy.json" require_file( fine_hierarchy_path, @@ -83,26 +85,20 @@ def _load_scope_asset_context( final_pages = ( [int(page) for page in recorded_pages] if isinstance(recorded_pages, list) - else toc_policy.filter_processing_pages( + else pages_excluding_toc( _derive_hierarchy_page_scope( skeletons=active_skeletons, page_count=page_count, - ) + ), + toc_pages, ) ) - recorded_excluded = prior_scope.get("excluded_toc_pages") - excluded_toc_pages = ( - [int(page) for page in recorded_excluded] - if isinstance(recorded_excluded, list) - else sorted(toc_policy.pure_toc_pages) - ) + final_pages = pages_excluding_toc(final_pages, toc_pages) scope_manifest = _scope_manifest( scope_id=scope_id, skeletons=active_skeletons, page_count=page_count, strategy="fine:assets", - processing_pages=final_pages, - excluded_toc_pages=excluded_toc_pages, ) return _ScopeAssetContext( scope_id=scope_id, @@ -135,7 +131,6 @@ def main() -> int: page_asset_summary_enabled, ) from app.services.page_memory.page_renderer import render_document_pages - from toc_page_policy import TocPagePolicy from shared.models.schemas.page_memory_config import PageMemoryConfig pdf_path, filename, out_dir = resolve_paths(args) @@ -151,7 +146,8 @@ def main() -> int: anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) page_count = anatomy.page_count page_features = anatomy.page_features if anatomy else [] - toc_policy = TocPagePolicy.from_anatomy(anatomy) + toc_result = getattr(anatomy, "toc_result", None) + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) scope_ids = resolve_debug_scope_ids( scopes_dir=scopes_dir, @@ -183,13 +179,16 @@ def main() -> int: scope_id=scope_id, scope_dir=scopes_dir / scope_id, page_count=page_count, - toc_policy=toc_policy, + toc_pages=toc_pages, ) if context is not None: scope_contexts.append(context) - union_pages = sorted( - {page for context in scope_contexts for page in context.pages} + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc + + union_pages = pages_excluding_toc( + sorted({page for context in scope_contexts for page in context.pages}), + toc_pages, ) logger.info( "🔬 document C5: {} unique pages across {} scopes", @@ -201,6 +200,7 @@ def main() -> int: pdf_path=pdf_path, page_count=page_count, output_dir=str(out_dir), + scope_id="c5_assets", pages=union_pages, page_features=page_features, page_texts=page_texts, diff --git a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py index 281040759..079a3349e 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py @@ -71,16 +71,14 @@ def _run_tagging_for_scope( page_count: int, page_texts: dict[int, str], page_features: list[Any], - toc_policy: Any, + toc_pages: list[int], args: Any, - token_cost_tracker: TokenCostTracker | None = None, ) -> ScopeResult: """Load Stage-3 combined tags and rehydrate renders for final assembly.""" + from app.services.document_agent.structure.toc_anchoring import pages_excluding_toc from app.services.page_memory.page_renderer import render_document_pages scope_stages: list[dict[str, Any]] = [] - if token_cost_tracker is not None: - token_cost_tracker.register_child_thread() fine_hierarchy_path = scope_dir / "fine_hierarchy.json" require_file(fine_hierarchy_path, hint=f"Run Stage 3 to produce {fine_hierarchy_path}") @@ -114,26 +112,20 @@ def _run_tagging_for_scope( final_pages = ( [int(page) for page in recorded_pages] if isinstance(recorded_pages, list) - else toc_policy.filter_processing_pages( + else pages_excluding_toc( _derive_hierarchy_page_scope( skeletons=active_skeletons, page_count=page_count, - ) + ), + toc_pages, ) ) - recorded_excluded = prior_scope.get("excluded_toc_pages") - excluded_toc_pages = ( - [int(page) for page in recorded_excluded] - if isinstance(recorded_excluded, list) - else sorted(toc_policy.pure_toc_pages) - ) + final_pages = pages_excluding_toc(final_pages, toc_pages) scope_manifest = _scope_manifest( scope_id=scope_id, skeletons=active_skeletons, page_count=page_count, strategy="fine:finalize", - processing_pages=final_pages, - excluded_toc_pages=excluded_toc_pages, ) logger.info( "🔬 [scope {}] loaded {} combined tags for {} processing pages", @@ -147,6 +139,7 @@ def _run_tagging_for_scope( pdf_path=pdf_path, page_count=page_count, output_dir=str(out_dir), + scope_id=scope_id, pages=final_pages, page_features=page_features, page_texts=page_texts, @@ -274,16 +267,7 @@ def main() -> int: args.finalize = True from app.services.document_agent.pdf_text import read_page_texts - from app.services.document_agent.structure.hierarchy_locator import extract_toc_nodes - from app.services.page_memory.memory_service import ( - _append_toc_nav_skeletons, - _merge_static_toc_tags, - ) - from app.services.page_memory.node_assembler import ( - build_node_chunks, - build_toc_node_chunks, - merge_chunks_by_first_page, - ) + from app.services.page_memory.node_assembler import build_node_rows from app.services.page_memory.skeleton_extractor import collapse_single_child_chains from shared.models.schemas.page_memory_config import PageMemoryConfig @@ -300,9 +284,9 @@ def main() -> int: anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) page_count = anatomy.page_count page_features = anatomy.page_features if anatomy else [] - from toc_page_policy import TocPagePolicy - - toc_policy = TocPagePolicy.from_anatomy(anatomy) + page_labels = anatomy.page_labels if anatomy else [] + toc_result = getattr(anatomy, "toc_result", None) + toc_pages = list(getattr(toc_result, "toc_pages", None) or []) scope_ids = resolve_debug_scope_ids( scopes_dir=scopes_dir, @@ -340,9 +324,8 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: page_count=page_count, page_texts=page_texts, page_features=page_features, - toc_policy=toc_policy, + toc_pages=toc_pages, args=args, - token_cost_tracker=token_cost_tracker, ) if args.max_workers > 1 and len(scope_ids) > 1: @@ -440,22 +423,17 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: ) active_skeletons = all_skeletons - tags = _merge_static_toc_tags(all_tags, toc_policy) - nav_skeletons = _append_toc_nav_skeletons( - body_skeletons=active_skeletons, - anatomy=anatomy, - filename=filename, - ) + tags = all_tags # Write top-level artifacts write_top_level_artifacts( out_dir=out_dir, - hierarchy=nav_skeletons, + hierarchy=active_skeletons, tags=tags, assets_by_page=all_assets if all_assets else None, ) - # ── C7: Node assembly ── + # ── C7: Node assembly (same entry as production memory_service) ── logger.info("=" * 70) logger.info("🧱 C7: assemble canonical chunks") logger.info("=" * 70) @@ -473,25 +451,48 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: image_path_by_page[page] = rend.image_path page_memory_config = PageMemoryConfig.default() - body_chunks = build_node_chunks( + label_map: dict[int, str] = {} + if page_labels: + for lbl in page_labels: + label_map[int(lbl.page)] = str(lbl.kind) + rows = build_node_rows( skeletons=active_skeletons, raw_text_by_page=raw_text_by_page, image_path_by_page=image_path_by_page, + output_dir=str(out_dir), + kind_by_page=label_map, tag_by_page=tag_map, filename=filename, + verdict="page", vlm_model=args.vlm_model or os.environ.get("IMAGE_MODEL"), page_assets_by_page=all_assets if all_assets else None, + node_summary_max_pages=page_memory_config.node_summary_max_pages, node_assembly_concurrency=page_memory_config.node_assembly_concurrency, - body_start_by_page=toc_policy.body_start_by_page(), - ) - toc_chunks = build_toc_node_chunks(anatomy=anatomy, filename=filename) - canonical_chunks = merge_chunks_by_first_page(toc_chunks, body_chunks) - from shared.services.chunks.canonical_chunk_builder import chunks_as_json - - chunks = cast( - list[dict[str, Any]], - chunks_as_json(canonical_chunks), ) + chunks = [ + { + "chunk_id": str(row.get("know_id") or ""), + "type": str(row.get("type") or "page"), + "content": str(row.get("content") or ""), + "path": str(row.get("path") or ""), + "metadata": { + "length": int(row.get("length") or 0), + "summary": str(row.get("summary") or ""), + "page_nums": [ + int(part) + for part in str(row.get("page_nums") or "").split(",") + if str(part).strip().isdigit() + ], + "keywords": [ + part.strip() + for part in str(row.get("keywords") or "").split(";") + if part.strip() + ], + **(row.get("extra_metadata") or {}), + }, + } + for row in rows + ] logger.info(f" C7: {len(chunks)} canonical chunks") record_stage( trace_stages, @@ -569,6 +570,10 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: ) # ── Final trace + cross-stage cost rollup ── + from app.services.document_agent.structure.hierarchy_locator import ( + extract_toc_nodes, + ) + toc_nodes = ( extract_toc_nodes(anatomy.toc_hierarchies) if anatomy.toc_hierarchies else [] ) diff --git a/apps/worker/scripts/page_memory/toc_page_policy.py b/apps/worker/scripts/page_memory/toc_page_policy.py deleted file mode 100644 index 4a15f5f50..000000000 --- a/apps/worker/scripts/page_memory/toc_page_policy.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Debug-only TOC page policy for staged page_memory scripts. - -Production page_memory no longer ships this helper; debug stages still need a -single place to read ``toc_result.toc_pages`` and filter processing ranges. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - - -@dataclass(frozen=True) -class TocPagePolicy: - pure_toc_pages: frozenset[int] = field(default_factory=frozenset) - - @classmethod - def from_anatomy(cls, anatomy: Any | None) -> TocPagePolicy: - toc_result = getattr(anatomy, "toc_result", None) if anatomy is not None else None - pages = getattr(toc_result, "toc_pages", None) or [] - pure: set[int] = set() - for page in pages: - try: - pure.add(int(page)) - except (TypeError, ValueError): - continue - return cls(pure_toc_pages=frozenset(pure)) - - def filter_processing_pages(self, pages: list[int]) -> list[int]: - return [page for page in pages if page not in self.pure_toc_pages] - - def body_start_by_page(self) -> dict[int, float]: - return {} diff --git a/apps/worker/tests/contract/test_calibration_scan_contract.py b/apps/worker/tests/contract/test_calibration_scan_contract.py index 34a5aa4d1..72e053aa6 100644 --- a/apps/worker/tests/contract/test_calibration_scan_contract.py +++ b/apps/worker/tests/contract/test_calibration_scan_contract.py @@ -60,10 +60,22 @@ def __call__(self, ctx: ToolContext, args: dict[str, Any]) -> ToolResult: ) +def _fake_render_pages( + ctx: ToolContext, + pages: list[int], + **kwargs: Any, +) -> list[dict[str, Any]]: + return [ + {"page": int(page), "png_path": f"/tmp/scan_page_{int(page)}.png"} + for page in pages + ] + + @pytest.fixture def patch_inspect(monkeypatch: pytest.MonkeyPatch): def _apply(fake: _FakeInspect) -> _FakeInspect: monkeypatch.setattr(scan_module, "inspect_pages", fake) + monkeypatch.setattr(scan_module, "render_pages", _fake_render_pages) return fake return _apply @@ -241,3 +253,88 @@ def __call__(self, ctx: ToolContext, args: dict[str, Any]) -> ToolResult: assert result.found_page == 13 assert sorted(fake.calls) == [[10], [11], [12], [13], [14], [15]] + + +def test_window_renders_once_serially_before_concurrent_vlm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """PDF→PNG is one serial batch per window; VLM stays one page per call.""" + render_calls: list[list[int]] = [] + inspect_args: list[dict[str, Any]] = [] + + def _recording_render( + ctx: ToolContext, + pages: list[int], + **kwargs: Any, + ) -> list[dict[str, Any]]: + render_calls.append(list(pages)) + return _fake_render_pages(ctx, pages, **kwargs) + + def _capture_inspect(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + inspect_args.append(dict(args)) + pages = list(args.get("pages") or []) + hit = 10 in pages + return ToolResult( + status="ok", + payload={ + "pages": pages, + "answer": "", + "fields": { + "found": hit, + "reason": "hit" if hit else "miss", + }, + }, + ) + + monkeypatch.setattr(scan_module, "inspect_pages", _capture_inspect) + monkeypatch.setattr(scan_module, "render_pages", _recording_render) + + result = scan_title_forward( + ctx=_ctx(), title="Chapter 1", start_page=10, page_count=60 + ) + + assert result.found_page == 10 + assert render_calls == [[10, 11]] + assert len(inspect_args) == 2 + assert sorted(args["pages"] for args in inspect_args) == [[10], [11]] + for args in inspect_args: + assert args["page_cap"] == 1 + assert len(args["rendered_pages"]) == 1 + assert args["rendered_pages"][0]["page"] == args["pages"][0] + + +def test_concurrent_inspect_records_tokens_on_parse_tracker( + patch_inspect, +) -> None: + from shared.services.ai.token_tracking import ( + cleanup_token_tracker, + init_token_tracker, + record_tokens, + ) + + class _RecordingInspect(_FakeInspect): + def __call__(self, ctx: ToolContext, args: dict[str, Any]) -> ToolResult: + record_tokens( + { + "prompt_tokens": 3, + "completion_tokens": 1, + "total_tokens": 4, + }, + model="test-vlm", + task="calibration.scan_title_forward", + ) + return super().__call__(ctx, args) + + tracker = init_token_tracker() + try: + patch_inspect(_RecordingInspect(hit_page=10)) + result = scan_title_forward( + ctx=_ctx(), title="Chapter 1", start_page=10, page_count=60 + ) + assert result.found_page == 10 + assert tracker["calls"] == 2 + assert tracker["total_tokens"] == 8 + assert tracker["by_task"]["calibration.scan_title_forward"]["total_tokens"] == 8 + assert tracker["by_model"]["test-vlm"]["total_tokens"] == 8 + finally: + cleanup_token_tracker() diff --git a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py index c78ba9dc5..d2eaf126a 100644 --- a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py +++ b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py @@ -53,6 +53,33 @@ def test_compute_fat_leaf_pages_uses_exclusive_boundaries() -> None: } +def test_compute_fat_leaf_pages_excludes_toc_pages_from_span() -> None: + """Closed range 1-6 with toc=[2,3,4] is only 3 body pages → not fat when min=4.""" + skeletons = [ + SectionSkeleton( + section_path="demo.pdf/Abbreviations", + level=1, + start_page=1, + end_page=6, + title="Abbreviations", + parent_path="demo.pdf", + ), + ] + assert ( + fine_hierarchy.compute_fat_leaf_pages( + skeletons, + min_pages=4, + toc_pages=[2, 3, 4], + ) + == set() + ) + assert fine_hierarchy.compute_fat_leaf_pages( + skeletons, + min_pages=2, + toc_pages=[2, 3, 4], + ) == {1, 5, 6} + + def test_refine_fat_leaf_skeletons_excludes_next_section_start_when_unordered( monkeypatch, ) -> None: diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index 58699fdeb..773f14a2e 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -113,6 +113,7 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={ 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), @@ -161,6 +162,7 @@ def _fake_compute_node_summary(**kwargs): skeletons=_ordered_page_skeletons(), raw_text_by_page={1: "text-1", 2: "text-2", 3: "text-3"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={}, filename="demo.pdf", @@ -202,6 +204,7 @@ def _fake_compute_node_summary(**kwargs): skeletons=_ordered_page_skeletons()[:1], raw_text_by_page={1: "text-1"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={}, filename="demo.pdf", @@ -212,14 +215,16 @@ def _fake_compute_node_summary(**kwargs): def test_build_node_rows_attaches_page_citation_assets_for_rendered_pages(tmp_path) -> None: - page_image = tmp_path / "pages" / "page-231.png" - page_image.parent.mkdir() + # Nested scope path matches production pages/{scope_id}/page-{n}.png. + page_image = tmp_path / "pages" / "p231-231" / "page-231.png" + page_image.parent.mkdir(parents=True) Image.new("RGB", (2, 3), color=(255, 255, 255)).save(page_image) rows = node_assembler.build_node_rows( skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={231: str(page_image)}, + output_dir=str(tmp_path), kind_by_page={}, tag_by_page={ 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), @@ -244,6 +249,8 @@ def test_build_node_rows_attaches_page_citation_assets_for_rendered_pages(tmp_pa } ] assert (tmp_path / "page_citation_assets" / "page-231.png").is_file() + # Must not land under the temporary pages/ tree. + assert not (tmp_path / "pages" / "page_citation_assets").exists() def test_build_node_rows_keeps_internal_section_body_pages() -> None: @@ -268,6 +275,7 @@ def test_build_node_rows_keeps_internal_section_body_pages() -> None: skeletons=[parent, child], raw_text_by_page={233: "parent body", 234: "child body"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={ 233: PageTagResult(page_index=233, summary="s233", keywords=["parent"]), @@ -350,6 +358,7 @@ def chat_completion_with_usage(self, **kwargs): skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={231: str(img), 232: str(img)}, + output_dir=str(tmp_path), kind_by_page={}, tag_by_page={ 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), @@ -390,6 +399,7 @@ def test_build_node_rows_prepends_asset_rows_and_links_page_nodes() -> None: skeletons=_same_page_sibling_skeletons(), raw_text_by_page={231: "text-231", 232: "text-232"}, image_path_by_page={}, + output_dir="/tmp/knowhere-test-node-assembler", kind_by_page={}, tag_by_page={ 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), diff --git a/apps/worker/tests/contract/test_page_memory_page_renderer_scope_path_contract.py b/apps/worker/tests/contract/test_page_memory_page_renderer_scope_path_contract.py new file mode 100644 index 000000000..6effce400 --- /dev/null +++ b/apps/worker/tests/contract/test_page_memory_page_renderer_scope_path_contract.py @@ -0,0 +1,64 @@ +"""Contract: page_memory renders into per-scope paths with page-{n}.png names.""" + +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from pathlib import Path + +import pytest + +import app.services.page_memory.page_renderer as page_renderer + + +def test_render_document_pages_requires_scope_id(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="scope_id"): + page_renderer.render_document_pages( + pdf_path=str(tmp_path / "missing.pdf"), + page_count=1, + output_dir=str(tmp_path), + scope_id=" ", + pages=[1], + ) + + +def test_render_document_pages_writes_under_scope_folder( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + captured: dict[str, object] = {} + + def _fake_render_pages(ctx, pages, **kwargs): + captured["folder_name"] = kwargs.get("folder_name") + captured["prefix"] = kwargs.get("prefix") + captured["output_dir"] = ctx.output_dir + folder = Path(ctx.output_dir) / str(kwargs["folder_name"]) + folder.mkdir(parents=True, exist_ok=True) + results = [] + for page in pages: + png_path = folder / f"page-{page}.png" + png_path.write_bytes(b"png") + results.append({"page": page, "png_path": str(png_path)}) + return results + + monkeypatch.setattr(page_renderer, "render_pages", _fake_render_pages) + + results = page_renderer.render_document_pages( + pdf_path=str(tmp_path / "doc.pdf"), + page_count=78, + output_dir=str(tmp_path), + scope_id="p58-60", + pages=[60], + ) + + assert captured["folder_name"] == "pages/p58-60" + assert captured["prefix"] == "" + assert len(results) == 1 + assert results[0].image_path.endswith("/pages/p58-60/page-60.png") + assert Path(results[0].image_path).is_file() diff --git a/apps/worker/tests/contract/test_parse_usage_task_contract.py b/apps/worker/tests/contract/test_parse_usage_task_contract.py new file mode 100644 index 000000000..c108e21d5 --- /dev/null +++ b/apps/worker/tests/contract/test_parse_usage_task_contract.py @@ -0,0 +1,67 @@ +"""Parse-pipeline LLM calls must carry an explicit usage_task.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +_WORKER_SERVICES = Path(__file__).resolve().parents[2] / "app" / "services" +_WORKER_SCRIPTS = Path(__file__).resolve().parents[2] / "scripts" +_PARSE_ROOTS = ( + _WORKER_SERVICES / "document_parser", + _WORKER_SERVICES / "document_agent", + _WORKER_SERVICES / "page_memory", + _WORKER_SERVICES / "connect_builder", + _WORKER_SCRIPTS, +) +_LLM_CALL_NAMES = { + "chat_completion", + "chat_completion_with_usage", + "chat_completion_raw_with_usage", + "summarize", + "transcribe", +} + + +def _call_name(node: ast.Call) -> str | None: + func = node.func + if isinstance(func, ast.Name): + return func.id + if isinstance(func, ast.Attribute): + return func.attr + return None + + +def _iter_parse_python_files() -> list[Path]: + files: list[Path] = [] + for root in _PARSE_ROOTS: + if not root.exists(): + continue + files.extend(path for path in root.rglob("*.py") if path.is_file()) + return files + + +def _relative_path(path: Path) -> str: + for anchor in (_WORKER_SERVICES, _WORKER_SCRIPTS): + try: + return str(path.relative_to(anchor)) + except ValueError: + continue + return str(path) + + +def test_parse_llm_calls_pass_explicit_usage_task() -> None: + missing: list[str] = [] + for path in _iter_parse_python_files(): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + name = _call_name(node) + if name not in _LLM_CALL_NAMES: + continue + if any(keyword.arg == "usage_task" for keyword in node.keywords): + continue + rel = _relative_path(path) + missing.append(f"{rel}:{node.lineno}:{name}") + assert missing == [] diff --git a/apps/worker/tests/contract/test_toc_rehome_contract.py b/apps/worker/tests/contract/test_toc_rehome_contract.py new file mode 100644 index 000000000..9e4f4b87f --- /dev/null +++ b/apps/worker/tests/contract/test_toc_rehome_contract.py @@ -0,0 +1,662 @@ +"""Contract tests for same-forest TOC paged-leaf rehome. Synthetic trees only.""" + +from __future__ import annotations + +import os +from unittest.mock import patch + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.document_agent.manifest import TocResult, ToolContext +from app.services.document_agent.state import ProfileBlackboard +from app.services.document_agent.structure.anchoring_primitives import ( + SkeletonAnchor, +) +from app.services.document_agent.structure.hierarchy_locator import ( + TitleMatch, + TitleNode, + resolve_hierarchy_page_ranges, +) +from app.services.document_agent.structure.toc_anchoring import run_toc_anchoring +from app.services.document_agent.structure.toc_rehome import rehome_forest + + +def _match(title: str, page: int) -> TitleMatch: + return TitleMatch( + page=page, + source="anchored", + matched_line=title, + candidates=[page], + evidence={}, + ) + + +def _overrides(pages: dict[tuple[str, ...], int]) -> dict[tuple[str, ...], TitleMatch]: + return {path: _match(path[-1], page) for path, page in pages.items()} + + +def _titles(nodes: list[TitleNode]) -> list[str]: + return [node.title for node in nodes] + + +def test_same_parent_backjump_inserts_after_nearest_prev_leaf() -> None: + nodes = [ + TitleNode(title="A", level=1, printed_page=10), + TitleNode(title="B", level=1, printed_page=20), + TitleNode(title="C", level=1, printed_page=15), + ] + result = rehome_forest( + nodes, + _overrides({("A",): 10, ("B",): 20, ("C",): 15}), + ) + # C@15 nearest prev leaf A@10 → after A under root. + assert _titles(result.nodes) == ["A", "C", "B"] + assert result.match_overrides[("C",)].page == 15 + assert not result.nodes[1].children + + +def test_no_prev_leaf_by_page_is_noop_for_that_backjump() -> None: + """Backjump with no TOC-earlier leaf page<=P is not moved (no fallback).""" + nodes = [ + TitleNode(title="A", level=1, printed_page=10), + TitleNode(title="B", level=1, printed_page=20), + TitleNode(title="C", level=1, printed_page=5), + ] + result = rehome_forest( + nodes, + _overrides({("A",): 10, ("B",): 20, ("C",): 5}), + ) + assert _titles(result.nodes) == ["A", "B", "C"] + assert result.events == [] + + +def test_equal_page_backjump_inserts_after_front_same_page_leaves() -> None: + """Rehome @116 with one or more front @116 leaves → after that same-page group.""" + nodes = [ + TitleNode(title="A116", level=1, printed_page=116), + TitleNode(title="B116", level=1, printed_page=116), + TitleNode(title="Late", level=1, printed_page=200), + TitleNode(title="C116", level=1, printed_page=116), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("A116",): 116, + ("B116",): 116, + ("Late",): 200, + ("C116",): 116, + } + ), + ) + assert _titles(result.nodes) == ["A116", "B116", "C116", "Late"] + assert result.events[0]["anchor_path"] == ["B116"] + + +def test_unpaged_shell_leaf_inserts_after_nearest_prev_and_drops_shell() -> None: + nodes = [ + TitleNode(title="Preamble", level=1, printed_page=6), + TitleNode(title="References", level=1, printed_page=61), + TitleNode( + title="List of tables", + level=1, + children=[ + TitleNode(title="Table 1", level=2, printed_page=7), + TitleNode(title="Table 2", level=2, printed_page=9), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Preamble",): 6, + ("References",): 61, + ("List of tables", "Table 1"): 7, + ("List of tables", "Table 2"): 9, + } + ), + ) + assert "List of tables" not in _titles(result.nodes) + # Table1 after Preamble@6; Table2 after Table1@7 — all stay leaves. + assert _titles(result.nodes) == ["Preamble", "Table 1", "Table 2", "References"] + assert not result.nodes[1].children + assert not result.nodes[2].children + assert result.match_overrides[("Table 1",)].page == 7 + assert result.match_overrides[("Table 2",)].page == 9 + + +def test_equal_page_points_at_same_first_segment_group() -> None: + """Both @117 leaves anchor only into the first monotonic segment.""" + nodes = [ + TitleNode(title="Prev", level=1, printed_page=116), + TitleNode(title="Late", level=1, printed_page=200), + TitleNode( + title="Shell", + level=1, + children=[ + TitleNode(title="T1", level=2, printed_page=117), + TitleNode(title="T2", level=2, printed_page=117), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Prev",): 116, + ("Late",): 200, + ("Shell", "T1"): 117, + ("Shell", "T2"): 117, + } + ), + ) + assert "Shell" not in _titles(result.nodes) + assert _titles(result.nodes) == ["Prev", "T1", "T2", "Late"] + assert [event["anchor_path"] for event in result.events] == [ + ["Prev"], + ["Prev"], + ] + + +def test_monotonic_forest_is_noop() -> None: + nodes = [ + TitleNode(title="A", level=1, printed_page=8), + TitleNode(title="B", level=1, printed_page=14), + TitleNode(title="C", level=1, printed_page=20), + ] + overrides = _overrides({("A",): 8, ("B",): 14, ("C",): 20}) + result = rehome_forest(nodes, overrides) + assert _titles(result.nodes) == ["A", "B", "C"] + assert result.match_overrides == overrides + assert result.events == [] + + +def test_every_leaf_in_post_break_segment_attempts_first_segment() -> None: + nodes = [ + TitleNode( + title="Chapter", + level=1, + printed_page=10, + children=[ + TitleNode(title="Intro", level=2, printed_page=11), + TitleNode(title="Late", level=2, printed_page=30), + TitleNode(title="Early", level=2, printed_page=12), + ], + ), + TitleNode(title="Next", level=1, printed_page=40), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Chapter",): 10, + ("Chapter", "Intro"): 11, + ("Chapter", "Late"): 30, + ("Chapter", "Early"): 12, + ("Next",): 40, + } + ), + ) + assert _titles(result.nodes) == ["Chapter"] + assert _titles(result.nodes[0].children) == ["Intro", "Early", "Late", "Next"] + assert ("Chapter", "Early") in result.match_overrides + assert ("Chapter", "Next") in result.match_overrides + + +def test_each_later_segment_anchors_only_to_first_segment() -> None: + nodes = [ + TitleNode(title="A", level=1, printed_page=10), + TitleNode(title="B", level=1, printed_page=20), + TitleNode( + title="Shell 1", + level=1, + children=[ + TitleNode(title="C", level=2, printed_page=15), + TitleNode(title="D", level=2, printed_page=17), + ], + ), + TitleNode( + title="Shell 2", + level=1, + children=[ + TitleNode(title="E", level=2, printed_page=12), + TitleNode(title="F", level=2, printed_page=16), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("A",): 10, + ("B",): 20, + ("Shell 1", "C"): 15, + ("Shell 1", "D"): 17, + ("Shell 2", "E"): 12, + ("Shell 2", "F"): 16, + } + ), + ) + + assert "Shell 1" not in _titles(result.nodes) + assert "Shell 2" not in _titles(result.nodes) + assert _titles(result.nodes) == ["A", "E", "C", "F", "D", "B"] + event_by_title = { + event["new_path"][-1]: event for event in result.events + } + assert event_by_title["C"]["anchor_path"] == ["A"] + assert event_by_title["D"]["anchor_path"] == ["A"] + assert event_by_title["E"]["anchor_path"] == ["A"] + assert event_by_title["F"]["anchor_path"] == ["A"] + + +def test_prune_duplicate_when_same_path_and_same_page_as_first_segment() -> None: + nodes = [ + TitleNode(title="Abbreviations", level=1, printed_page=4), + TitleNode(title="Body", level=1, printed_page=10), + TitleNode(title="Late", level=1, printed_page=50), + TitleNode(title="Abbreviations", level=1, printed_page=4), + ] + overrides = _overrides( + { + ("Abbreviations",): 4, + ("Body",): 10, + ("Late",): 50, + } + ) + result = rehome_forest(nodes, overrides) + assert _titles(result.nodes) == ["Abbreviations", "Body", "Late"] + assert len(result.events) == 1 + assert result.events[0]["action"] == "pruned" + assert result.events[0]["source_path"] == ["Abbreviations"] + assert result.match_overrides[("Abbreviations",)].page == 4 + assert "toc_rehome" not in (result.match_overrides[("Abbreviations",)].evidence or {}) + + +def test_prune_same_path_page_even_when_page_is_toc_page() -> None: + """same_path_page is judged before toc_page skip and still prunes.""" + nodes = [ + TitleNode(title="Abbreviations", level=1, printed_page=4), + TitleNode(title="Body", level=1, printed_page=10), + TitleNode(title="Late", level=1, printed_page=50), + TitleNode(title="Abbreviations", level=1, printed_page=4), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Abbreviations",): 4, + ("Body",): 10, + ("Late",): 50, + } + ), + toc_pages=[4], + ) + assert _titles(result.nodes) == ["Abbreviations", "Body", "Late"] + assert result.events[0]["action"] == "pruned" + assert "toc_rehome" not in (result.match_overrides[("Abbreviations",)].evidence or {}) + + +def test_skip_rehome_when_physical_page_is_toc_page() -> None: + nodes = [ + TitleNode(title="Lead", level=1, printed_page=5), + TitleNode(title="Body", level=1, printed_page=20), + TitleNode( + title="Shell", + level=1, + children=[TitleNode(title="Early", level=2, printed_page=3)], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Lead",): 5, + ("Body",): 20, + ("Shell", "Early"): 3, + } + ), + toc_pages=[2, 3, 4], + ) + assert result.events == [] + assert "Shell" in _titles(result.nodes) + assert _titles(result.nodes[2].children) == ["Early"] + assert "toc_rehome" not in ( + result.match_overrides[("Shell", "Early")].evidence or {} + ) + + +def test_rehome_leaves_attach_to_resolved_scope_without_cutting_boundaries() -> None: + nodes = [ + TitleNode(title="Lead", level=1, printed_page=21), + TitleNode( + title="Chapter 7", + level=1, + printed_page=22, + children=[ + TitleNode(title="7.1", level=2, printed_page=24), + TitleNode(title="7.2", level=2, printed_page=26), + ], + ), + TitleNode(title="Recommendation", level=1, printed_page=22), + TitleNode(title="Table 7", level=1, printed_page=23), + TitleNode(title="Chapter 8", level=1, printed_page=27), + ] + overrides = _overrides( + { + ("Lead",): 21, + ("Chapter 7",): 22, + ("Chapter 7", "7.1"): 24, + ("Chapter 7", "7.2"): 26, + ("Recommendation",): 22, + ("Table 7",): 23, + ("Chapter 8",): 27, + } + ) + for path in (("Recommendation",), ("Table 7",)): + overrides[path] = TitleMatch( + page=overrides[path].page, + source=overrides[path].source, + matched_line=overrides[path].matched_line, + candidates=overrides[path].candidates, + evidence={"toc_rehome": {"segment_index": 1}}, + ) + + ranges = resolve_hierarchy_page_ranges( + nodes, + page_count=30, + match_overrides=overrides, + ) + by_title = {item.title: item for item in ranges} + + assert by_title["7.1"].match is not None + assert by_title["7.1"].start_page == 24 + assert by_title["7.1"].evidence["source"] == "anchored" + assert by_title["7.2"].start_page == 26 + assert (by_title["Recommendation"].start_page, by_title["Recommendation"].end_page) == ( + 22, + 24, + ) + assert (by_title["Table 7"].start_page, by_title["Table 7"].end_page) == ( + 22, + 24, + ) + assert by_title["Recommendation"].evidence["status"] == "rehome_attached" + assert by_title["Table 7"].evidence["status"] == "rehome_attached" + + +def test_rehome_leaf_attaches_by_physical_page_after_structural_resolve() -> None: + nodes = [ + TitleNode( + title="Chapter 3", + level=1, + printed_page=7, + children=[ + TitleNode(title="3.1", level=2, printed_page=7), + TitleNode(title="3.2", level=2, printed_page=8), + TitleNode(title="Table 3", level=2, printed_page=9), + ], + ), + TitleNode(title="Chapter 4", level=1, printed_page=8), + TitleNode(title="Chapter 5", level=1, printed_page=12), + ] + overrides = _overrides( + { + ("Chapter 3",): 7, + ("Chapter 3", "3.1"): 7, + ("Chapter 3", "3.2"): 8, + ("Chapter 3", "Table 3"): 9, + ("Chapter 4",): 8, + ("Chapter 5",): 12, + } + ) + table_path = ("Chapter 3", "Table 3") + table_match = overrides[table_path] + overrides[table_path] = TitleMatch( + page=table_match.page, + source=table_match.source, + matched_line=table_match.matched_line, + candidates=table_match.candidates, + evidence={"toc_rehome": {"segment_index": 1}}, + ) + + ranges = resolve_hierarchy_page_ranges( + nodes, + page_count=15, + match_overrides=overrides, + ) + by_title = {item.title: item for item in ranges} + + assert (by_title["3.2"].start_page, by_title["3.2"].end_page) == (8, 8) + assert (by_title["Table 3"].start_page, by_title["Table 3"].end_page) == ( + 8, + 12, + ) + assert by_title["Table 3"].evidence["scope_host_path"] == ["Chapter 4"] + assert [item.title for item in ranges] == [ + "3.1", + "3.2", + "Table 3", + "Chapter 4", + "Chapter 5", + ] + + +def test_nested_leaf_inserts_after_nearest_root_leaf() -> None: + nodes = [ + TitleNode(title="Chapter", level=1, printed_page=10), + TitleNode( + title="List", + level=1, + children=[ + TitleNode( + title="FrontMatter", + level=2, + printed_page=50, + children=[ + TitleNode(title="Late", level=3, printed_page=80), + TitleNode(title="Early", level=3, printed_page=55), + ], + ), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Chapter",): 10, + ("List", "FrontMatter"): 50, + ("List", "FrontMatter", "Late"): 80, + ("List", "FrontMatter", "Early"): 55, + } + ), + ) + # Early@55 after Chapter@10 at root; List (parent shell with Late) stays. + assert _titles(result.nodes) == ["Chapter", "List", "Early"] + assert not result.nodes[2].children + front = result.nodes[1].children[0] + assert front.title == "FrontMatter" + assert _titles(front.children) == ["Late"] + assert result.match_overrides[("Early",)].page == 55 + + +def test_non_leaf_not_rehomed_even_if_page_looks_early() -> None: + nodes = [ + TitleNode(title="Chapter", level=1, printed_page=10), + TitleNode( + title="List", + level=1, + children=[ + TitleNode( + title="FrontMatter", + level=2, + printed_page=5, + children=[ + TitleNode(title="Child", level=3, printed_page=50), + ], + ), + ], + ), + ] + result = rehome_forest( + nodes, + _overrides( + { + ("Chapter",): 10, + ("List", "FrontMatter"): 5, + ("List", "FrontMatter", "Child"): 50, + } + ), + ) + assert _titles(result.nodes) == ["Chapter", "List"] + assert result.nodes[1].children[0].title == "FrontMatter" + assert result.events == [] + + +def _ctx(*, hierarchies: list[dict], page_count: int = 80) -> ToolContext: + ctx = ToolContext( + pdf_path="/tmp/rehome.pdf", + job_id="rehome-test", + blackboard=ProfileBlackboard(page_count=page_count), + trace=None, + settings={}, + ) + ctx.blackboard.toc_hierarchies = hierarchies + ctx.blackboard.toc_result = TocResult(method="vlm_batch", toc_pages=[1]) + ctx.blackboard.page_full_text_cache = { + page: "body" for page in range(1, page_count + 1) + } + return ctx + + +def test_run_toc_anchoring_single_toc_still_runs_rehome() -> None: + hierarchies = [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Alpha", "level": 1, "page_number": 10}, + {"heading": "Bravo", "level": 1, "page_number": 20}, + {"heading": "Charlie", "level": 1, "page_number": 15}, + ], + } + ] + ctx = _ctx(hierarchies=hierarchies, page_count=30) + + def _fake_anchor(**kwargs): + nodes = [ + TitleNode(title="Alpha", level=1, printed_page=10), + TitleNode(title="Bravo", level=1, printed_page=20), + TitleNode(title="Charlie", level=1, printed_page=15), + ] + anchor = SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides=_overrides( + {("Alpha",): 10, ("Bravo",): 20, ("Charlie",): 15} + ), + null_page_report=[], + bulk_count=3, + ) + return nodes, anchor + + with patch( + "app.services.document_agent.calibration.orchestrator.anchor_hierarchy", + side_effect=_fake_anchor, + ): + run_toc_anchoring(ctx) + + assert ctx.blackboard.skeleton_anchor is not None + assert ctx.blackboard.skeleton_nodes is not None + assert ctx.blackboard.pending_skeleton_anchors == [] + assert [node["title"] for node in ctx.blackboard.skeleton_nodes] == [ + "Alpha", + "Charlie", + "Bravo", + ] + + +def test_run_toc_anchoring_rehome_before_classify_keeps_contained_graft() -> None: + """Pending calibrate must succeed without opening pdf_path; graft after rehome. + + Stub the calibration *internals* (same pattern as toc_graft contracts), not + ``_calibrate_pending_tocs`` itself. Patching that helper by string path is + brittle across the full contract suite and lets the real pending path try to + render ``pdf_path`` (``/tmp/rehome.pdf``), then drop the pending record. + """ + hierarchies = [ + { + "toc_range": [1, 1], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Host", "level": 1, "page_number": 10}, + {"heading": "Tail", "level": 1, "page_number": 40}, + ], + }, + { + "toc_range": [5, 5], + "toc_range_unit": "page", + "toc_with_level": [ + {"heading": "Inner", "level": 1, "page_number": 12}, + ], + }, + ] + ctx = _ctx(hierarchies=hierarchies, page_count=50) + + def _fake_anchor(**kwargs): + nodes = [ + TitleNode(title="Host", level=1, printed_page=10), + TitleNode(title="Tail", level=1, printed_page=40), + ] + anchor = SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides=_overrides({("Host",): 10, ("Tail",): 40}), + null_page_report=[], + bulk_count=2, + ) + return nodes, anchor + + inner = TitleNode(title="Inner", level=1, printed_page=12) + inner_anchor = SkeletonAnchor( + offset=0, + offset_status="ok", + match_overrides=_overrides({("Inner",): 12}), + null_page_report=[], + bulk_count=1, + ) + + with ( + patch( + "app.services.document_agent.calibration.orchestrator.anchor_hierarchy", + side_effect=_fake_anchor, + ), + patch( + "app.services.document_agent.calibration.service.calibrate_offset", + return_value=object(), + ), + patch( + "app.services.document_agent.calibration.procedure.pick_primary_offset", + return_value=0, + ), + patch( + "app.services.document_agent.calibration.procedure.finalize_calibration_result", + return_value=([inner], inner_anchor, True), + ), + ): + run_toc_anchoring(ctx) + + record = ctx.blackboard.pending_skeleton_anchors[0] + assert record["relationship"] == "contained" + assert record.get("grafted") is True + host = next( + node for node in ctx.blackboard.skeleton_nodes if node["title"] == "Host" + ) + assert any(child["title"] == "Inner" for child in host["children"]) diff --git a/apps/worker/tests/unit/test_debug_text_track_token_ledger.py b/apps/worker/tests/unit/test_debug_text_track_token_ledger.py new file mode 100644 index 000000000..4d7078322 --- /dev/null +++ b/apps/worker/tests/unit/test_debug_text_track_token_ledger.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +_SCRIPTS_ROOT = Path(__file__).resolve().parents[2] / "scripts" +sys.path.insert(0, str(_SCRIPTS_ROOT)) + +from _debug_token_ledger import ( + aggregate_stage_deltas, + empty_token_usage, + token_usage_delta, +) + +_TEXT_TRACK_STAGES = ("profile", "mineru", "hierarchy", "full") + + +def test_token_usage_delta_and_reuse_merge_match_stage_costs_semantics() -> None: + prev = empty_token_usage() + after_profile = { + "prompt_tokens": 10, + "completion_tokens": 4, + "total_tokens": 14, + "calls": 1, + "by_model": {"vlm": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14, "calls": 1}}, + "by_task": { + "document_agent.coarse_profile": { + "prompt_tokens": 10, + "completion_tokens": 4, + "total_tokens": 14, + "calls": 1, + } + }, + } + profile_delta = token_usage_delta(prev, after_profile) + assert profile_delta["total_tokens"] == 14 + assert profile_delta["by_task"]["document_agent.coarse_profile"]["calls"] == 1 + + # reuse-profile starts a fresh tracker; only the new stage exists in this run. + reuse_run = token_usage_delta(empty_token_usage(), { + "prompt_tokens": 8, + "completion_tokens": 2, + "total_tokens": 10, + "calls": 1, + "by_model": { + "text": {"prompt_tokens": 8, "completion_tokens": 2, "total_tokens": 10, "calls": 1} + }, + "by_task": { + "parser.heading_hierarchy": { + "prompt_tokens": 8, + "completion_tokens": 2, + "total_tokens": 10, + "calls": 1, + } + }, + }) + ledger = { + "stages": { + "profile": {"token_usage": profile_delta}, + "hierarchy": {"token_usage": reuse_run}, + } + } + merged = aggregate_stage_deltas(ledger, _TEXT_TRACK_STAGES) + assert merged["total_tokens"] == 24 + assert merged["calls"] == 2 + assert merged["by_task"]["document_agent.coarse_profile"]["total_tokens"] == 14 + assert merged["by_task"]["parser.heading_hierarchy"]["total_tokens"] == 10 diff --git a/apps/worker/tests/unit/test_processing_metadata_persist.py b/apps/worker/tests/unit/test_processing_metadata_persist.py index 4cd3d8200..340426e49 100644 --- a/apps/worker/tests/unit/test_processing_metadata_persist.py +++ b/apps/worker/tests/unit/test_processing_metadata_persist.py @@ -104,3 +104,52 @@ def test_record_processing_completion_persists_token_usage_to_job_row( assert updates["stages"]["token_usage"]["total_tokens"] == 10 assert "processing_completed_at" in updates assert "processing_duration_ms" in updates + + +def test_run_parse_job_failure_persists_token_usage(monkeypatch: object) -> None: + import app.services.document_ingestion.processing_run as processing_run + + persisted: dict[str, object] = {} + + def fake_persist(*, job_id: str, job_context: object, metadata_updates: dict[str, object]) -> None: + persisted.update(metadata_updates) + + monkeypatch.setattr(processing_run, "persist_job_metadata_updates", fake_persist) + monkeypatch.setattr( + processing_run, + "init_token_tracker", + lambda: {"prompt_tokens": 6, "completion_tokens": 1, "total_tokens": 7, "calls": 1}, + ) + monkeypatch.setattr(processing_run, "init_stage_tracker", lambda: {"worker.parse.document": 12}) + monkeypatch.setattr(processing_run, "init_llm_overrides", lambda *_args: None) + monkeypatch.setattr(processing_run, "cleanup_llm_overrides", lambda: None) + monkeypatch.setattr(processing_run, "cleanup_token_tracker", lambda: None) + monkeypatch.setattr(processing_run, "cleanup_stage_tracker", lambda: None) + + def _boom(**_kwargs: object) -> object: + raise RuntimeError("parse failed after LLM work") + + monkeypatch.setattr(processing_run, "prepare_source_file", _boom) + + lifecycle = Mock() + job_context = _job_context() + try: + processing_run._run_parse_job( + job_id="job-fail", + job_context=job_context, + lifecycle_service=lifecycle, + task_workspace=SimpleNamespace( + input_dir="/tmp", + output_dir="/tmp", + root_dir="/tmp", + ), + ) + except RuntimeError as exc: + assert "parse failed after LLM work" in str(exc) + else: + raise AssertionError("expected parse failure") + + stages = persisted["stages"] + assert isinstance(stages, dict) + assert stages["token_usage"]["total_tokens"] == 7 + assert stages["timing_ms"]["worker.parse.document"] == 12 diff --git a/packages/shared-python/shared/services/ai/summary/engine.py b/packages/shared-python/shared/services/ai/summary/engine.py index 2066ab795..970a08657 100644 --- a/packages/shared-python/shared/services/ai/summary/engine.py +++ b/packages/shared-python/shared/services/ai/summary/engine.py @@ -7,7 +7,6 @@ - prompt construction (via the shared ``build_prompt`` registry), - the text-or-vision LLM call (with optional page/asset images), - JSON parsing with one retry, -- budget reserve/commit/refund for visual calls, - deterministic language locking for text input. Two public functions: @@ -38,18 +37,20 @@ BodySummary, Entity, ) -from shared.utils.token_estimate import estimate_tokens SummaryMode = Literal["text", "page", "asset"] _MAX_JSON_RETRIES = 1 -_IMAGE_TOKEN_EST = 800 def _read_image_b64(image_path: str) -> str | None: try: with open(image_path, "rb") as handle: - return base64.b64encode(handle.read()).decode() + raw = handle.read() + if not raw: + logger.warning("[summary] empty image file {}", image_path) + return None + return base64.b64encode(raw).decode() except Exception as exc: logger.warning("[summary] failed to read image {}: {}", image_path, exc) return None @@ -105,17 +106,12 @@ def _call_llm( image_paths: list[str], usage_task: str, expect_json: bool, - budget: Any | None, - budget_pool: str, - budget_stage: str | None, channel: Literal["text", "vision"] = "text", ) -> Any | None: - """One text-or-vision call with budget accounting and a single JSON retry. + """One text-or-vision call with a single JSON retry. Returns the parsed object (``expect_json``) or the raw string, or ``None`` on - failure / exhausted budget. Budget is reserved before the call, committed on - success, refunded on failure — matching the prior per-caller bookkeeping but - in one place. + failure. Token usage is recorded by the shared client into the parse tracker. ``channel`` selects BYOK text vs vision credentials when overrides are active. """ @@ -136,13 +132,6 @@ def _call_llm( # Every requested image failed to load; nothing to send. return None - est = estimate_tokens(prompt) + _IMAGE_TOKEN_EST * max(1, len(image_paths)) - if budget is not None and not budget.try_reserve( - budget_pool, est, stage=budget_stage - ): - logger.debug("[summary] budget exhausted for task {}", usage_task) - return None - api_kwargs: dict[str, Any] = {} if expect_json: api_kwargs["response_format"] = {"type": "json_object"} @@ -150,8 +139,6 @@ def _call_llm( resolve = resolve_vision if channel == "vision" else resolve_text effective_model, api_key, api_url = resolve(model) if not effective_model: - if budget is not None: - budget.refund(budget_pool, est=est, stage=budget_stage) return None client = _client_mod.get_openai_client( @@ -161,7 +148,7 @@ def _call_llm( ) for attempt in range(_MAX_JSON_RETRIES + 1): try: - raw, usage = client.chat_completion_with_usage( + raw, _ = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=effective_model, temperature=temperature, @@ -170,14 +157,6 @@ def _call_llm( usage_task=usage_task, **api_kwargs, ) - if budget is not None: - budget.commit( - budget_pool, - actual=usage.get("total_tokens", est), - est=est, - stage=budget_stage, - ) - budget = None # commit once even across the retry loop if not expect_json: if isinstance(raw, str) and raw.strip().lower() in ("null", "none"): return None @@ -195,15 +174,11 @@ def _call_llm( continue return None except UnavailableException: - if budget is not None: - budget.refund(budget_pool, est=est, stage=budget_stage) if usage_task.startswith("page_memory."): raise return None except Exception as exc: logger.warning("[summary] LLM call failed for {}: {}", usage_task, exc) - if budget is not None: - budget.refund(budget_pool, est=est, stage=budget_stage) return None return None @@ -218,9 +193,6 @@ def summarize( max_keywords: int = ..., model: str | None = ..., usage_task: str | None = ..., - budget: Any | None = ..., - budget_pool: str = ..., - budget_stage: str | None = ..., asset_title_hint: str = ..., prompt_task: str | None = ..., prompt_paras: dict[str, Any] | None = ..., @@ -238,9 +210,6 @@ def summarize( max_keywords: int = ..., model: str | None = ..., usage_task: str | None = ..., - budget: Any | None = ..., - budget_pool: str = ..., - budget_stage: str | None = ..., asset_title_hint: str = ..., prompt_task: str | None = ..., prompt_paras: dict[str, Any] | None = ..., @@ -257,13 +226,6 @@ def summarize( max_keywords: int = 5, model: str | None = None, usage_task: str | None = None, - # TODO(parse-budget-cleanup): no live caller passes a non-None budget after - # PROFILE BudgetTracker removal. Drop budget/budget_pool/budget_stage once - # remaining formats stop needing this duck-typed hook, or redirect any - # future limit to token_tracking instead. - budget: Any | None = None, - budget_pool: str = "visual", - budget_stage: str | None = None, asset_title_hint: str = "", prompt_task: str | None = None, prompt_paras: dict[str, Any] | None = None, @@ -279,9 +241,6 @@ def summarize( image_paths: Page or asset image(s). Required for ``page``/``asset`` modes that render from an image; ignored for plain ``text``. - budget: - Optional external reservation ledger. Visual calls reserve from - ``budget_stage``. prompt_task / prompt_paras: Override the prompt used for the image-based page path. Lets a bounded node summary (``page-memory-node-summary`` with ``node_title`` / @@ -300,9 +259,6 @@ def summarize( summary_len=summary_len, model=model, usage_task=usage_task or "summary.asset", - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, asset_title_hint=asset_title_hint, ) return _summarize_body( @@ -313,9 +269,6 @@ def summarize( max_keywords=max_keywords, model=model, usage_task=usage_task or f"summary.{mode}", - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, prompt_task=prompt_task, prompt_paras=prompt_paras, ) @@ -330,9 +283,6 @@ def _summarize_body( max_keywords: int, model: str | None, usage_task: str, - budget: Any | None, - budget_pool: str, - budget_stage: str | None, prompt_task: str | None = None, prompt_paras: dict[str, Any] | None = None, ) -> BodySummary: @@ -355,9 +305,6 @@ def _summarize_body( image_paths=image_paths, usage_task=usage_task, expect_json=True, - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, channel="vision", ) else: @@ -385,9 +332,6 @@ def _summarize_body( image_paths=[], usage_task=usage_task, expect_json=True, - budget=budget, - budget_pool="plan", - budget_stage=None, channel="text", ) @@ -408,9 +352,6 @@ def _summarize_asset( summary_len: int, model: str | None, usage_task: str, - budget: Any | None, - budget_pool: str, - budget_stage: str | None, asset_title_hint: str, ) -> AssetSummary: if not image_paths and not text.strip(): @@ -436,9 +377,6 @@ def _summarize_asset( image_paths=[], usage_task=usage_task, expect_json=False, - budget=budget, - budget_pool="plan", - budget_stage=None, channel="text", ) if isinstance(raw, str) and raw.strip(): @@ -465,9 +403,6 @@ def _summarize_asset( image_paths=image_paths, usage_task=usage_task, expect_json=True, - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, channel="vision", ) if isinstance(parsed, dict): @@ -486,9 +421,6 @@ def transcribe( model: str | None = None, max_tokens: int = 1500, usage_task: str = "summary.transcribe", - budget: Any | None = None, - budget_pool: str = "visual", - budget_stage: str | None = None, ) -> str: """Single OCR primitive (§4.2): transcribe page/image text verbatim. @@ -514,9 +446,6 @@ def transcribe( image_paths=image_paths, usage_task=usage_task, expect_json=True, - budget=budget, - budget_pool=budget_pool, - budget_stage=budget_stage, channel="vision", ) if isinstance(parsed, dict): From 79317ab42a6afb67a78a35726597f2148b6ff9d4 Mon Sep 17 00:00:00 2001 From: cqboy1993 <167045138+EricNGOntos@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:57:53 +0800 Subject: [PATCH 12/14] fix: omit LLM cost_estimate from ZIP manifests (#338) Production and debug ZIP packages no longer embed token cost estimates. Local debug manifests still enrich cost_estimate for operator inspection; ZipResultService strips any residual cost fields before writing manifest.json. Co-authored-by: Cursor --- apps/worker/scripts/debug_text_track.py | 9 ++- .../debug_pm_stage5_tagging_finalize.py | 11 ++- .../services/storage/zip_manifest_schema.py | 41 ++++++++++- .../services/storage/zip_result_service.py | 2 + .../shared/tests/test_zip_manifest_schema.py | 73 +++++++++++++++++++ 5 files changed, 131 insertions(+), 5 deletions(-) create mode 100644 packages/shared-python/shared/tests/test_zip_manifest_schema.py diff --git a/apps/worker/scripts/debug_text_track.py b/apps/worker/scripts/debug_text_track.py index e14176a1f..d85fc9a99 100644 --- a/apps/worker/scripts/debug_text_track.py +++ b/apps/worker/scripts/debug_text_track.py @@ -120,6 +120,10 @@ def _apply_token_usage_to_outputs( trace: dict[str, Any], usage: dict[str, Any], ) -> None: + from shared.services.storage.zip_manifest_schema import ( + enrich_manifest_with_token_cost_estimate, + ) + trace["token_usage"] = usage _write_json(out_dir / "trace.json", trace) manifest_path = out_dir / "manifest.json" @@ -131,7 +135,10 @@ def _apply_token_usage_to_outputs( processing = manifest.setdefault("processing", {}) if isinstance(processing, dict): processing["token_usage"] = usage - _write_json(manifest_path, manifest) + _write_json( + manifest_path, + enrich_manifest_with_token_cost_estimate(manifest), + ) # ── Stage 1: Profile + Shard Plan (PDF only) ─────────────────────────────── diff --git a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py index 079a3349e..09113e4aa 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage5_tagging_finalize.py @@ -611,7 +611,10 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: report_path.write_text(report, encoding="utf-8") if args.finalize: - from shared.services.storage.zip_manifest_schema import ZipManifestBuilder + from shared.services.storage.zip_manifest_schema import ( + ZipManifestBuilder, + enrich_manifest_with_token_cost_estimate, + ) manifest = ZipManifestBuilder().generate_manifest( job_id=filename, @@ -625,7 +628,11 @@ def _load_selected_scope(scope_id: str) -> ScopeResult: hierarchy=hierarchy_dict, ) (out_dir / "manifest.json").write_text( - json.dumps(manifest, ensure_ascii=False, indent=2), + json.dumps( + enrich_manifest_with_token_cost_estimate(manifest), + ensure_ascii=False, + indent=2, + ), encoding="utf-8", ) diff --git a/packages/shared-python/shared/services/storage/zip_manifest_schema.py b/packages/shared-python/shared/services/storage/zip_manifest_schema.py index 4190fe4e2..adc76513a 100644 --- a/packages/shared-python/shared/services/storage/zip_manifest_schema.py +++ b/packages/shared-python/shared/services/storage/zip_manifest_schema.py @@ -2,12 +2,51 @@ from __future__ import annotations +from copy import deepcopy from typing import Any from shared.services.ai.token_costing import build_token_cost_estimate from shared.utils.utc_now import utc_now_naive +def extract_manifest_token_usage(manifest: dict[str, Any]) -> dict[str, Any]: + """Return token usage embedded in a manifest, if present.""" + processing = manifest.get("processing") + if not isinstance(processing, dict): + return {} + + stages = processing.get("stages") + if isinstance(stages, dict): + usage = stages.get("token_usage") + if isinstance(usage, dict): + return usage + + usage = processing.get("token_usage") + if isinstance(usage, dict): + return usage + return {} + + +def enrich_manifest_with_token_cost_estimate(manifest: dict[str, Any]) -> dict[str, Any]: + """Add LLM cost_estimate to a local debug manifest copy.""" + enriched = deepcopy(manifest) + processing = enriched.setdefault("processing", {}) + if not isinstance(processing, dict): + return enriched + token_usage = extract_manifest_token_usage(enriched) + processing["cost_estimate"] = build_token_cost_estimate(token_usage) + return enriched + + +def strip_manifest_cost_fields(manifest: dict[str, Any]) -> dict[str, Any]: + """Remove internal LLM cost fields before writing manifest into a ZIP.""" + stripped = deepcopy(manifest) + processing = stripped.get("processing") + if isinstance(processing, dict): + processing.pop("cost_estimate", None) + return stripped + + class ZipManifestBuilder: def generate_manifest( self, @@ -20,7 +59,6 @@ def generate_manifest( hierarchy: dict[str, Any] | None = None, ) -> dict[str, Any]: stages = job_metadata.get("stages", {}) - token_usage = stages.get("token_usage") if isinstance(stages, dict) else {} return { "version": "2.0", "job_id": job_id, @@ -34,7 +72,6 @@ def generate_manifest( "micro_dollars": job_metadata.get("billing_amount_micro_dollars"), "credits": job_metadata.get("billing_credits"), }, - "cost_estimate": build_token_cost_estimate(token_usage), "timing": { "started_at": job_metadata.get("processing_started_at"), "completed_at": job_metadata.get("processing_completed_at"), diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index c1808002f..160223a9e 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -24,6 +24,7 @@ ) from shared.services.storage.zip_result_resources import ZipResourceCollector from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder +from shared.services.storage.zip_manifest_schema import strip_manifest_cost_fields class ZipResultService: @@ -83,6 +84,7 @@ def generate_zip_package( job_metadata=job_metadata, hierarchy=hierarchy, ) + manifest = strip_manifest_cost_fields(manifest) parse_track = str((job_metadata or {}).get("parse_track") or "") artifact = self._writer.write( ZipPackageWriteRequest( diff --git a/packages/shared-python/shared/tests/test_zip_manifest_schema.py b/packages/shared-python/shared/tests/test_zip_manifest_schema.py new file mode 100644 index 000000000..ad64bd7d2 --- /dev/null +++ b/packages/shared-python/shared/tests/test_zip_manifest_schema.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.services.storage.zip_manifest_schema import ( + ZipManifestBuilder, + enrich_manifest_with_token_cost_estimate, + strip_manifest_cost_fields, +) + + +def test_generate_manifest_omits_llm_cost_estimate() -> None: + manifest = ZipManifestBuilder().generate_manifest( + job_id="job_test", + data_id="data_test", + source_file_name="sample.pdf", + statistics={"total_chunks": 1}, + job_metadata={ + "page_count": 10, + "billing_status": "charged", + "billing_amount_micro_dollars": 150_000, + "billing_credits": 0.15, + "stages": { + "token_usage": { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "calls": 2, + "by_model": {"deepseek-chat": {"total_tokens": 120, "calls": 2}}, + } + }, + }, + hierarchy={"Root": {}}, + ) + + processing = manifest["processing"] + assert "cost_estimate" not in processing + assert processing["cost"]["credits"] == 0.15 + assert processing["stages"]["token_usage"]["by_model"] + + +def test_enrich_and_strip_manifest_cost_fields() -> None: + manifest = ZipManifestBuilder().generate_manifest( + job_id="job_test", + data_id=None, + source_file_name="sample.pdf", + statistics={}, + job_metadata={ + "stages": { + "token_usage": { + "prompt_tokens": 10, + "completion_tokens": 4, + "total_tokens": 14, + "calls": 1, + "by_model": {}, + "by_task": {}, + } + } + }, + ) + + enriched = enrich_manifest_with_token_cost_estimate(manifest) + assert "cost_estimate" in enriched["processing"] + assert enriched["processing"]["cost_estimate"]["currency"] == "USD" + + stripped = strip_manifest_cost_fields(enriched) + assert "cost_estimate" not in stripped["processing"] + assert stripped["processing"]["stages"]["token_usage"]["total_tokens"] == 14 From ccfb464b60344c301990b017e843e773a5295d77 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 00:04:50 +0800 Subject: [PATCH 13/14] Potential fix for pull request finding 'CodeQL / Statement has no effect' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- .../shared-python/shared/services/retrieval/nav_snapshot.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index e8c327c38..43e58c7b7 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -36,7 +36,7 @@ class SnapshotSession(Protocol): """Minimal database interface required by the snapshot loader.""" async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: - ... + pass @dataclass(frozen=True) From a15afe7039363a515ab237493d22aa12cd22d600 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 00:09:41 +0800 Subject: [PATCH 14/14] fix: satisfy snapshot session type contract --- .../shared-python/shared/services/retrieval/nav_snapshot.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index 43e58c7b7..7dc60c6f3 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -36,7 +36,7 @@ class SnapshotSession(Protocol): """Minimal database interface required by the snapshot loader.""" async def execute(self, statement: Executable) -> Result[tuple[object, ...]]: - pass + raise NotImplementedError @dataclass(frozen=True)